Examples with QueryId io.trino.spi.QueryId used on opensource projects

Example 36 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class LowMemoryKillerTestingUtils method tasksMemoryInfoForNode.

private static ListMultimap<QueryId, TaskMemoryInfo> tasksMemoryInfoForNode(String nodeIdentifier, Map<String, Map<String, Map<String, Long>>> tasks) {
    ImmutableListMultimap.Builder<QueryId, TaskMemoryInfo> result = ImmutableListMultimap.builder();
    for (Map.Entry<String, Map<String, Map<String, Long>>> queryNodesEntry : tasks.entrySet()) {
        QueryId query = QueryId.valueOf(queryNodesEntry.getKey());
        for (Map.Entry<String, Map<String, Long>> nodeTasksEntry : queryNodesEntry.getValue().entrySet()) {
            if (!nodeIdentifier.equals(nodeTasksEntry.getKey())) {
                continue;
            }
            for (Map.Entry<String, Long> taskReservationEntry : nodeTasksEntry.getValue().entrySet()) {
                TaskId taskId = TaskId.valueOf(taskReservationEntry.getKey());
                long taskReservation = taskReservationEntry.getValue();
                result.put(query, new TaskMemoryInfo(taskId, taskReservation));
            }
        }
    }
    return result.build();
}

Example 37 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class TestMemoryPools method testTaggedAllocations.

@Test
public void testTaggedAllocations() {
    QueryId testQuery = new QueryId("test_query");
    MemoryPool testPool = new MemoryPool(DataSize.ofBytes(1000));
    testPool.reserve(testQuery, "test_tag", 10);
    Map<String, Long> allocations = testPool.getTaggedMemoryAllocations().get(testQuery);
    assertEquals(allocations, ImmutableMap.of("test_tag", 10L));
    // free 5 bytes for test_tag
    testPool.free(testQuery, "test_tag", 5);
    assertEquals(allocations, ImmutableMap.of("test_tag", 5L));
    testPool.reserve(testQuery, "test_tag2", 20);
    assertEquals(allocations, ImmutableMap.of("test_tag", 5L, "test_tag2", 20L));
    // free the remaining 5 bytes for test_tag
    testPool.free(testQuery, "test_tag", 5);
    assertEquals(allocations, ImmutableMap.of("test_tag2", 20L));
    // free all for test_tag2
    testPool.free(testQuery, "test_tag2", 20);
    assertEquals(testPool.getTaggedMemoryAllocations().size(), 0);
}

Example 38 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class BaseDynamicPartitionPruningTest method getQueryInputPositions.

private long getQueryInputPositions(Session session, @Language("SQL") String sql, int expectedRowCount) {
    DistributedQueryRunner runner = (DistributedQueryRunner) getQueryRunner();
    ResultWithQueryId<MaterializedResult> result = runner.executeWithQueryId(session, sql);
    assertThat(result.getResult().getRowCount()).isEqualTo(expectedRowCount);
    QueryId queryId = result.getQueryId();
    QueryStats stats = runner.getCoordinator().getQueryManager().getFullQueryInfo(queryId).getQueryStats();
    return stats.getPhysicalInputPositions();
}

Example 39 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class AbstractTestingTrinoClient method execute.

public ResultWithQueryId<T> execute(Session session, @Language("SQL") String sql) {
    ResultsSession<T> resultsSession = getResultSession(session);
    ClientSession clientSession = toClientSession(session, trinoServer.getBaseUrl(), new Duration(2, TimeUnit.MINUTES));
    try (StatementClient client = newStatementClient(httpClient, clientSession, sql)) {
        while (client.isRunning()) {
            resultsSession.addResults(client.currentStatusInfo(), client.currentData());
            client.advance();
        }
        checkState(client.isFinished());
        QueryError error = client.finalStatusInfo().getError();
        if (error == null) {
            QueryStatusInfo results = client.finalStatusInfo();
            if (results.getUpdateType() != null) {
                resultsSession.setUpdateType(results.getUpdateType());
            }
            if (results.getUpdateCount() != null) {
                resultsSession.setUpdateCount(results.getUpdateCount());
            }
            resultsSession.setWarnings(results.getWarnings());
            resultsSession.setStatementStats(results.getStats());
            T result = resultsSession.build(client.getSetSessionProperties(), client.getResetSessionProperties());
            return new ResultWithQueryId<>(new QueryId(results.getId()), result);
        }
        if (error.getFailureInfo() != null) {
            RuntimeException remoteException = error.getFailureInfo().toException();
            throw new RuntimeException(Optional.ofNullable(remoteException.getMessage()).orElseGet(remoteException::toString), remoteException);
        }
        throw new RuntimeException("Query failed: " + error.getMessage());
    // dump query info to console for debugging (NOTE: not pretty printed)
    // JsonCodec<QueryInfo> queryInfoJsonCodec = createCodecFactory().prettyPrint().jsonCodec(QueryInfo.class);
    // log.info("\n" + queryInfoJsonCodec.toJson(queryInfo));
    }
}

Example 40 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class TestPendingStageState method testPendingState.

@Test(timeOut = 30_000)
public void testPendingState() throws Exception {
    QueryId queryId = createQuery(queryRunner, TEST_SESSION, "SELECT * FROM tpch.sf1000.lineitem limit 1");
    waitForQueryState(queryRunner, queryId, RUNNING);
    // wait for the query to finish producing results, but don't poll them
    assertEventually(new Duration(10, SECONDS), () -> assertEquals(queryRunner.getCoordinator().getFullQueryInfo(queryId).getOutputStage().get().getState(), StageState.RUNNING));
    // wait for the sub stages to go to pending state
    assertEventually(new Duration(10, SECONDS), () -> assertEquals(queryRunner.getCoordinator().getFullQueryInfo(queryId).getOutputStage().get().getSubStages().get(0).getState(), StageState.PENDING));
    QueryInfo queryInfo = queryRunner.getCoordinator().getFullQueryInfo(queryId);
    assertEquals(queryInfo.getState(), RUNNING);
    assertEquals(queryInfo.getOutputStage().get().getState(), StageState.RUNNING);
    assertEquals(queryInfo.getOutputStage().get().getSubStages().size(), 1);
    assertEquals(queryInfo.getOutputStage().get().getSubStages().get(0).getState(), StageState.PENDING);
}