Examples with ListState org.apache.flink.api.common.state.ListState used on opensource projects

Example 11 with ListState

use of org.apache.flink.api.common.state.ListState in project beam by apache.

the class BufferingDoFnRunnerTest method createBufferingDoFnRunner.

private static BufferingDoFnRunner createBufferingDoFnRunner(int concurrentCheckpoints, List<BufferingDoFnRunner.CheckpointIdentifier> notYetAcknowledgeCheckpoints) throws Exception {
    DoFnRunner doFnRunner = Mockito.mock(DoFnRunner.class);
    OperatorStateBackend operatorStateBackend = Mockito.mock(OperatorStateBackend.class);
    // Setup not yet acknowledged checkpoint union list state
    ListState unionListState = Mockito.mock(ListState.class);
    Mockito.when(operatorStateBackend.getUnionListState(Mockito.any())).thenReturn(unionListState);
    Mockito.when(unionListState.get()).thenReturn(notYetAcknowledgeCheckpoints);
    // Setup buffer list state
    Mockito.when(operatorStateBackend.getListState(Mockito.any())).thenReturn(Mockito.mock(ListState.class));
    return BufferingDoFnRunner.create(doFnRunner, "stable-input", StringUtf8Coder.of(), WindowedValue.getFullCoder(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE), operatorStateBackend, null, concurrentCheckpoints, new SerializablePipelineOptions(FlinkPipelineOptions.defaults()));
}

Example 12 with ListState

use of org.apache.flink.api.common.state.ListState in project flink by apache.

the class StateBackendTestBase method testConcurrentMapIfQueryable.

/**
 * Tests that {@link AbstractHeapState} instances respect the queryable flag and create
 * concurrent variants for internal state structures.
 */
@SuppressWarnings("unchecked")
protected void testConcurrentMapIfQueryable() throws Exception {
    final int numberOfKeyGroups = 1;
    final CheckpointableKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE, numberOfKeyGroups, new KeyGroupRange(0, 0), new DummyEnvironment());
    try {
        {
            // ValueState
            ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("value-state", Integer.class, -1);
            desc.setQueryable("my-query");
            ValueState<Integer> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
            InternalKvState<Integer, VoidNamespace, Integer> kvState = (InternalKvState<Integer, VoidNamespace, Integer>) state;
            assertTrue(kvState instanceof AbstractHeapState);
            kvState.setCurrentNamespace(VoidNamespace.INSTANCE);
            backend.setCurrentKey(1);
            state.update(121818273);
            assertNotNull("State not set", ((AbstractHeapState<?, ?, ?>) kvState).getStateTable());
        }
        {
            // ListState
            ListStateDescriptor<Integer> desc = new ListStateDescriptor<>("list-state", Integer.class);
            desc.setQueryable("my-query");
            ListState<Integer> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
            InternalKvState<Integer, VoidNamespace, Integer> kvState = (InternalKvState<Integer, VoidNamespace, Integer>) state;
            assertTrue(kvState instanceof AbstractHeapState);
            kvState.setCurrentNamespace(VoidNamespace.INSTANCE);
            backend.setCurrentKey(1);
            state.add(121818273);
            assertNotNull("State not set", ((AbstractHeapState<?, ?, ?>) kvState).getStateTable());
        }
        {
            // ReducingState
            ReducingStateDescriptor<Integer> desc = new ReducingStateDescriptor<>("reducing-state", new ReduceFunction<Integer>() {

                @Override
                public Integer reduce(Integer value1, Integer value2) throws Exception {
                    return value1 + value2;
                }
            }, Integer.class);
            desc.setQueryable("my-query");
            ReducingState<Integer> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
            InternalKvState<Integer, VoidNamespace, Integer> kvState = (InternalKvState<Integer, VoidNamespace, Integer>) state;
            assertTrue(kvState instanceof AbstractHeapState);
            kvState.setCurrentNamespace(VoidNamespace.INSTANCE);
            backend.setCurrentKey(1);
            state.add(121818273);
            assertNotNull("State not set", ((AbstractHeapState<?, ?, ?>) kvState).getStateTable());
        }
        {
            // MapState
            MapStateDescriptor<Integer, String> desc = new MapStateDescriptor<>("map-state", Integer.class, String.class);
            desc.setQueryable("my-query");
            MapState<Integer, String> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
            InternalKvState<Integer, VoidNamespace, Map<Integer, String>> kvState = (InternalKvState<Integer, VoidNamespace, Map<Integer, String>>) state;
            assertTrue(kvState instanceof AbstractHeapState);
            kvState.setCurrentNamespace(VoidNamespace.INSTANCE);
            backend.setCurrentKey(1);
            state.put(121818273, "121818273");
            int keyGroupIndex = KeyGroupRangeAssignment.assignToKeyGroup(1, numberOfKeyGroups);
            StateTable stateTable = ((AbstractHeapState) kvState).getStateTable();
            assertNotNull("State not set", stateTable.get(keyGroupIndex));
        }
    } finally {
        IOUtils.closeQuietly(backend);
        backend.dispose();
    }
}