Search in sources :

Example 1 with OperatorStateInputSplit

use of org.apache.flink.state.api.input.splits.OperatorStateInputSplit in project flink by apache.

the class BroadcastStateInputFormatTest method testReadBroadcastState.

@Test
public void testReadBroadcastState() throws Exception {
    try (TwoInputStreamOperatorTestHarness<Void, Integer, Void> testHarness = getTestHarness()) {
        testHarness.open();
        testHarness.processElement2(new StreamRecord<>(1));
        testHarness.processElement2(new StreamRecord<>(2));
        testHarness.processElement2(new StreamRecord<>(3));
        OperatorSubtaskState subtaskState = testHarness.snapshot(0, 0);
        OperatorState state = new OperatorState(OperatorIDGenerator.fromUid("uid"), 1, 4);
        state.putState(0, subtaskState);
        OperatorStateInputSplit split = new OperatorStateInputSplit(subtaskState.getManagedOperatorState(), 0);
        BroadcastStateInputFormat<Integer, Integer> format = new BroadcastStateInputFormat<>(state, new Configuration(), null, descriptor);
        format.setRuntimeContext(new MockStreamingRuntimeContext(false, 1, 0));
        format.open(split);
        Map<Integer, Integer> results = new HashMap<>(3);
        while (!format.reachedEnd()) {
            Tuple2<Integer, Integer> entry = format.nextRecord(null);
            results.put(entry.f0, entry.f1);
        }
        Map<Integer, Integer> expected = new HashMap<>(3);
        expected.put(1, 1);
        expected.put(2, 2);
        expected.put(3, 3);
        Assert.assertEquals("Failed to read correct list state from state backend", expected, results);
    }
}
Also used : MockStreamingRuntimeContext(org.apache.flink.streaming.util.MockStreamingRuntimeContext) Configuration(org.apache.flink.configuration.Configuration) HashMap(java.util.HashMap) OperatorSubtaskState(org.apache.flink.runtime.checkpoint.OperatorSubtaskState) OperatorStateInputSplit(org.apache.flink.state.api.input.splits.OperatorStateInputSplit) OperatorState(org.apache.flink.runtime.checkpoint.OperatorState) Test(org.junit.Test)

Example 2 with OperatorStateInputSplit

use of org.apache.flink.state.api.input.splits.OperatorStateInputSplit in project flink by apache.

the class ListStateInputFormatTest method testReadListOperatorState.

@Test
public void testReadListOperatorState() throws Exception {
    try (OneInputStreamOperatorTestHarness<Integer, Void> testHarness = getTestHarness()) {
        testHarness.open();
        testHarness.processElement(1, 0);
        testHarness.processElement(2, 0);
        testHarness.processElement(3, 0);
        OperatorSubtaskState subtaskState = testHarness.snapshot(0, 0);
        OperatorState state = new OperatorState(OperatorIDGenerator.fromUid("uid"), 1, 4);
        state.putState(0, subtaskState);
        OperatorStateInputSplit split = new OperatorStateInputSplit(subtaskState.getManagedOperatorState(), 0);
        ListStateInputFormat<Integer> format = new ListStateInputFormat<>(state, new Configuration(), null, descriptor);
        format.setRuntimeContext(new MockStreamingRuntimeContext(false, 1, 0));
        format.open(split);
        List<Integer> results = new ArrayList<>();
        while (!format.reachedEnd()) {
            results.add(format.nextRecord(0));
        }
        results.sort(Comparator.naturalOrder());
        Assert.assertEquals("Failed to read correct list state from state backend", Arrays.asList(1, 2, 3), results);
    }
}
Also used : MockStreamingRuntimeContext(org.apache.flink.streaming.util.MockStreamingRuntimeContext) OperatorStateInputSplit(org.apache.flink.state.api.input.splits.OperatorStateInputSplit) Configuration(org.apache.flink.configuration.Configuration) ArrayList(java.util.ArrayList) OperatorState(org.apache.flink.runtime.checkpoint.OperatorState) OperatorSubtaskState(org.apache.flink.runtime.checkpoint.OperatorSubtaskState) Test(org.junit.Test)

Example 3 with OperatorStateInputSplit

use of org.apache.flink.state.api.input.splits.OperatorStateInputSplit in project flink by apache.

the class UnionStateInputFormatTest method testReadUnionOperatorState.

@Test
public void testReadUnionOperatorState() throws Exception {
    try (OneInputStreamOperatorTestHarness<Integer, Void> testHarness = getTestHarness()) {
        testHarness.open();
        testHarness.processElement(1, 0);
        testHarness.processElement(2, 0);
        testHarness.processElement(3, 0);
        OperatorSubtaskState subtaskState = testHarness.snapshot(0, 0);
        OperatorState state = new OperatorState(OperatorIDGenerator.fromUid("uid"), 1, 4);
        state.putState(0, subtaskState);
        OperatorStateInputSplit split = new OperatorStateInputSplit(subtaskState.getManagedOperatorState(), 0);
        UnionStateInputFormat<Integer> format = new UnionStateInputFormat<>(state, new Configuration(), null, descriptor);
        format.setRuntimeContext(new MockStreamingRuntimeContext(false, 1, 0));
        format.open(split);
        List<Integer> results = new ArrayList<>();
        while (!format.reachedEnd()) {
            results.add(format.nextRecord(0));
        }
        results.sort(Comparator.naturalOrder());
        Assert.assertEquals("Failed to read correct list state from state backend", Arrays.asList(1, 2, 3), results);
    }
}
Also used : MockStreamingRuntimeContext(org.apache.flink.streaming.util.MockStreamingRuntimeContext) OperatorStateInputSplit(org.apache.flink.state.api.input.splits.OperatorStateInputSplit) Configuration(org.apache.flink.configuration.Configuration) ArrayList(java.util.ArrayList) OperatorState(org.apache.flink.runtime.checkpoint.OperatorState) OperatorSubtaskState(org.apache.flink.runtime.checkpoint.OperatorSubtaskState) Test(org.junit.Test)

Aggregations

Configuration (org.apache.flink.configuration.Configuration)3 OperatorState (org.apache.flink.runtime.checkpoint.OperatorState)3 OperatorSubtaskState (org.apache.flink.runtime.checkpoint.OperatorSubtaskState)3 OperatorStateInputSplit (org.apache.flink.state.api.input.splits.OperatorStateInputSplit)3 MockStreamingRuntimeContext (org.apache.flink.streaming.util.MockStreamingRuntimeContext)3 Test (org.junit.Test)3 ArrayList (java.util.ArrayList)2 HashMap (java.util.HashMap)1