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

Example 21 with StateStore

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

the class SplitCacheStateInitializer method updateLocal.

private void updateLocal() {
    StateStore stateStore = provider.getStateStore();
    if (stateStore == null) {
        log.debug("State store not yet initialized. Will retry after %s milli seconds until %s", delay.toMillis(), timeout.toString(TimeUnit.SECONDS));
        return;
    }
    if (initializationState.get() == InitializationStatus.COMPLETED) {
        log.debug("Split cache map already initialized.");
        return;
    }
    // create state store collection
    if (stateStore.getStateCollection(StateStoreConstants.SPLIT_CACHE_METADATA_NAME) == null) {
        stateStore.createStateMap(StateStoreConstants.SPLIT_CACHE_METADATA_NAME, new SplitCacheStateStoreChangesListener(SplitCacheMap.getInstance(), mapper));
    }
    StateMap<String, String> stateMap = (StateMap<String, String>) stateStore.getStateCollection(StateStoreConstants.SPLIT_CACHE_METADATA_NAME);
    stateMap.getAll().forEach((fqTableName, stateJson) -> {
        try {
            TableCacheInfo cacheInfo = mapper.readerFor(TableCacheInfo.class).readValue(stateJson);
            log.info("Retrieving cache info for table %s from state store and updating on local copy.", fqTableName);
            splitCacheMap.setTableCacheInfo(fqTableName, cacheInfo);
        } catch (Exception e) {
            log.error(e, "Unable to update local split cache map from state store.");
        }
    });
    initializationState.set(InitializationStatus.COMPLETED);
}

Example 22 with StateStore

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

the class QueryTracker method removeQueryInStateStore.

private void removeQueryInStateStore(QueryId queryId) {
    StateStore stateStore = stateStoreProvider.getStateStore();
    if (stateStore == null) {
        return;
    }
    StateCollection stateCollection = stateStore.getStateCollection(StateStoreConstants.FINISHED_QUERY_STATE_COLLECTION_NAME);
    if (stateCollection != null && stateCollection.getType().equals(StateCollection.Type.MAP)) {
        ((StateMap<String, String>) stateCollection).remove(queryId.getId());
    }
}

Example 23 with StateStore

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

the class DynamicFilterService method hasMergeCondition.

private boolean hasMergeCondition(String filterKey, String queryId) {
    int finishedDynamicFilterNumber = 0;
    final StateStore stateStore = stateStoreProvider.getStateStore();
    StateCollection temp = stateStore.getStateCollection(createKey(DynamicFilterUtils.TASKSPREFIX, filterKey, queryId));
    if (temp != null) {
        finishedDynamicFilterNumber = temp.size();
    }
    return finishedDynamicFilterNumber > 0 && finishedDynamicFilterNumber == dynamicFiltersToTask.get(filterKey + "-" + queryId).size();
}

Example 24 with StateStore

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

the class LocalDynamicFilter method addPartialFilterToStateStore.

private void addPartialFilterToStateStore() {
    StateStore stateStore = stateStoreProvider.getStateStore();
    if (stateStore == null) {
        return;
    }
    DynamicFilter.DataType dataType = getDynamicFilterDataType(type, dynamicFilterDataType);
    for (Map.Entry<String, Set> filter : result.entrySet()) {
        DynamicFilterSourceOperator.Channel channel = channels.get(filter.getKey());
        Set filterValues = filter.getValue();
        String filterId = channel.getFilterId();
        String key = createKey(PARTIALPREFIX, filterId, channel.getQueryId());
        if (dataType == BLOOM_FILTER) {
            byte[] finalOutput = convertBloomFilterToByteArray(createBloomFilterFromSet(channel, filterValues, bloomFilterFpp));
            if (finalOutput != null) {
                ((StateSet) stateStore.getOrCreateStateCollection(key, SET)).add(finalOutput);
            }
        } else {
            ((StateSet) stateStore.getOrCreateStateCollection(key, SET)).add(filterValues);
        }
        ((StateSet) stateStore.getOrCreateStateCollection(createKey(TASKSPREFIX, filterId, channel.getQueryId()), SET)).add(taskId.toString());
        log.debug("creating new " + dataType + " dynamic filter for size of: " + result.size() + ", key: " + key + ", taskId: " + taskId);
    }
}

Example 25 with StateStore

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

the class TestStateUpdater method setUp.

@BeforeMethod
public void setUp() {
    updateInterval = new Duration(MINIMUM_UPDATE_INTERVAL, MILLISECONDS);
    stateStore = Mockito.mock(StateStore.class);
}