use of io.prestosql.sql.DynamicFilters 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.sql.DynamicFilters in project hetu-core by openlookeng.
the class TestDynamicFilterServiceWithBloomFilter method testRegisterAndMergeDynamicFilters.
@Test
public void testRegisterAndMergeDynamicFilters() throws InterruptedException {
registerDf(filterId, session, PARTITIONED, dynamicFilterService);
// Test getDynamicFilterSupplier
VariableReferenceExpression mockExpression = mock(VariableReferenceExpression.class);
when(mockExpression.getName()).thenReturn("name");
ColumnHandle mockColumnHandle = mock(ColumnHandle.class);
Supplier<List<Set<DynamicFilter>>> dynamicFilterSupplier = DynamicFilterService.getDynamicFilterSupplier(session.getQueryId(), ImmutableList.of(ImmutableList.of(new DynamicFilters.Descriptor(filterId, mockExpression))), ImmutableMap.of(new Symbol("name"), mockColumnHandle));
assertTrue(dynamicFilterSupplier.get().isEmpty(), "should return empty dynamic filter set when dynamic filters are not available");
mockLocalDynamicFilter("task1.0", filterId, session.getQueryId().toString(), Arrays.asList("1", "2", "3", "4"));
mockLocalDynamicFilter("task1.1", filterId, session.getQueryId().toString(), Arrays.asList("5", "6", "7", "8"));
Thread.sleep(3000);
BloomFilter bf = fetchDynamicFilter(filterId, session.getQueryId().toString());
for (int i = 1; i < 9; i++) {
assertTrue(bf.test((String.valueOf(i).getBytes(StandardCharsets.UTF_8))));
}
assertFalse(bf.test("10".getBytes(StandardCharsets.UTF_8)));
// Test getDynamicFilterSupplier
dynamicFilterSupplier = DynamicFilterService.getDynamicFilterSupplier(session.getQueryId(), ImmutableList.of(ImmutableList.of(new DynamicFilters.Descriptor(filterId, mockExpression))), ImmutableMap.of(new Symbol("name"), mockColumnHandle));
List<Set<DynamicFilter>> dynamicFilters = dynamicFilterSupplier.get();
assertFalse(dynamicFilters == null, "dynamic filters should be ready");
assertEquals(dynamicFilters.size(), 1, "there should be 1 dynamic filter in supplier");
DynamicFilter dynamicFilter = dynamicFilters.get(0).iterator().next();
for (int i = 1; i < 9; i++) {
assertTrue(dynamicFilter.contains(String.valueOf(i)));
}
assertFalse(dynamicFilter.contains("10"));
dynamicFilterSupplier = DynamicFilterService.getDynamicFilterSupplier(new QueryId("invalid"), ImmutableList.of(ImmutableList.of(new DynamicFilters.Descriptor(filterId, mockExpression))), ImmutableMap.of(new Symbol("name"), mockColumnHandle));
assertTrue(dynamicFilterSupplier.get().isEmpty(), "should return empty dynamic filter set for invalid or non-existing queryId");
String queryId = session.getQueryId().getId();
assertEquals(stateStoreProvider.getStateStore().getStateCollection(createKey(DynamicFilterUtils.PARTIALPREFIX, filterId, queryId)).size(), 2);
assertEquals(stateStoreProvider.getStateStore().getStateCollection(createKey(DynamicFilterUtils.TASKSPREFIX, filterId, queryId)).size(), 2);
dynamicFilterService.clearDynamicFiltersForQuery(queryId);
Thread.sleep(1000);
assertEquals(stateStoreProvider.getStateStore().getStateCollection(createKey(DynamicFilterUtils.PARTIALPREFIX, filterId, queryId)).size(), 0);
assertEquals(stateStoreProvider.getStateStore().getStateCollection(createKey(DynamicFilterUtils.TASKSPREFIX, filterId, queryId)).size(), 0);
}
use of io.prestosql.sql.DynamicFilters in project hetu-core by openlookeng.
the class TestDynamicFilterServiceWithHashSet method testRegisterAndMergeDynamicFiltersHashSet.
@Test
public void testRegisterAndMergeDynamicFiltersHashSet() throws InterruptedException {
setUpHashSet();
filterId = "df2";
registerDf(filterId, session, PARTITIONED, dynamicFilterService);
// Test getDynamicFilterSupplier
VariableReferenceExpression mockExpression = mock(VariableReferenceExpression.class);
when(mockExpression.getName()).thenReturn("name");
ColumnHandle mockColumnHandle = mock(ColumnHandle.class);
Supplier<List<Set<DynamicFilter>>> dynamicFilterSupplier = DynamicFilterService.getDynamicFilterSupplier(session.getQueryId(), ImmutableList.of(ImmutableList.of(new DynamicFilters.Descriptor(filterId, mockExpression))), ImmutableMap.of(new Symbol("name"), mockColumnHandle));
assertTrue(dynamicFilterSupplier.get().isEmpty(), "should return empty dynamic filter set when dynamic filters are not available");
mockLocalDynamicFilterHashSet("task1.0", filterId, session.getQueryId().toString(), Arrays.asList("11", "12", "13", "14"));
mockLocalDynamicFilterHashSet("task1.1", filterId, session.getQueryId().toString(), Arrays.asList("15", "16", "17", "18"));
Thread.sleep(2000);
Set hs = fetchDynamicFilterHashSet(filterId, session.getQueryId().toString());
for (int i = 11; i < 19; i++) {
Assert.assertEquals(true, hs.contains(i + ""));
}
Assert.assertEquals(false, hs.contains("10"));
// Test getDynamicFilterSupplier
dynamicFilterSupplier = DynamicFilterService.getDynamicFilterSupplier(session.getQueryId(), ImmutableList.of(ImmutableList.of(new DynamicFilters.Descriptor(filterId, mockExpression))), ImmutableMap.of(new Symbol("name"), mockColumnHandle));
List<Set<DynamicFilter>> dynamicFilters = dynamicFilterSupplier.get();
assertFalse(dynamicFilters == null, "dynamic filters should be ready");
assertEquals(dynamicFilters.size(), 1, "there should be 1 dynamic filter in supplier");
HashSetDynamicFilter hsDF = ((HashSetDynamicFilter) dynamicFilters.get(0).toArray()[0]);
assertEquals(hs, hsDF.getSetValues(), "dynamic filter in supplier should be the same as the one merged");
dynamicFilterSupplier = DynamicFilterService.getDynamicFilterSupplier(new QueryId("invalid"), ImmutableList.of(ImmutableList.of(new DynamicFilters.Descriptor(filterId, mockExpression))), ImmutableMap.of(new Symbol("name"), mockColumnHandle));
assertTrue(dynamicFilterSupplier.get().isEmpty(), "should return empty dynamic filter set for invalid or non-existing queryId");
String queryId = session.getQueryId().getId();
assertEquals(stateStoreProvider.getStateStore().getStateCollection(createKey(DynamicFilterUtils.PARTIALPREFIX, filterId, queryId)).size(), 2);
assertEquals(stateStoreProvider.getStateStore().getStateCollection(createKey(DynamicFilterUtils.TASKSPREFIX, filterId, queryId)).size(), 2);
dynamicFilterService.clearDynamicFiltersForQuery(queryId);
Thread.sleep(1000);
assertEquals(stateStoreProvider.getStateStore().getStateCollection(createKey(DynamicFilterUtils.PARTIALPREFIX, filterId, queryId)).size(), 0);
assertEquals(stateStoreProvider.getStateStore().getStateCollection(createKey(DynamicFilterUtils.TASKSPREFIX, filterId, queryId)).size(), 0);
}
use of io.prestosql.sql.DynamicFilters in project hetu-core by openlookeng.
the class DynamicFilterService method extractSourceExpressionSymbols.
private static Map<String, Symbol> extractSourceExpressionSymbols(List<DynamicFilters.Descriptor> dynamicFilters) {
ImmutableMap.Builder<String, Symbol> resultBuilder = ImmutableMap.builder();
for (DynamicFilters.Descriptor descriptor : dynamicFilters) {
RowExpression expression = descriptor.getInput();
// Extract the column symbols from CAST expressions
while (expression instanceof CallExpression) {
expression = ((CallExpression) expression).getArguments().get(0);
}
if (!(expression instanceof VariableReferenceExpression)) {
continue;
}
resultBuilder.put(descriptor.getId(), new Symbol(((VariableReferenceExpression) expression).getName()));
}
return resultBuilder.build();
}
Aggregations