Examples with StateStore io.prestosql.spi.statestore.StateStore used on opensource projects

Example 16 with StateStore

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

the class TestSplitCacheStateInitializer method testCreateStateCollectionIfNotExists.

@Test
public void testCreateStateCollectionIfNotExists() {
    StateStoreProvider provider = mock(StateStoreProvider.class);
    StateStore stateStore = mock(StateStore.class);
    when(stateStore.getName()).thenReturn("mock");
    when(stateStore.getStateCollection(StateStoreConstants.SPLIT_CACHE_METADATA_NAME)).thenReturn(null).thenReturn(new MockStateMap<>(StateStoreConstants.SPLIT_CACHE_METADATA_NAME, new HashMap<>()));
    // simulate behaviour wait until state store is ready
    when(provider.getStateStore()).thenReturn(stateStore);
    SplitCacheMap splitCacheMap = createNew();
    assertFalse(splitCacheMap.cacheExists(table1QN));
    assertFalse(splitCacheMap.cacheExists(table2QN));
    AtomicReference<SplitCacheStateInitializer.InitializationStatus> status = new AtomicReference<>(SplitCacheStateInitializer.InitializationStatus.INITIALIZING);
    SplitCacheStateInitializer initializer = new SplitCacheStateInitializer(provider, splitCacheMap, new Duration(300, TimeUnit.MILLISECONDS), new Duration(1, TimeUnit.SECONDS), objectMapper, status);
    try {
        initializer.start();
        // sleep for sometime till local split cache map is initialized
        try {
            Thread.sleep(1000);
        } catch (InterruptedException e) {
        // ignore - nothing to do here
        }
    } finally {
        initializer.stop();
    }
    verify(provider, atLeastOnce()).getStateStore();
    verify(stateStore, times(1)).createStateMap(eq(StateStoreConstants.SPLIT_CACHE_METADATA_NAME), any(SplitCacheStateStoreChangesListener.class));
    assertEquals(status.get(), SplitCacheStateInitializer.InitializationStatus.COMPLETED);
}

Example 17 with StateStore

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

the class DynamicFilterService method registerTasksHelper.

private void registerTasksHelper(PlanNode node, Symbol buildSymbol, Map<String, Symbol> dynamicFiltersMap, Set<TaskId> taskIds, Set<InternalNode> workers, StageStateMachine stateMachine) {
    final StateStore stateStore = stateStoreProvider.getStateStore();
    String queryId = stateMachine.getSession().getQueryId().toString();
    for (Map.Entry<String, Symbol> entry : dynamicFiltersMap.entrySet()) {
        Symbol buildSymbolToCheck = buildSymbol != null ? buildSymbol : node.getOutputSymbols().contains(entry.getValue()) ? entry.getValue() : null;
        if (buildSymbolToCheck != null && entry.getValue().getName().equals(buildSymbol.getName())) {
            String filterId = entry.getKey();
            stateStore.createStateCollection(createKey(DynamicFilterUtils.TASKSPREFIX, filterId, queryId), SET);
            stateStore.createStateCollection(createKey(DynamicFilterUtils.PARTIALPREFIX, filterId, queryId), SET);
            dynamicFilters.putIfAbsent(queryId, new ConcurrentHashMap<>());
            Map<String, DynamicFilterRegistryInfo> filters = dynamicFilters.get(queryId);
            if (node instanceof JoinNode) {
                filters.put(filterId, extractDynamicFilterRegistryInfo((JoinNode) node, stateMachine.getSession(), filterId));
            } else if (node instanceof SemiJoinNode) {
                filters.put(filterId, extractDynamicFilterRegistryInfo((SemiJoinNode) node, stateMachine.getSession()));
            }
            dynamicFiltersToTask.putIfAbsent(filterId + "-" + queryId, new CopyOnWriteArraySet<>());
            CopyOnWriteArraySet<TaskId> taskSet = dynamicFiltersToTask.get(filterId + "-" + queryId);
            taskSet.addAll(taskIds);
            log.debug("registerTasks source " + filterId + " filters:" + filters + ", workers: " + workers.stream().map(x -> x.getNodeIdentifier()).collect(Collectors.joining(",")) + ", taskIds: " + taskIds.stream().map(TaskId::toString).collect(Collectors.joining(",")));
        }
    }
}

Example 18 with StateStore

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

the class DynamicFilterService method mergeDynamicFilters.

/**
 * Global Dynamic Filter merging, periodically looks for dynamic filters that can be merged and merges them
 */
private void mergeDynamicFilters() {
    final StateStore stateStore = stateStoreProvider.getStateStore();
    for (Map.Entry<String, Map<String, DynamicFilterRegistryInfo>> queryToDynamicFiltersEntry : dynamicFilters.entrySet()) {
        final String queryId = queryToDynamicFiltersEntry.getKey();
        if (!cachedDynamicFilters.containsKey(queryId)) {
            cachedDynamicFilters.put(queryId, new ConcurrentHashMap<>());
        }
        Map<String, DynamicFilter> cachedDynamicFiltersForQuery = cachedDynamicFilters.get(queryId);
        StateMap mergedDynamicFilters = (StateMap) stateStore.getOrCreateStateCollection(DynamicFilterUtils.MERGED_DYNAMIC_FILTERS, MAP);
        for (Map.Entry<String, DynamicFilterRegistryInfo> columnToDynamicFilterEntry : queryToDynamicFiltersEntry.getValue().entrySet()) {
            if (columnToDynamicFilterEntry.getValue().isMerged()) {
                continue;
            }
            final String filterId = columnToDynamicFilterEntry.getKey();
            final Type filterType = columnToDynamicFilterEntry.getValue().getType();
            final DataType filterDataType = columnToDynamicFilterEntry.getValue().getDataType();
            final Optional<Predicate<List>> dfFilter = columnToDynamicFilterEntry.getValue().getFilter();
            final Symbol column = columnToDynamicFilterEntry.getValue().getSymbol();
            final String filterKey = createKey(DynamicFilterUtils.FILTERPREFIX, filterId, queryId);
            if (!hasMergeCondition(filterId, queryId)) {
                continue;
            }
            Collection<Object> results = ((StateSet) stateStore.getStateCollection(createKey(DynamicFilterUtils.PARTIALPREFIX, filterId, queryId))).getAll();
            try {
                DynamicFilter mergedFilter;
                if (filterDataType == BLOOM_FILTER) {
                    BloomFilter mergedBloomFilter = mergeBloomFilters(results);
                    if (mergedBloomFilter.expectedFpp() > DynamicFilterUtils.BLOOM_FILTER_EXPECTED_FPP) {
                        throw new PrestoException(GENERIC_INTERNAL_ERROR, "FPP too high: " + mergedBloomFilter.approximateElementCount());
                    }
                    mergedFilter = new BloomFilterDynamicFilter(filterKey, null, mergedBloomFilter, filterType);
                    if (filterType == GLOBAL) {
                        try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
                            mergedBloomFilter.writeTo(out);
                            byte[] filter = out.toByteArray();
                            mergedDynamicFilters.put(filterKey, filter);
                        }
                    }
                } else if (filterDataType == HASHSET) {
                    Set mergedSet = mergeHashSets(results);
                    mergedFilter = DynamicFilterFactory.create(filterKey, null, mergedSet, filterType, dfFilter, Optional.empty());
                    if (filterType == GLOBAL) {
                        mergedDynamicFilters.put(filterKey, mergedSet);
                    }
                } else {
                    throw new PrestoException(GENERIC_INTERNAL_ERROR, "Unsupported filter data type: " + filterDataType);
                }
                log.debug("Merged successfully dynamic filter id: " + filterId + "-" + queryId + " type: " + filterDataType + ", column: " + column + ", item count: " + mergedFilter.getSize());
                cachedDynamicFiltersForQuery.put(filterId, mergedFilter);
            } catch (IOException | PrestoException e) {
                log.warn("Could not merge dynamic filter: " + e.getLocalizedMessage());
            } finally {
                // for each dynamic filter we only try to merge it once
                columnToDynamicFilterEntry.getValue().setMerged();
            }
        }
    }
}

Example 19 with StateStore

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

the class DynamicFilterService method removeFinishedQuery.

private void removeFinishedQuery() {
    List<String> handledQuery = new ArrayList<>();
    StateStore stateStore = stateStoreProvider.getStateStore();
    StateMap mergedStateCollection = (StateMap) stateStore.getOrCreateStateCollection(DynamicFilterUtils.MERGED_DYNAMIC_FILTERS, MAP);
    // Clear registered dynamic filter tasks
    synchronized (finishedQuery) {
        for (String queryId : finishedQuery) {
            Map<String, DynamicFilterRegistryInfo> filters = dynamicFilters.get(queryId);
            if (filters != null) {
                for (Entry<String, DynamicFilterRegistryInfo> entry : filters.entrySet()) {
                    String filterId = entry.getKey();
                    clearPartialResults(filterId, queryId);
                    if (entry.getValue().isMerged()) {
                        String filterKey = createKey(DynamicFilterUtils.FILTERPREFIX, filterId, queryId);
                        mergedStateCollection.remove(filterKey);
                    }
                }
            }
            List<String> collectionKeys = stateStore.getStateCollections().keySet().stream().filter(key -> key.contains(queryId)).collect(Collectors.toList());
            for (String key : collectionKeys) {
                clearStatesInStateStore(stateStore, key);
            }
            dynamicFilters.remove(queryId);
            cachedDynamicFilters.remove(queryId);
            handledQuery.add(queryId);
        }
        finishedQuery.removeAll(handledQuery);
    }
}

Example 20 with StateStore

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

the class DynamicFilterService method clearPartialResults.

/**
 * Cleaning local cache and state store cache
 *
 * @param filterId, part of the id for dynamic filter that will be cleaned
 * @param queryId, query id for part of the id of dynamic filter that will be cleaned
 */
private void clearPartialResults(String filterId, String queryId) {
    StateStore stateStore = stateStoreProvider.getStateStore();
    if (stateStore != null) {
        clearStatesInStateStore(stateStore, createKey(DynamicFilterUtils.PARTIALPREFIX, filterId, queryId));
        clearStatesInStateStore(stateStore, createKey(DynamicFilterUtils.TASKSPREFIX, filterId, queryId));
    }
    dynamicFiltersToTask.remove(filterId + "-" + queryId);
}