Search in sources :

Example 1 with MockStateMap

use of io.prestosql.statestore.MockStateMap 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");
}
Also used : HashMap(java.util.HashMap) Symbol(io.prestosql.spi.plan.Symbol) MockStateMap(io.prestosql.statestore.MockStateMap) StateMap(io.prestosql.spi.statestore.StateMap) StateStoreListenerManager(io.prestosql.statestore.listener.StateStoreListenerManager) StateStoreProvider(io.prestosql.statestore.StateStoreProvider) MockStateMap(io.prestosql.statestore.MockStateMap) TestingColumnHandle(io.prestosql.spi.connector.TestingColumnHandle) DynamicFilters(io.prestosql.sql.DynamicFilters) DynamicFilterListener(io.prestosql.dynamicfilter.DynamicFilterListener) DynamicFilterCacheManager(io.prestosql.dynamicfilter.DynamicFilterCacheManager) TaskContext(io.prestosql.operator.TaskContext) DynamicFilter(io.prestosql.spi.dynamicfilter.DynamicFilter) HashSetDynamicFilter(io.prestosql.spi.dynamicfilter.HashSetDynamicFilter) QueryId(io.prestosql.spi.QueryId) StateStore(io.prestosql.spi.statestore.StateStore) HashSetDynamicFilter(io.prestosql.spi.dynamicfilter.HashSetDynamicFilter) TableScanNode(io.prestosql.spi.plan.TableScanNode) VariableReferenceExpression(io.prestosql.spi.relation.VariableReferenceExpression) MockStateMap(io.prestosql.statestore.MockStateMap) HashMap(java.util.HashMap) Map(java.util.Map) ImmutableMap(com.google.common.collect.ImmutableMap) StateMap(io.prestosql.spi.statestore.StateMap) Session(io.prestosql.Session) Test(org.testng.annotations.Test)

Example 2 with MockStateMap

use of io.prestosql.statestore.MockStateMap in project hetu-core by openlookeng.

the class TestDistributedResourceGroup method testStateStoreFetchAndUpdate.

@Test
public void testStateStoreFetchAndUpdate() throws JsonProcessingException {
    synchronized (lock) {
        DistributedResourceGroupTemp root = new DistributedResourceGroupTemp(Optional.empty(), "root", (group, export) -> {
        }, directExecutor(), statestore, internalNodeManager);
        resourceGroupBasicSetUp(root, ONE_MEGABYTE, 1, 1);
        MockManagedQueryExecution query1 = new MockManagedQueryExecution(100);
        query1.setResourceGroupId(root.getId());
        MockManagedQueryExecution query2 = new MockManagedQueryExecution(0);
        query2.setResourceGroupId(root.getId());
        Map<String, String> mockMap = new HashMap<>();
        MockStateMap<String, String> mockStateMap = new MockStateMap<>("127.0.0.1-resourceaggrstats", mockMap);
        when(statestore.getOrCreateStateCollection(anyString(), anyObject())).thenReturn(mockStateMap);
        when(internalNodeManager.getCurrentNode()).thenReturn(new InternalNode("node1", URI.create("local://127.0.0.1"), NodeVersion.UNKNOWN, true));
        when(internalNodeManager.getCoordinators()).thenReturn(ImmutableSet.of(new InternalNode("node1", URI.create("local://127.0.0.1"), NodeVersion.UNKNOWN, true)));
        root.run(query1);
        DistributedResourceGroupAggrStats rootStats = MAPPER.readerFor(DistributedResourceGroupAggrStats.class).readValue(mockMap.get("root"));
        assertEquals(rootStats.getRunningQueries(), 1);
        root.run(query2);
        rootStats = MAPPER.readerFor(DistributedResourceGroupAggrStats.class).readValue(mockMap.get("root"));
        assertEquals(rootStats.getQueuedQueries(), 1);
        assertEquals(rootStats.getCachedMemoryUsageBytes(), 100);
        when(internalNodeManager.getCurrentNode()).thenReturn(new InternalNode("node2", URI.create("local://127.0.0.2"), NodeVersion.UNKNOWN, true));
        assertEquals(root.getGlobalCachedMemoryUsageBytes(), 200);
    }
}
Also used : MockStateMap(io.prestosql.statestore.MockStateMap) MockManagedQueryExecution(io.prestosql.execution.MockManagedQueryExecution) HashMap(java.util.HashMap) Matchers.anyString(org.mockito.Matchers.anyString) InternalNode(io.prestosql.metadata.InternalNode) Test(org.testng.annotations.Test)

Aggregations

MockStateMap (io.prestosql.statestore.MockStateMap)2 HashMap (java.util.HashMap)2 Test (org.testng.annotations.Test)2 ImmutableMap (com.google.common.collect.ImmutableMap)1 Session (io.prestosql.Session)1 DynamicFilterCacheManager (io.prestosql.dynamicfilter.DynamicFilterCacheManager)1 DynamicFilterListener (io.prestosql.dynamicfilter.DynamicFilterListener)1 MockManagedQueryExecution (io.prestosql.execution.MockManagedQueryExecution)1 InternalNode (io.prestosql.metadata.InternalNode)1 TaskContext (io.prestosql.operator.TaskContext)1 QueryId (io.prestosql.spi.QueryId)1 TestingColumnHandle (io.prestosql.spi.connector.TestingColumnHandle)1 DynamicFilter (io.prestosql.spi.dynamicfilter.DynamicFilter)1 HashSetDynamicFilter (io.prestosql.spi.dynamicfilter.HashSetDynamicFilter)1 Symbol (io.prestosql.spi.plan.Symbol)1 TableScanNode (io.prestosql.spi.plan.TableScanNode)1 VariableReferenceExpression (io.prestosql.spi.relation.VariableReferenceExpression)1 StateMap (io.prestosql.spi.statestore.StateMap)1 StateStore (io.prestosql.spi.statestore.StateStore)1 DynamicFilters (io.prestosql.sql.DynamicFilters)1