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

Example 1 with StateMap

use of io.prestosql.spi.statestore.StateMap 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");
}

Example 2 with StateMap

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

the class TestDynamicFilterUtil method setupMockStateStore.

public static StateStore setupMockStateStore(Map mergeMap, Map<String, String> dfTypeMap, Set<String> tasks, Set partial, String queryId, String filterId) {
    StateMap mockMergeMap = mock(StateMap.class);
    StateMap mockDFTypeMap = mock(StateMap.class);
    StateSet mockPartialSet = mock(StateSet.class);
    StateSet mockTasksSet = mock(StateSet.class);
    StateStore stateStore = mock(StateStore.class);
    when(mockMergeMap.put(anyString(), any())).thenAnswer(i -> mergeMap.put(i.getArguments()[0], i.getArguments()[1]));
    when(mockDFTypeMap.put(anyString(), anyString())).thenAnswer(i -> dfTypeMap.put((String) i.getArguments()[0], (String) i.getArguments()[1]));
    when(mockTasksSet.add(anyString())).thenAnswer(i -> tasks.add((String) i.getArguments()[0]));
    when(mockPartialSet.add(any())).thenAnswer(i -> partial.add(i.getArguments()[0]));
    when(mockMergeMap.get(anyString())).thenAnswer(i -> mergeMap.get(i.getArguments()[0]));
    when(mockDFTypeMap.get(anyString())).thenAnswer(i -> dfTypeMap.get(i.getArguments()[0]));
    when(mockMergeMap.getAll()).thenReturn(mergeMap);
    when(mockDFTypeMap.getAll()).thenReturn(dfTypeMap);
    when(mockPartialSet.getAll()).thenReturn(partial);
    when(mockPartialSet.size()).thenAnswer(i -> partial.size());
    when(mockTasksSet.size()).thenAnswer(i -> tasks.size());
    when(stateStore.getStateCollection(DynamicFilterUtils.MERGED_DYNAMIC_FILTERS)).thenReturn(mockMergeMap);
    when(stateStore.createStateCollection(DynamicFilterUtils.MERGED_DYNAMIC_FILTERS, StateCollection.Type.MAP)).thenReturn(mockMergeMap);
    when(stateStore.getOrCreateStateCollection(DynamicFilterUtils.MERGED_DYNAMIC_FILTERS, StateCollection.Type.MAP)).thenReturn(mockMergeMap);
    when(stateStore.getStateCollection(DynamicFilterUtils.createKey(DynamicFilterUtils.TASKSPREFIX, filterId, queryId))).thenReturn(mockTasksSet);
    when(stateStore.createStateCollection(DynamicFilterUtils.createKey(DynamicFilterUtils.TASKSPREFIX, filterId, queryId), StateCollection.Type.SET)).thenReturn(mockTasksSet);
    when(stateStore.getStateCollection(DynamicFilterUtils.createKey(DynamicFilterUtils.PARTIALPREFIX, filterId, queryId))).thenReturn(mockPartialSet);
    when(stateStore.createStateCollection(DynamicFilterUtils.createKey(DynamicFilterUtils.PARTIALPREFIX, filterId, queryId), StateCollection.Type.SET)).thenReturn(mockPartialSet);
    // In statestore, set and map are destroyed and set to null after query finishes, however, in the UT we just assume the set and map to be empty.
    doAnswer(i -> {
        tasks.clear();
        return null;
    }).when(mockTasksSet).destroy();
    doAnswer(i -> {
        partial.clear();
        return null;
    }).when(mockPartialSet).destroy();
    return stateStore;
}

Example 3 with StateMap

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

the class HetuServiceInventory method getServiceDescriptors.

@Override
public Iterable<ServiceDescriptor> getServiceDescriptors(String type) {
    if ("discovery".equals(type)) {
        if (hetuConfig.isMultipleCoordinatorEnabled()) {
            try {
                StateStore stateStore = stateStoreProvider.getStateStore();
                if (stateStore == null) {
                    throw new PrestoException(GENERIC_INTERNAL_ERROR, "State store has not been loaded yet");
                }
                Map.Entry<String, String> entry = ((StateMap<String, String>) stateStore.getStateCollection(StateStoreConstants.DISCOVERY_SERVICE_COLLECTION_NAME)).getAll().entrySet().stream().findFirst().get();
                ImmutableMap.Builder properties = new ImmutableMap.Builder();
                if (internalCommunicationConfig != null && internalCommunicationConfig.isHttpsRequired()) {
                    properties.put("https", "https://" + entry.getKey() + ":" + entry.getValue());
                } else {
                    properties.put("http", "http://" + entry.getKey() + ":" + entry.getValue());
                }
                return ImmutableList.of(new ServiceDescriptor(UUID.randomUUID(), nodeId, "discovery", null, null, ServiceState.RUNNING, properties.build()));
            } catch (Exception e) {
                logServerError("Select service from state store failed: " + e);
                // return an empty list here.
                return ImmutableList.of();
            }
        } else if (seedStoreManager != null && seedStoreManager.isSeedStoreOnYarnEnabled()) {
            try {
                boolean httpsRequired = (internalCommunicationConfig != null && internalCommunicationConfig.isHttpsRequired());
                String location = seedStoreManager.getLatestSeedLocation(SeedStoreSubType.ON_YARN, httpsRequired);
                if (location == null || location.isEmpty()) {
                    throw new PrestoException(GENERIC_INTERNAL_ERROR, "Seed store has not been initialized yet");
                }
                ImmutableMap.Builder properties = new ImmutableMap.Builder();
                if (httpsRequired) {
                    properties.put("https", location);
                } else {
                    properties.put("http", location);
                }
                return ImmutableList.of(new ServiceDescriptor(UUID.randomUUID(), nodeId, "discovery", null, null, ServiceState.RUNNING, properties.build()));
            } catch (Exception e) {
                logServerError("Select service from seed store failed: " + e);
                // return an empty list here.
                return ImmutableList.of();
            }
        }
    }
    return super.getServiceDescriptors(type);
}

Example 4 with StateMap

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

the class StateFetcher method fetchRunningQueryStates.

/**
 * Fetch states only related to running queries from state store to cache store
 * This is used when not all the states are needed for better states fetching performance
 *
 * @throws IOException exception when failed to deserialize states
 */
public void fetchRunningQueryStates(StateStore stateStore) throws IOException {
    synchronized (this) {
        long start = System.currentTimeMillis();
        LOG.debug("fetchStates starts at current time milliseconds: %s, at format HH:mm:ss:SSS:%s", start, new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(start)));
        StateCollection cpuUsageCollection = stateStore.getStateCollection(CPU_USAGE_STATE_COLLECTION_NAME);
        StateCollection queryStateCollection = stateStore.getStateCollection(QUERY_STATE_COLLECTION_NAME);
        StateCacheStore.get().setCachedStates(CPU_USAGE_STATE_COLLECTION_NAME, ((StateMap) cpuUsageCollection).getAll());
        Map<String, String> states = ((StateMap<String, String>) queryStateCollection).getAll();
        StateCacheStore.get().setCachedStates(QUERY_STATE_COLLECTION_NAME, deserializeFetchedStates(states));
        long end = System.currentTimeMillis();
        LOG.debug("fetchStates ends at current time milliseconds: %s, at format HH:mm:ss:SSS:%s, total time use: %s", end, new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(end)), end - start);
    }
}

Example 5 with StateMap

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

the class StateFetcher method fetchAllStates.

/**
 * Fetch all states from state store to cache store
 *
 * @throws IOException exception when failed to deserialize states
 */
public void fetchAllStates() throws IOException {
    synchronized (this) {
        // State store hasn't been loaded yet
        if (stateStoreProvider.getStateStore() == null) {
            return;
        }
        long start = System.currentTimeMillis();
        LOG.debug("fetchStates starts at current time milliseconds: %s, at format HH:mm:ss:SSS:%s", start, new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(start)));
        for (String stateCollectionName : stateCollections) {
            StateCollection stateCollection = stateStoreProvider.getStateStore().getStateCollection(stateCollectionName);
            if (stateCollection == null) {
                continue;
            }
            if (stateCollectionName.equals(CPU_USAGE_STATE_COLLECTION_NAME)) {
                StateCacheStore.get().setCachedStates(stateCollectionName, ((StateMap) stateCollection).getAll());
                continue;
            }
            if (stateCollection.getType() == StateCollection.Type.MAP) {
                Map<String, String> states = ((StateMap<String, String>) stateCollection).getAll();
                StateCacheStore.get().setCachedStates(stateCollectionName, deserializeFetchedStates(states));
            } else {
                LOG.warn("Unsupported state collection type: %s", stateCollection.getType());
            }
        }
        long end = System.currentTimeMillis();
        LOG.debug("fetchStates ends at current time milliseconds: %s, at format HH:mm:ss:SSS:%s, total time use: %s", end, new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(end)), end - start);
    }
}