use of io.prestosql.statestore.StateStoreProvider 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.StateStoreProvider in project hetu-core by openlookeng.
the class LocalDynamicFilter method create.
public static Optional<LocalDynamicFilter> create(JoinNode planNode, int partitionCount, Session session, TaskId taskId, StateStoreProvider stateStoreProvider) {
Set<String> joinDynamicFilters = planNode.getDynamicFilters().keySet();
// Mapping from probe-side dynamic filters' IDs to their matching probe symbols.
Multimap<String, Symbol> localProbeSymbols = MultimapBuilder.treeKeys().arrayListValues().build();
PlanNode buildNode;
DynamicFilter.Type localType = DynamicFilter.Type.LOCAL;
List<FilterNode> filterNodes = findFilterNodeInStage(planNode);
if (filterNodes.isEmpty()) {
buildNode = planNode.getRight();
mapProbeSymbolsFromCriteria(planNode.getDynamicFilters(), localProbeSymbols, planNode.getCriteria());
localType = DynamicFilter.Type.GLOBAL;
} else {
buildNode = planNode.getRight();
for (FilterNode filterNode : filterNodes) {
mapProbeSymbols(filterNode.getPredicate(), joinDynamicFilters, localProbeSymbols);
}
}
final List<Symbol> buildSideSymbols = buildNode.getOutputSymbols();
Map<String, Integer> localBuildChannels = planNode.getDynamicFilters().entrySet().stream().filter(entry -> localProbeSymbols.containsKey(entry.getKey())).collect(toMap(// Dynamic filter ID
entry -> entry.getKey(), // Build-side channel index
entry -> {
Symbol buildSymbol = entry.getValue();
int buildChannelIndex = buildSideSymbols.indexOf(buildSymbol);
verify(buildChannelIndex >= 0);
return buildChannelIndex;
}));
if (localBuildChannels.isEmpty()) {
return Optional.empty();
}
return Optional.of(new LocalDynamicFilter(localProbeSymbols, localBuildChannels, partitionCount, localType, session, taskId, stateStoreProvider));
}
use of io.prestosql.statestore.StateStoreProvider in project hetu-core by openlookeng.
the class TestHetuServiceInventory method createMockStateStoreProvider.
private StateStoreProvider createMockStateStoreProvider() {
StateStoreProvider mockStateStoreProvider = mock(LocalStateStoreProvider.class);
StateStore mockStateStore = mock(StateStore.class);
when(mockStateStoreProvider.getStateStore()).thenReturn(mockStateStore);
StateMap<String, String> mockDiscoveryInfo = mock(StateMap.class);
Map<String, String> mockDiscoveryMap = new HashMap<>();
mockDiscoveryMap.put(TESTING_HOST, TESTING_PORT);
when(mockDiscoveryInfo.getAll()).thenReturn(mockDiscoveryMap);
when(mockStateStore.getStateCollection(StateStoreConstants.DISCOVERY_SERVICE_COLLECTION_NAME)).thenReturn(mockDiscoveryInfo);
return mockStateStoreProvider;
}
use of io.prestosql.statestore.StateStoreProvider in project hetu-core by openlookeng.
the class TestHetuMetastoreCacheLocal method createMockStateStore.
private StateStore createMockStateStore() {
StateStoreProvider mockStateStoreProvider = mock(LocalStateStoreProvider.class);
StateStore mockStateStore = mock(StateStore.class);
when(mockStateStoreProvider.getStateStore()).thenReturn(mockStateStore);
StateMap<String, String> mockDiscoveryInfo = mock(StateMap.class);
Map<String, String> mockDiscoveryMap = new HashMap<>();
mockDiscoveryMap.put(TESTING_HOST, TESTING_PORT);
when(mockDiscoveryInfo.getAll()).thenReturn(mockDiscoveryMap);
when(mockStateStore.getStateCollection(StateStoreConstants.DISCOVERY_SERVICE_COLLECTION_NAME)).thenReturn(mockDiscoveryInfo);
return mockStateStoreProvider.getStateStore();
}
use of io.prestosql.statestore.StateStoreProvider in project hetu-core by openlookeng.
the class TaskTestUtils method createTestingPlanner.
public static LocalExecutionPlanner createTestingPlanner() {
Metadata metadata = createTestMetadataManager();
PageSourceManager pageSourceManager = new PageSourceManager();
HetuMetaStoreManager hetuMetaStoreManager = new HetuMetaStoreManager();
FeaturesConfig featuresConfig = new FeaturesConfig();
CubeManager cubeManager = new CubeManager(featuresConfig, hetuMetaStoreManager);
pageSourceManager.addConnectorPageSourceProvider(CONNECTOR_ID, new TestingPageSourceProvider());
// we don't start the finalizer so nothing will be collected, which is ok for a test
FinalizerService finalizerService = new FinalizerService();
NodeScheduler nodeScheduler = new NodeScheduler(new LegacyNetworkTopology(), new InMemoryNodeManager(), new NodeSchedulerConfig().setIncludeCoordinator(true), new NodeTaskMap(finalizerService));
NodePartitioningManager nodePartitioningManager = new NodePartitioningManager(nodeScheduler);
PageFunctionCompiler pageFunctionCompiler = new PageFunctionCompiler(metadata, 0);
NodeInfo nodeInfo = new NodeInfo("test");
FileSystemClientManager fileSystemClientManager = new FileSystemClientManager();
SeedStoreManager seedStoreManager = new SeedStoreManager(fileSystemClientManager);
StateStoreProvider stateStoreProvider = new LocalStateStoreProvider(seedStoreManager);
HeuristicIndexerManager heuristicIndexerManager = new HeuristicIndexerManager(new FileSystemClientManager(), new HetuMetaStoreManager());
return new LocalExecutionPlanner(metadata, new TypeAnalyzer(new SqlParser(), metadata), Optional.empty(), pageSourceManager, new IndexManager(), nodePartitioningManager, new PageSinkManager(), new MockExchangeClientSupplier(), new ExpressionCompiler(metadata, pageFunctionCompiler), pageFunctionCompiler, new JoinFilterFunctionCompiler(metadata), new IndexJoinLookupStats(), new TaskManagerConfig(), new GenericSpillerFactory((types, spillContext, memoryContext) -> {
throw new UnsupportedOperationException();
}), (types, spillContext, memoryContext) -> {
throw new UnsupportedOperationException();
}, (types, partitionFunction, spillContext, memoryContext) -> {
throw new UnsupportedOperationException();
}, new PagesIndex.TestingFactory(false), new JoinCompiler(metadata), new LookupJoinOperators(), new OrderingCompiler(), nodeInfo, stateStoreProvider, new StateStoreListenerManager(stateStoreProvider), new DynamicFilterCacheManager(), heuristicIndexerManager, cubeManager);
}
Aggregations