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");
}
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);
}
}
Aggregations