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

Example 1 with QueryId

use of io.prestosql.spi.QueryId in project hetu-core by openlookeng.

the class TestDynamicFiltersCollector method TestCollectingGlobalDynamicFilters.

@Test
public void TestCollectingGlobalDynamicFilters() throws InterruptedException {
    final QueryId queryId = new QueryId("test_query");
    final String filterId = "1";
    final String columnName = "column";
    final TestingColumnHandle columnHandle = new TestingColumnHandle(columnName);
    final Set<String> valueSet = ImmutableSet.of("1", "2", "3");
    TaskContext taskContext = mock(TaskContext.class);
    Session session = testSessionBuilder().setQueryId(queryId).setSystemProperty(ENABLE_DYNAMIC_FILTERING, "true").setSystemProperty(DYNAMIC_FILTERING_DATA_TYPE, "HASHSET").build();
    when(taskContext.getSession()).thenReturn(session);
    // set up state store and merged dynamic filters map
    Map mockMap = new HashMap<>();
    StateStoreProvider stateStoreProvider = mock(StateStoreProvider.class);
    StateStore stateStore = mock(StateStore.class);
    StateMap stateMap = new MockStateMap<>("test-map", mockMap);
    when(stateStoreProvider.getStateStore()).thenReturn(stateStore);
    when(stateStore.getStateCollection(any())).thenReturn(stateMap);
    when(stateStore.createStateMap(any())).thenReturn(stateMap);
    when(stateStore.getOrCreateStateCollection(any(), any())).thenReturn(stateMap);
    // set up state store listener and dynamic filter cache
    StateStoreListenerManager stateStoreListenerManager = new StateStoreListenerManager(stateStoreProvider);
    DynamicFilterCacheManager dynamicFilterCacheManager = new DynamicFilterCacheManager();
    stateStoreListenerManager.addStateStoreListener(new DynamicFilterListener(dynamicFilterCacheManager), MERGED_DYNAMIC_FILTERS);
    LocalDynamicFiltersCollector collector = new LocalDynamicFiltersCollector(taskContext, Optional.empty(), dynamicFilterCacheManager);
    TableScanNode tableScan = mock(TableScanNode.class);
    when(tableScan.getAssignments()).thenReturn(ImmutableMap.of(new Symbol(columnName), columnHandle));
    List<DynamicFilters.Descriptor> dynamicFilterDescriptors = ImmutableList.of(new DynamicFilters.Descriptor(filterId, new VariableReferenceExpression(columnName, BIGINT)));
    collector.initContext(ImmutableList.of(dynamicFilterDescriptors), SymbolUtils.toLayOut(tableScan.getOutputSymbols()));
    assertTrue(collector.getDynamicFilters(tableScan).isEmpty(), "there should be no dynamic filter available");
    // put some values in state store as a new dynamic filter
    // and wait for the listener to process the event
    stateMap.put(createKey(DynamicFilterUtils.FILTERPREFIX, filterId, queryId.getId()), valueSet);
    TimeUnit.MILLISECONDS.sleep(100);
    // get available dynamic filter and verify it
    List<Map<ColumnHandle, DynamicFilter>> dynamicFilters = collector.getDynamicFilters(tableScan);
    assertEquals(dynamicFilters.size(), 1, "there should be a new dynamic filter");
    assertEquals(dynamicFilters.size(), 1);
    DynamicFilter dynamicFilter = dynamicFilters.get(0).get(columnHandle);
    assertTrue(dynamicFilter instanceof HashSetDynamicFilter, "new dynamic filter should be hashset");
    assertEquals(dynamicFilter.getSize(), valueSet.size(), "new dynamic filter should have correct size");
    for (String value : valueSet) {
        assertTrue(dynamicFilter.contains(value), "new dynamic filter should contain correct values");
    }
    // clean up when task finishes
    collector.removeDynamicFilter(true);
    DynamicFilter cachedFilter = dynamicFilterCacheManager.getDynamicFilter(DynamicFilterCacheManager.createCacheKey(filterId, queryId.getId()));
    assertNull(cachedFilter, "cached dynamic filter should have been removed");
}

Example 2 with QueryId

use of io.prestosql.spi.QueryId in project hetu-core by openlookeng.

the class TestQueuesDb method testBasic.

@Test(timeOut = 60_000)
public void testBasic() throws Exception {
    // submit first "dashboard" query
    QueryId firstDashboardQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY);
    // wait for the first "dashboard" query to start
    waitForQueryState(queryRunner, firstDashboardQuery, RUNNING);
    waitForRunningQueryCount(queryRunner, 1);
    // submit second "dashboard" query
    QueryId secondDashboardQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY);
    MILLISECONDS.sleep(2000);
    // wait for the second "dashboard" query to be queued ("dashboard.${USER}" queue strategy only allows one "dashboard" query to be accepted for execution)
    waitForQueryState(queryRunner, secondDashboardQuery, QUEUED);
    waitForRunningQueryCount(queryRunner, 1);
    // Update db to allow for 1 more running query in dashboard resource group
    // Hetu: add parameters softReservedMemory and hardReservedConcurrency
    dao.updateResourceGroup(3, "user-${USER}", "1MB", "10%", 3, 4, 4, 4, null, null, null, null, null, "RECENT_QUERIES", 1L, TEST_ENVIRONMENT);
    dao.updateResourceGroup(5, "dashboard-${USER}", "1MB", "10%", 1, 2, 2, 2, null, null, null, null, null, "RECENT_QUERIES", 3L, TEST_ENVIRONMENT);
    waitForQueryState(queryRunner, secondDashboardQuery, RUNNING);
    QueryId thirdDashboardQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY);
    waitForQueryState(queryRunner, thirdDashboardQuery, QUEUED);
    waitForRunningQueryCount(queryRunner, 2);
    // submit first non "dashboard" query
    QueryId firstNonDashboardQuery = createQuery(queryRunner, adhocSession(), LONG_LASTING_QUERY);
    // wait for the first non "dashboard" query to start
    waitForQueryState(queryRunner, firstNonDashboardQuery, RUNNING);
    waitForRunningQueryCount(queryRunner, 3);
    // submit second non "dashboard" query
    QueryId secondNonDashboardQuery = createQuery(queryRunner, adhocSession(), LONG_LASTING_QUERY);
    // wait for the second non "dashboard" query to start
    waitForQueryState(queryRunner, secondNonDashboardQuery, RUNNING);
    waitForRunningQueryCount(queryRunner, 4);
    // cancel first "dashboard" query, the second "dashboard" query and second non "dashboard" query should start running
    cancelQuery(queryRunner, firstDashboardQuery);
    waitForQueryState(queryRunner, firstDashboardQuery, FAILED);
    waitForQueryState(queryRunner, thirdDashboardQuery, RUNNING);
    waitForRunningQueryCount(queryRunner, 4);
    waitForCompleteQueryCount(queryRunner, 1);
}

Example 3 with QueryId

use of io.prestosql.spi.QueryId in project hetu-core by openlookeng.

the class TestQueuesDb method testQuerySystemTableResourceGroup.

@Test(timeOut = 60_000)
public void testQuerySystemTableResourceGroup() throws Exception {
    QueryId firstQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY);
    waitForQueryState(queryRunner, firstQuery, RUNNING);
    MaterializedResult result = queryRunner.execute("SELECT resource_group_id FROM system.runtime.queries WHERE source = 'dashboard'");
    assertEquals(result.getOnlyValue(), ImmutableList.of("global", "user-user", "dashboard-user"));
}

Example 4 with QueryId

use of io.prestosql.spi.QueryId in project hetu-core by openlookeng.

the class TestQueuesDb method assertResourceGroupWithClientTags.

private void assertResourceGroupWithClientTags(Set<String> clientTags, ResourceGroupId expectedResourceGroup) throws InterruptedException {
    Session session = testSessionBuilder().setCatalog("tpch").setSchema("sf100000").setSource("client_tags").setClientTags(clientTags).build();
    QueryId queryId = createQuery(queryRunner, session, LONG_LASTING_QUERY);
    waitForQueryState(queryRunner, queryId, ImmutableSet.of(RUNNING, FINISHED));
    Optional<ResourceGroupId> resourceGroupId = queryRunner.getCoordinator().getQueryManager().getFullQueryInfo(queryId).getResourceGroupId();
    assertTrue(resourceGroupId.isPresent(), "Query should have a resource group");
    assertEquals(resourceGroupId.get(), expectedResourceGroup, format("Expected: '%s' resource group, found: %s", expectedResourceGroup, resourceGroupId.get()));
}

Example 5 with QueryId

use of io.prestosql.spi.QueryId in project hetu-core by openlookeng.

the class TestQueuesDb method testRejection.

@Test(timeOut = 60_000)
public void testRejection() throws Exception {
    InternalResourceGroupManager<?> manager = queryRunner.getCoordinator().getResourceGroupManager().get();
    DbResourceGroupConfigurationManager dbConfigurationManager = (DbResourceGroupConfigurationManager) manager.getConfigurationManager();
    // Verify the query cannot be submitted
    QueryId queryId = createQuery(queryRunner, rejectingSession(), LONG_LASTING_QUERY);
    waitForQueryState(queryRunner, queryId, FAILED);
    DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager();
    assertEquals(dispatchManager.getQueryInfo(queryId).getErrorCode(), QUERY_REJECTED.toErrorCode());
    int selectorCount = getSelectors(queryRunner).size();
    dao.insertSelector(4, 100_000, "user.*", "(?i).*reject.*", null, null, null);
    dbConfigurationManager.load();
    assertEquals(getSelectors(queryRunner).size(), selectorCount + 1);
    // Verify the query can be submitted
    queryId = createQuery(queryRunner, rejectingSession(), LONG_LASTING_QUERY);
    waitForQueryState(queryRunner, queryId, RUNNING);
    dao.deleteSelector(4, "user.*", "(?i).*reject.*", null);
    dbConfigurationManager.load();
    // Verify the query cannot be submitted
    queryId = createQuery(queryRunner, rejectingSession(), LONG_LASTING_QUERY);
    waitForQueryState(queryRunner, queryId, FAILED);
}