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

Example 1 with StateSet

use of io.prestosql.spi.statestore.StateSet 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 2 with StateSet

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

the class TestDynamicFilterServiceWithBloomFilter method mockLocalDynamicFilter.

private void mockLocalDynamicFilter(String taskId, String filterId, String queryId, List<String> values) {
    BloomFilter bloomFilter = new BloomFilter(1024 * 1024, 0.1);
    for (String val : values) {
        bloomFilter.add(val.getBytes(StandardCharsets.UTF_8));
    }
    String key = DynamicFilterUtils.createKey(DynamicFilterUtils.PARTIALPREFIX, filterId, queryId);
    try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
        bloomFilter.writeTo(out);
        byte[] finalOutput = out.toByteArray();
        ((StateSet) stateStoreProvider.getStateStore().getStateCollection(key)).add(finalOutput);
        ((StateSet) stateStoreProvider.getStateStore().getStateCollection(DynamicFilterUtils.createKey(DynamicFilterUtils.TASKSPREFIX, filterId, queryId))).add(taskId);
    } catch (IOException e) {
        Assert.fail("could not register finish filter, Exception happened:" + e.getMessage());
    }
}

Example 3 with StateSet

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

the class TestDynamicFilterSourceOperator method testCollectMultipleColumns.

@Test
public void testCollectMultipleColumns() {
    String filterId1 = "multiple_columns_1";
    String filterId2 = "multiple_columns_2";
    OperatorFactory operatorFactory = createOperatorFactory(LOCAL, HASHSET, 1, channel(0, BOOLEAN, filterId1), channel(1, DOUBLE, filterId2));
    verifyPassthrough(createOperator((DynamicFilterSourceOperatorFactory) operatorFactory), ImmutableList.of(BOOLEAN, DOUBLE), new Page(createBooleansBlock(true, 2), createDoublesBlock(1.5, 3.0)), new Page(createBooleansBlock(false, 1), createDoublesBlock(4.5)));
    operatorFactory.noMoreOperators();
    String key1 = DynamicFilterUtils.createKey(PARTIALPREFIX, filterId1, TEST_SESSION.getQueryId().toString());
    Set<Boolean> set1 = new HashSet<>(Arrays.asList(true, false));
    StateSet states1 = ((StateSet) stateStoreProvider.getStateStore().getStateCollection(key1));
    for (Object bfSerialized : states1.getAll()) {
        assertEquals((Set) bfSerialized, set1);
    }
    String key2 = DynamicFilterUtils.createKey(PARTIALPREFIX, filterId2, TEST_SESSION.getQueryId().toString());
    Set<Double> set2 = new HashSet<>(Arrays.asList(1.5, 3.0, 4.5));
    StateSet states2 = ((StateSet) stateStoreProvider.getStateStore().getStateCollection(key2));
    for (Object bfSerialized : states2.getAll()) {
        assertEquals((Set) bfSerialized, set2);
    }
}

Example 4 with StateSet

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

the class TestDynamicFilterSourceOperator method testCollectDeduplication.

@Test
public void testCollectDeduplication() {
    String filterId = "deduplication";
    final int maxRowCount = getDynamicFilteringMaxPerDriverValueCount(pipelineContext.getSession());
    // lots of zeros
    Page largePage = new Page(createLongRepeatBlock(7, maxRowCount * 10));
    // lots of nulls
    Page nullsPage = new Page(createLongsBlock(Arrays.asList(new Long[maxRowCount * 10])));
    OperatorFactory operatorFactory = createOperatorFactory(LOCAL, BLOOM_FILTER, 1, channel(0, BIGINT, filterId));
    verifyPassthrough(createOperator((DynamicFilterSourceOperatorFactory) operatorFactory), ImmutableList.of(BIGINT), largePage, nullsPage);
    operatorFactory.noMoreOperators();
    String key = DynamicFilterUtils.createKey(PARTIALPREFIX, filterId, TEST_SESSION.getQueryId().toString());
    Set<Long> set = new HashSet<>();
    set.add(7L);
    StateSet states = ((StateSet) stateStoreProvider.getStateStore().getStateCollection(key));
    for (Object bfSerialized : states.getAll()) {
        assertEquals((Set) bfSerialized, set);
    }
}

Example 5 with StateSet

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

the class TestDynamicFilterSourceOperator method testCollectMultipleOperators.

@Test
public void testCollectMultipleOperators() {
    String filterId = "0";
    DynamicFilterSourceOperatorFactory operatorFactory = createOperatorFactory(LOCAL, HASHSET, 2, channel(0, BIGINT, filterId));
    // will finish before noMoreOperators()
    DynamicFilterSourceOperator op1 = createOperator(operatorFactory);
    verifyPassthrough(op1, ImmutableList.of(BIGINT), new Page(createLongsBlock(1, 2)), new Page(createLongsBlock(3, 5)));
    // will finish after noMoreOperators()
    Operator op2 = createOperator(operatorFactory);
    operatorFactory.noMoreOperators();
    verifyPassthrough(op2, ImmutableList.of(BIGINT), new Page(createLongsBlock(2, 3)), new Page(createLongsBlock(1, 4)));
    String key = DynamicFilterUtils.createKey(PARTIALPREFIX, filterId, TEST_SESSION.getQueryId().toString());
    Set<Long> set = new HashSet<>(Arrays.asList(1L, 2L, 3L, 4L, 5L));
    StateSet states = ((StateSet) stateStoreProvider.getStateStore().getStateCollection(key));
    for (Object bfSerialized : states.getAll()) {
        assertEquals(set, (Set) bfSerialized);
    }
}