This is an automated email from the ASF dual-hosted git repository. guoweijie pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/flink.git
commit 9a47e375d97d5895205edef0b967d5e3595b2d03 Author: Jeyhun Karimov <[email protected]> AuthorDate: Wed Jun 5 21:41:20 2024 +0200 [FLINK-34977][API] Introduce State Access on DataStream API V2 --- .../flink/datastream/api/context/StateManager.java | 67 +++++++ .../impl/context/DefaultPartitionedContext.java | 10 +- .../impl/context/DefaultStateManager.java | 110 ++++++++++- .../datastream/impl/operators/ProcessOperator.java | 7 +- .../TwoInputBroadcastProcessOperator.java | 7 +- .../TwoInputNonBroadcastProcessOperator.java | 10 +- .../impl/operators/TwoOutputProcessOperator.java | 9 +- .../context/DefaultNonPartitionedContextTest.java | 10 +- .../impl/context/DefaultStateManagerTest.java | 219 ++++++++++++++++++++- .../DefaultTwoOutputNonPartitionedContextTest.java | 10 +- .../operators/MockFreqCountProcessFunction.java | 61 ++++++ .../MockGlobalDecuplicateCountProcessFunction.java | 67 +++++++ .../MockGlobalListAppenderProcessFunction.java | 78 ++++++++ .../operators/MockListAppenderProcessFunction.java | 77 ++++++++ .../operators/MockMultiplierProcessFunction.java | 59 ++++++ .../MockRecudingMultiplierProcessFunction.java | 68 +++++++ .../operators/MockSumAggregateProcessFunction.java | 83 ++++++++ 17 files changed, 938 insertions(+), 14 deletions(-) diff --git a/flink-datastream-api/src/main/java/org/apache/flink/datastream/api/context/StateManager.java b/flink-datastream-api/src/main/java/org/apache/flink/datastream/api/context/StateManager.java index 38b03e5a2cf..e92de2856b0 100644 --- a/flink-datastream-api/src/main/java/org/apache/flink/datastream/api/context/StateManager.java +++ b/flink-datastream-api/src/main/java/org/apache/flink/datastream/api/context/StateManager.java @@ -19,6 +19,20 @@ package org.apache.flink.datastream.api.context; import org.apache.flink.annotation.Experimental; +import org.apache.flink.api.common.state.AggregatingState; +import org.apache.flink.api.common.state.AggregatingStateDeclaration; +import org.apache.flink.api.common.state.BroadcastState; +import org.apache.flink.api.common.state.BroadcastStateDeclaration; +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.state.ListStateDeclaration; +import org.apache.flink.api.common.state.MapState; +import org.apache.flink.api.common.state.MapStateDeclaration; +import org.apache.flink.api.common.state.ReducingState; +import org.apache.flink.api.common.state.ReducingStateDeclaration; +import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.common.state.ValueStateDeclaration; + +import java.util.Optional; /** This is responsibility for managing runtime information related to state of process function. */ @Experimental @@ -31,4 +45,57 @@ public interface StateManager { * instance, get the key from a non-keyed partition stream. */ <K> K getCurrentKey() throws UnsupportedOperationException; + + /** + * Get the specific list state. + * + * @param stateDeclaration of this state. + * @return the list state corresponds to the state declaration. + */ + <T> Optional<ListState<T>> getState(ListStateDeclaration<T> stateDeclaration) throws Exception; + + /** + * Get the specific value state. + * + * @param stateDeclaration of this state. + * @return the value state corresponds to the state declaration. + */ + <T> Optional<ValueState<T>> getState(ValueStateDeclaration<T> stateDeclaration) + throws Exception; + + /** + * Get the specific map state. + * + * @param stateDeclaration of this state. + * @return the map state corresponds to the state declaration. + */ + <K, V> Optional<MapState<K, V>> getState(MapStateDeclaration<K, V> stateDeclaration) + throws Exception; + + /** + * Get the specific reducing state. + * + * @param stateDeclaration of this state. + * @return the reducing state corresponds to the state declaration. + */ + <T> Optional<ReducingState<T>> getState(ReducingStateDeclaration<T> stateDeclaration) + throws Exception; + + /** + * Get the specific aggregating state. + * + * @param stateDeclaration of this state. + * @return the aggregating state corresponds to the state declaration. + */ + <IN, ACC, OUT> Optional<AggregatingState<IN, OUT>> getState( + AggregatingStateDeclaration<IN, ACC, OUT> stateDeclaration) throws Exception; + + /** + * Get the specific broadcast state. + * + * @param stateDeclaration of this state. + * @return the broadcast state corresponds to the state declaration. + */ + <K, V> Optional<BroadcastState<K, V>> getState(BroadcastStateDeclaration<K, V> stateDeclaration) + throws Exception; } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultPartitionedContext.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultPartitionedContext.java index 61fc9501213..b3c7e4b6323 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultPartitionedContext.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultPartitionedContext.java @@ -18,12 +18,14 @@ package org.apache.flink.datastream.impl.context; +import org.apache.flink.api.common.state.OperatorStateStore; import org.apache.flink.datastream.api.context.JobInfo; import org.apache.flink.datastream.api.context.PartitionedContext; import org.apache.flink.datastream.api.context.ProcessingTimeManager; import org.apache.flink.datastream.api.context.RuntimeContext; import org.apache.flink.datastream.api.context.TaskInfo; import org.apache.flink.metrics.MetricGroup; +import org.apache.flink.streaming.api.operators.StreamingRuntimeContext; import java.util.function.Consumer; import java.util.function.Supplier; @@ -40,9 +42,13 @@ public class DefaultPartitionedContext implements PartitionedContext { RuntimeContext context, Supplier<Object> currentKeySupplier, Consumer<Object> currentKeySetter, - ProcessingTimeManager processingTimeManager) { + ProcessingTimeManager processingTimeManager, + StreamingRuntimeContext operatorContext, + OperatorStateStore operatorStateStore) { this.context = context; - this.stateManager = new DefaultStateManager(currentKeySupplier, currentKeySetter); + this.stateManager = + new DefaultStateManager( + currentKeySupplier, currentKeySetter, operatorContext, operatorStateStore); this.processingTimeManager = processingTimeManager; } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultStateManager.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultStateManager.java index a6ce75361f4..bcc74520403 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultStateManager.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultStateManager.java @@ -18,8 +18,30 @@ package org.apache.flink.datastream.impl.context; +import org.apache.flink.api.common.state.AggregatingState; +import org.apache.flink.api.common.state.AggregatingStateDeclaration; +import org.apache.flink.api.common.state.AggregatingStateDescriptor; +import org.apache.flink.api.common.state.BroadcastState; +import org.apache.flink.api.common.state.BroadcastStateDeclaration; +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.state.ListStateDeclaration; +import org.apache.flink.api.common.state.ListStateDescriptor; +import org.apache.flink.api.common.state.MapState; +import org.apache.flink.api.common.state.MapStateDeclaration; +import org.apache.flink.api.common.state.MapStateDescriptor; +import org.apache.flink.api.common.state.OperatorStateStore; +import org.apache.flink.api.common.state.ReducingState; +import org.apache.flink.api.common.state.ReducingStateDeclaration; +import org.apache.flink.api.common.state.ReducingStateDescriptor; +import org.apache.flink.api.common.state.StateDeclaration; +import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.common.state.ValueStateDeclaration; +import org.apache.flink.api.common.state.ValueStateDescriptor; import org.apache.flink.datastream.api.context.StateManager; +import org.apache.flink.streaming.api.operators.StreamingRuntimeContext; +import org.apache.flink.util.Preconditions; +import java.util.Optional; import java.util.function.Consumer; import java.util.function.Supplier; @@ -37,10 +59,19 @@ public class DefaultStateManager implements StateManager { private final Consumer<Object> currentKeySetter; + protected final StreamingRuntimeContext operatorContext; + + protected final OperatorStateStore operatorStateStore; + public DefaultStateManager( - Supplier<Object> currentKeySupplier, Consumer<Object> currentKeySetter) { + Supplier<Object> currentKeySupplier, + Consumer<Object> currentKeySetter, + StreamingRuntimeContext operatorContext, + OperatorStateStore operatorStateStore) { this.currentKeySupplier = currentKeySupplier; this.currentKeySetter = currentKeySetter; + this.operatorContext = Preconditions.checkNotNull(operatorContext); + this.operatorStateStore = Preconditions.checkNotNull(operatorStateStore); } @Override @@ -49,6 +80,83 @@ public class DefaultStateManager implements StateManager { return (K) currentKeySupplier.get(); } + @Override + public <T> Optional<ValueState<T>> getState(ValueStateDeclaration<T> stateDeclaration) + throws Exception { + ValueStateDescriptor<T> valueStateDescriptor = + new ValueStateDescriptor<>( + stateDeclaration.getName(), + stateDeclaration.getTypeDescriptor().getTypeClass()); + return Optional.ofNullable(operatorContext.getState(valueStateDescriptor)); + } + + @Override + public <T> Optional<ListState<T>> getState(ListStateDeclaration<T> stateDeclaration) + throws Exception { + + ListStateDescriptor<T> listStateDescriptor = + new ListStateDescriptor<>( + stateDeclaration.getName(), + stateDeclaration.getTypeDescriptor().getTypeClass()); + + if (stateDeclaration.getRedistributionMode() + == StateDeclaration.RedistributionMode.REDISTRIBUTABLE) { + if (stateDeclaration.getRedistributionStrategy() + == ListStateDeclaration.RedistributionStrategy.UNION) { + return Optional.ofNullable( + operatorStateStore.getUnionListState(listStateDescriptor)); + } else { + return Optional.ofNullable(operatorStateStore.getListState(listStateDescriptor)); + } + } else { + return Optional.ofNullable(operatorContext.getListState(listStateDescriptor)); + } + } + + @Override + public <K, V> Optional<MapState<K, V>> getState(MapStateDeclaration<K, V> stateDeclaration) + throws Exception { + MapStateDescriptor<K, V> mapStateDescriptor = + new MapStateDescriptor<>( + stateDeclaration.getName(), + stateDeclaration.getKeyTypeDescriptor().getTypeClass(), + stateDeclaration.getValueTypeDescriptor().getTypeClass()); + return Optional.ofNullable(operatorContext.getMapState(mapStateDescriptor)); + } + + @Override + public <T> Optional<ReducingState<T>> getState(ReducingStateDeclaration<T> stateDeclaration) + throws Exception { + ReducingStateDescriptor<T> reducingStateDescriptor = + new ReducingStateDescriptor<>( + stateDeclaration.getName(), + stateDeclaration.getReduceFunction(), + stateDeclaration.getTypeDescriptor().getTypeClass()); + return Optional.ofNullable(operatorContext.getReducingState(reducingStateDescriptor)); + } + + @Override + public <IN, ACC, OUT> Optional<AggregatingState<IN, OUT>> getState( + AggregatingStateDeclaration<IN, ACC, OUT> stateDeclaration) throws Exception { + AggregatingStateDescriptor<IN, ACC, OUT> aggregatingStateDescriptor = + new AggregatingStateDescriptor<>( + stateDeclaration.getName(), + stateDeclaration.getAggregateFunction(), + stateDeclaration.getTypeDescriptor().getTypeClass()); + return Optional.ofNullable(operatorContext.getAggregatingState(aggregatingStateDescriptor)); + } + + @Override + public <K, V> Optional<BroadcastState<K, V>> getState( + BroadcastStateDeclaration<K, V> stateDeclaration) throws Exception { + MapStateDescriptor<K, V> mapStateDescriptor = + new MapStateDescriptor<>( + stateDeclaration.getName(), + stateDeclaration.getKeyTypeDescriptor().getTypeClass(), + stateDeclaration.getValueTypeDescriptor().getTypeClass()); + return Optional.ofNullable(operatorStateStore.getBroadcastState(mapStateDescriptor)); + } + /** * This method should be used to run a block of code with a specific key context. The original * key must be reset after the block is executed. diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/ProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/ProcessOperator.java index 368b5bed247..2511b5db740 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/ProcessOperator.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/ProcessOperator.java @@ -69,7 +69,12 @@ public class ProcessOperator<IN, OUT> operatorContext.getMetricGroup()); partitionedContext = new DefaultPartitionedContext( - context, this::currentKey, this::setCurrentKey, getProcessingTimeManager()); + context, + this::currentKey, + this::setCurrentKey, + getProcessingTimeManager(), + operatorContext, + getOperatorStateBackend()); outputCollector = getOutputCollector(); nonPartitionedContext = getNonPartitionedContext(); } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperator.java index 4e6f663adf8..a6e34a1daca 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperator.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperator.java @@ -73,7 +73,12 @@ public class TwoInputBroadcastProcessOperator<IN1, IN2, OUT> operatorContext.getMetricGroup()); this.partitionedContext = new DefaultPartitionedContext( - context, this::currentKey, this::setCurrentKey, getProcessingTimeManager()); + context, + this::currentKey, + this::setCurrentKey, + getProcessingTimeManager(), + operatorContext, + getOperatorStateBackend()); this.nonPartitionedContext = getNonPartitionedContext(); } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperator.java index 14eed8ad419..31cc7d48dab 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperator.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperator.java @@ -28,6 +28,7 @@ import org.apache.flink.datastream.impl.context.DefaultNonPartitionedContext; import org.apache.flink.datastream.impl.context.DefaultPartitionedContext; import org.apache.flink.datastream.impl.context.DefaultRuntimeContext; import org.apache.flink.datastream.impl.context.UnsupportedProcessingTimeManager; +import org.apache.flink.runtime.state.OperatorStateBackend; import org.apache.flink.streaming.api.operators.AbstractUdfStreamOperator; import org.apache.flink.streaming.api.operators.BoundedMultiInput; import org.apache.flink.streaming.api.operators.ChainingStrategy; @@ -62,6 +63,8 @@ public class TwoInputNonBroadcastProcessOperator<IN1, IN2, OUT> super.open(); this.collector = getOutputCollector(); StreamingRuntimeContext operatorContext = getRuntimeContext(); + OperatorStateBackend operatorStateBackend = getOperatorStateBackend(); + TaskInfo taskInfo = operatorContext.getTaskInfo(); this.context = new DefaultRuntimeContext( @@ -73,7 +76,12 @@ public class TwoInputNonBroadcastProcessOperator<IN1, IN2, OUT> operatorContext.getMetricGroup()); this.partitionedContext = new DefaultPartitionedContext( - context, this::currentKey, this::setCurrentKey, getProcessingTimeManager()); + context, + this::currentKey, + this::setCurrentKey, + getProcessingTimeManager(), + operatorContext, + operatorStateBackend); this.nonPartitionedContext = getNonPartitionedContext(); } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperator.java index 43681f4a1b5..af24d1839cd 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperator.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperator.java @@ -19,6 +19,7 @@ package org.apache.flink.datastream.impl.operators; import org.apache.flink.api.common.TaskInfo; +import org.apache.flink.api.common.state.OperatorStateStore; import org.apache.flink.datastream.api.context.ProcessingTimeManager; import org.apache.flink.datastream.api.context.TwoOutputNonPartitionedContext; import org.apache.flink.datastream.api.function.TwoOutputStreamProcessFunction; @@ -72,6 +73,7 @@ public class TwoOutputProcessOperator<IN, OUT_MAIN, OUT_SIDE> this.mainCollector = getMainCollector(); this.sideCollector = getSideCollector(); StreamingRuntimeContext operatorContext = getRuntimeContext(); + OperatorStateStore operatorStateStore = getOperatorStateBackend(); TaskInfo taskInfo = operatorContext.getTaskInfo(); this.context = new DefaultRuntimeContext( @@ -83,7 +85,12 @@ public class TwoOutputProcessOperator<IN, OUT_MAIN, OUT_SIDE> operatorContext.getMetricGroup()); this.partitionedContext = new DefaultPartitionedContext( - context, this::currentKey, this::setCurrentKey, getProcessingTimeManager()); + context, + this::currentKey, + this::setCurrentKey, + getProcessingTimeManager(), + operatorContext, + operatorStateStore); this.nonPartitionedContext = getNonPartitionedContext(); } diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultNonPartitionedContextTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultNonPartitionedContextTest.java index 1bd740f6453..b4fa8c62531 100644 --- a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultNonPartitionedContextTest.java +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultNonPartitionedContextTest.java @@ -20,6 +20,7 @@ package org.apache.flink.datastream.impl.context; import org.apache.flink.datastream.impl.common.TestingTimestampCollector; import org.apache.flink.streaming.api.operators.StreamingRuntimeContext; +import org.apache.flink.streaming.api.operators.collect.utils.MockOperatorStateStore; import org.junit.jupiter.api.Test; @@ -55,6 +56,7 @@ class DefaultNonPartitionedContextTest { 2, "mock-task", operatorRuntimeContext.getMetricGroup()); + DefaultNonPartitionedContext<Integer> nonPartitionedContext = new DefaultNonPartitionedContext<>( runtimeContext, @@ -62,7 +64,9 @@ class DefaultNonPartitionedContextTest { runtimeContext, Optional::empty, (key) -> cf.complete(null), - UnsupportedProcessingTimeManager.INSTANCE), + UnsupportedProcessingTimeManager.INSTANCE, + ContextTestUtils.createStreamingRuntimeContext(), + new MockOperatorStateStore()), collector, false, null); @@ -109,7 +113,9 @@ class DefaultNonPartitionedContextTest { runtimeContext, currentKey::get, (key) -> currentKey.set((Integer) key), - UnsupportedProcessingTimeManager.INSTANCE), + UnsupportedProcessingTimeManager.INSTANCE, + ContextTestUtils.createStreamingRuntimeContext(), + new MockOperatorStateStore()), collector, true, allKeys); diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultStateManagerTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultStateManagerTest.java index 7528d8c0b16..9be18a26768 100644 --- a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultStateManagerTest.java +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultStateManagerTest.java @@ -18,8 +18,25 @@ package org.apache.flink.datastream.impl.context; +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.datastream.impl.operators.KeyedProcessOperator; +import org.apache.flink.datastream.impl.operators.MockFreqCountProcessFunction; +import org.apache.flink.datastream.impl.operators.MockGlobalDecuplicateCountProcessFunction; +import org.apache.flink.datastream.impl.operators.MockGlobalListAppenderProcessFunction; +import org.apache.flink.datastream.impl.operators.MockListAppenderProcessFunction; +import org.apache.flink.datastream.impl.operators.MockMultiplierProcessFunction; +import org.apache.flink.datastream.impl.operators.MockRecudingMultiplierProcessFunction; +import org.apache.flink.datastream.impl.operators.MockSumAggregateProcessFunction; +import org.apache.flink.streaming.api.operators.collect.utils.MockOperatorStateStore; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness; +import org.apache.flink.streaming.util.MockStreamingRuntimeContext; + import org.junit.jupiter.api.Test; +import java.util.Collection; +import java.util.List; import java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; @@ -30,7 +47,12 @@ class DefaultStateManagerTest { @Test void testGetCurrentKey() { final String key = "key"; - DefaultStateManager stateManager = new DefaultStateManager(() -> key, ignore -> {}); + DefaultStateManager stateManager = + new DefaultStateManager( + () -> key, + ignore -> {}, + new MockStreamingRuntimeContext(false, 1, 0), + new MockOperatorStateStore()); assertThat((String) stateManager.getCurrentKey()).isEqualTo(key); } @@ -41,7 +63,9 @@ class DefaultStateManagerTest { () -> { throw new RuntimeException("Expected Error"); }, - ignore -> {}); + ignore -> {}, + new MockStreamingRuntimeContext(false, 1, 0), + new MockOperatorStateStore()); assertThatThrownBy(stateManager::getCurrentKey) .isInstanceOf(RuntimeException.class) .hasMessageContaining("Expected Error"); @@ -54,8 +78,197 @@ class DefaultStateManagerTest { // -1 as unset value AtomicInteger setKey = new AtomicInteger(-1); DefaultStateManager stateManager = - new DefaultStateManager(() -> oldKey, k -> setKey.set((Integer) k)); + new DefaultStateManager( + () -> oldKey, + k -> setKey.set((Integer) k), + new MockStreamingRuntimeContext(false, 1, 0), + new MockOperatorStateStore()); stateManager.executeInKeyContext(() -> assertThat(setKey).hasValue(newKey), newKey); assertThat(setKey).hasValue(oldKey); } + + @Test + void testListState() throws Exception { + MockListAppenderProcessFunction function = new MockListAppenderProcessFunction(); + KeyedProcessOperator<Integer, Integer, Integer> processOperator = + new KeyedProcessOperator<>( + function, (KeySelector<Integer, Integer>) value -> value); + + try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Integer>) value -> value, + Types.INT)) { + testHarness.open(); + testHarness.processElement(new StreamRecord<>(1)); + testHarness.processElement(new StreamRecord<>(1)); + testHarness.processElement(new StreamRecord<>(1)); + testHarness.processElement(new StreamRecord<>(1)); + + List<Integer> listState = function.getResultingState(); + assertThat(listState).containsExactly(1, 1, 1, 1); + } + } + + @Test + void testAggState() throws Exception { + MockSumAggregateProcessFunction function = new MockSumAggregateProcessFunction(); + KeyedProcessOperator<Integer, Integer, Integer> processOperator = + new KeyedProcessOperator<>(function); + + try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Integer>) value -> value, + Types.INT)) { + testHarness.open(); + testHarness.processElement(new StreamRecord<>(1)); + testHarness.processElement(new StreamRecord<>(1)); + + testHarness.processElement(new StreamRecord<>(2)); + testHarness.processElement(new StreamRecord<>(2)); + + Collection<StreamRecord<Integer>> recordOutput = testHarness.getRecordOutput(); + assertThat(recordOutput) + .containsExactly( + new StreamRecord<>(1), + new StreamRecord<>(2), + new StreamRecord<>(2), + new StreamRecord<>(4)); + } + } + + @Test + void testValueState() throws Exception { + MockMultiplierProcessFunction function = new MockMultiplierProcessFunction(); + KeyedProcessOperator<Integer, Integer, Integer> processOperator = + new KeyedProcessOperator<>(function); + + try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Integer>) value -> value, + Types.INT)) { + testHarness.open(); + testHarness.processElement(new StreamRecord<>(3)); + testHarness.processElement(new StreamRecord<>(3)); + testHarness.processElement(new StreamRecord<>(4)); + testHarness.processElement(new StreamRecord<>(4)); + + Collection<StreamRecord<Integer>> recordOutput = testHarness.getRecordOutput(); + assertThat(recordOutput) + .containsExactly( + new StreamRecord<>(3), + new StreamRecord<>(9), + new StreamRecord<>(4), + new StreamRecord<>(16)); + } + } + + @Test + void testMapState() throws Exception { + MockFreqCountProcessFunction function = new MockFreqCountProcessFunction(); + KeyedProcessOperator<Integer, Integer, Integer> processOperator = + new KeyedProcessOperator<>(function); + + try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Integer>) value -> value, + Types.INT)) { + testHarness.open(); + testHarness.processElement(new StreamRecord<>(1)); + testHarness.processElement(new StreamRecord<>(1)); + testHarness.processElement(new StreamRecord<>(1)); + + testHarness.processElement(new StreamRecord<>(3)); + Collection<StreamRecord<Integer>> recordOutput = testHarness.getRecordOutput(); + assertThat(recordOutput) + .containsExactly( + new StreamRecord<>(1), + new StreamRecord<>(2), + new StreamRecord<>(3), + new StreamRecord<>(1)); + } + } + + @Test + void testReducingState() throws Exception { + MockRecudingMultiplierProcessFunction function = + new MockRecudingMultiplierProcessFunction(); + KeyedProcessOperator<Integer, Integer, Integer> processOperator = + new KeyedProcessOperator<>(function); + + try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Integer>) value -> value, + Types.INT)) { + testHarness.open(); + testHarness.processElement(new StreamRecord<>(2)); + testHarness.processElement(new StreamRecord<>(2)); + testHarness.processElement(new StreamRecord<>(3)); + testHarness.processElement(new StreamRecord<>(2)); + + Collection<StreamRecord<Integer>> recordOutput = testHarness.getRecordOutput(); + assertThat(recordOutput) + .containsExactly( + new StreamRecord<>(2), + new StreamRecord<>(4), + new StreamRecord<>(3), + new StreamRecord<>(8)); + } + } + + @Test + void testBroadcastMapState() throws Exception { + MockGlobalDecuplicateCountProcessFunction function = + new MockGlobalDecuplicateCountProcessFunction(); + KeyedProcessOperator<Integer, Integer, Integer> processOperator = + new KeyedProcessOperator<>(function); + + try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Integer>) value -> value, + Types.INT)) { + testHarness.open(); + testHarness.processElement(new StreamRecord<>(1)); + testHarness.processElement(new StreamRecord<>(2)); + testHarness.processElement(new StreamRecord<>(3)); + testHarness.processElement(new StreamRecord<>(4)); + + Collection<StreamRecord<Integer>> recordOutput = testHarness.getRecordOutput(); + assertThat(recordOutput) + .containsExactly( + new StreamRecord<>(1), + new StreamRecord<>(2), + new StreamRecord<>(3), + new StreamRecord<>(4)); + } + } + + @Test + void testBroadcastListState() throws Exception { + MockGlobalListAppenderProcessFunction function = + new MockGlobalListAppenderProcessFunction(); + KeyedProcessOperator<Integer, Integer, Integer> processOperator = + new KeyedProcessOperator<>( + function, (KeySelector<Integer, Integer>) value -> value); + + try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Integer>) value -> value, + Types.INT)) { + testHarness.open(); + testHarness.processElement(new StreamRecord<>(1)); + testHarness.processElement(new StreamRecord<>(2)); + testHarness.processElement(new StreamRecord<>(3)); + testHarness.processElement(new StreamRecord<>(4)); + + List<Integer> listState = function.getResultingState(); + assertThat(listState).containsExactly(1, 2, 3, 4); + } + } } diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultTwoOutputNonPartitionedContextTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultTwoOutputNonPartitionedContextTest.java index fbad9f4bcba..137527b6542 100644 --- a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultTwoOutputNonPartitionedContextTest.java +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultTwoOutputNonPartitionedContextTest.java @@ -20,6 +20,7 @@ package org.apache.flink.datastream.impl.context; import org.apache.flink.datastream.impl.common.TestingTimestampCollector; import org.apache.flink.streaming.api.operators.StreamingRuntimeContext; +import org.apache.flink.streaming.api.operators.collect.utils.MockOperatorStateStore; import org.junit.jupiter.api.Test; @@ -67,11 +68,14 @@ class DefaultTwoOutputNonPartitionedContextTest { runtimeContext, Optional::empty, (key) -> cf.complete(null), - UnsupportedProcessingTimeManager.INSTANCE), + UnsupportedProcessingTimeManager.INSTANCE, + ContextTestUtils.createStreamingRuntimeContext(), + new MockOperatorStateStore()), firstCollector, secondCollector, false, null); + nonPartitionedContext.applyToAllPartitions( (firstOutput, secondOutput, ctx) -> { counter.incrementAndGet(); @@ -122,7 +126,9 @@ class DefaultTwoOutputNonPartitionedContextTest { runtimeContext, currentKey::get, (key) -> currentKey.set((Integer) key), - UnsupportedProcessingTimeManager.INSTANCE), + UnsupportedProcessingTimeManager.INSTANCE, + ContextTestUtils.createStreamingRuntimeContext(), + new MockOperatorStateStore()), firstCollector, secondCollector, true, diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockFreqCountProcessFunction.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockFreqCountProcessFunction.java new file mode 100644 index 00000000000..8bbd1e6d9b6 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockFreqCountProcessFunction.java @@ -0,0 +1,61 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.datastream.impl.operators; + +import org.apache.flink.api.common.state.MapState; +import org.apache.flink.api.common.state.MapStateDeclaration; +import org.apache.flink.api.common.state.StateDeclaration; +import org.apache.flink.api.common.state.StateDeclarations; +import org.apache.flink.api.common.typeinfo.TypeDescriptors; +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.PartitionedContext; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; + +import java.util.Collections; +import java.util.HashSet; +import java.util.Optional; +import java.util.Set; + +public class MockFreqCountProcessFunction + implements OneInputStreamProcessFunction<Integer, Integer> { + + private final MapStateDeclaration<Integer, Integer> mapStateDeclaration = + StateDeclarations.mapState("map-state", TypeDescriptors.INT, TypeDescriptors.INT); + + @Override + public Set<StateDeclaration> usesStates() { + return new HashSet<>(Collections.singletonList(mapStateDeclaration)); + } + + @Override + public void processRecord(Integer record, Collector<Integer> output, PartitionedContext ctx) + throws Exception { + Optional<MapState<Integer, Integer>> stateOptional = + ctx.getStateManager().getState(mapStateDeclaration); + if (!stateOptional.isPresent()) { + throw new RuntimeException("State is not available"); + } + MapState<Integer, Integer> state = stateOptional.get(); + Integer oldFreq = state.get(record); + Integer newFreq = oldFreq == null ? 1 : oldFreq + 1; + state.put(record, newFreq); + + output.collect(newFreq); + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockGlobalDecuplicateCountProcessFunction.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockGlobalDecuplicateCountProcessFunction.java new file mode 100644 index 00000000000..aca147b3ae5 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockGlobalDecuplicateCountProcessFunction.java @@ -0,0 +1,67 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.datastream.impl.operators; + +import org.apache.flink.api.common.state.BroadcastState; +import org.apache.flink.api.common.state.BroadcastStateDeclaration; +import org.apache.flink.api.common.state.StateDeclaration; +import org.apache.flink.api.common.state.StateDeclarations; +import org.apache.flink.api.common.typeinfo.TypeDescriptors; +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.PartitionedContext; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; + +import java.util.Collections; +import java.util.HashSet; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +public class MockGlobalDecuplicateCountProcessFunction + implements OneInputStreamProcessFunction<Integer, Integer> { + + private final BroadcastStateDeclaration<Integer, Integer> broadcastStateDeclaration = + StateDeclarations.mapStateBuilder( + "broadcast-state", TypeDescriptors.INT, TypeDescriptors.INT) + .buildBroadcast(); + + @Override + public Set<StateDeclaration> usesStates() { + return new HashSet<>(Collections.singletonList(broadcastStateDeclaration)); + } + + @Override + public void processRecord(Integer record, Collector<Integer> output, PartitionedContext ctx) + throws Exception { + Optional<BroadcastState<Integer, Integer>> stateOptional = + ctx.getStateManager().getState(broadcastStateDeclaration); + if (!stateOptional.isPresent()) { + throw new RuntimeException("State is not available"); + } + BroadcastState<Integer, Integer> state = stateOptional.get(); + state.put(record, record); + + int len = 0; + for (Map.Entry<Integer, Integer> entry : state.entries()) { + len++; + } + + output.collect(len); + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockGlobalListAppenderProcessFunction.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockGlobalListAppenderProcessFunction.java new file mode 100644 index 00000000000..413af0827eb --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockGlobalListAppenderProcessFunction.java @@ -0,0 +1,78 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.datastream.impl.operators; + +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.state.ListStateDeclaration; +import org.apache.flink.api.common.state.StateDeclaration; +import org.apache.flink.api.common.state.StateDeclarations; +import org.apache.flink.api.common.typeinfo.TypeDescriptors; +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.PartitionedContext; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Optional; +import java.util.Set; + +public class MockGlobalListAppenderProcessFunction + implements OneInputStreamProcessFunction<Integer, Integer> { + + private final ListStateDeclaration<Integer> listStateDeclaration = + StateDeclarations.listStateBuilder("list-state", TypeDescriptors.INT) + .redistributeWithMode(StateDeclaration.RedistributionMode.IDENTICAL) + .redistributeBy(ListStateDeclaration.RedistributionStrategy.SPLIT) + .build(); + + private List<Integer> resultingListState = new ArrayList<>(); + + @Override + public Set<StateDeclaration> usesStates() { + return new HashSet<>(Collections.singletonList(listStateDeclaration)); + } + + public List<Integer> getResultingState() { + return resultingListState; + } + + private void updateResultingState(ListState<Integer> state) throws Exception { + resultingListState = new ArrayList<>(); + for (Integer integer : state.get()) { + resultingListState.add(integer); + } + } + + @Override + public void processRecord(Integer record, Collector<Integer> output, PartitionedContext ctx) + throws Exception { + Optional<ListState<Integer>> stateOptional = + ctx.getStateManager().getState(listStateDeclaration); + if (!stateOptional.isPresent()) { + throw new RuntimeException("State is not available"); + } + ListState<Integer> state = stateOptional.get(); + state.add(record); + updateResultingState(state); + // forward the record to check input key. + output.collect(record); + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockListAppenderProcessFunction.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockListAppenderProcessFunction.java new file mode 100644 index 00000000000..eca3b295598 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockListAppenderProcessFunction.java @@ -0,0 +1,77 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.datastream.impl.operators; + +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.state.ListStateDeclaration; +import org.apache.flink.api.common.state.StateDeclaration; +import org.apache.flink.api.common.state.StateDeclarations; +import org.apache.flink.api.common.typeinfo.TypeDescriptors; +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.PartitionedContext; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Optional; +import java.util.Set; + +public class MockListAppenderProcessFunction + implements OneInputStreamProcessFunction<Integer, Integer> { + + private final ListStateDeclaration<Integer> listStateDeclaration = + StateDeclarations.listStateBuilder("list-state", TypeDescriptors.INT) + .redistributeBy(ListStateDeclaration.RedistributionStrategy.SPLIT) + .build(); + + private List<Integer> resultingListState = new ArrayList<>(); + + @Override + public Set<StateDeclaration> usesStates() { + return new HashSet<>(Collections.singletonList(listStateDeclaration)); + } + + public List<Integer> getResultingState() { + return resultingListState; + } + + private void updateResultingState(ListState<Integer> state) throws Exception { + resultingListState = new ArrayList<>(); + for (Integer integer : state.get()) { + resultingListState.add(integer); + } + } + + @Override + public void processRecord(Integer record, Collector<Integer> output, PartitionedContext ctx) + throws Exception { + Optional<ListState<Integer>> stateOptional = + ctx.getStateManager().getState(listStateDeclaration); + if (!stateOptional.isPresent()) { + throw new RuntimeException("State is not available"); + } + ListState<Integer> state = stateOptional.get(); + state.add(record); + updateResultingState(state); + // forward the record to check input key. + output.collect(record); + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockMultiplierProcessFunction.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockMultiplierProcessFunction.java new file mode 100644 index 00000000000..050844ba2f9 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockMultiplierProcessFunction.java @@ -0,0 +1,59 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.datastream.impl.operators; + +import org.apache.flink.api.common.state.StateDeclaration; +import org.apache.flink.api.common.state.StateDeclarations; +import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.common.state.ValueStateDeclaration; +import org.apache.flink.api.common.typeinfo.TypeDescriptors; +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.PartitionedContext; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; + +import java.util.Collections; +import java.util.HashSet; +import java.util.Optional; +import java.util.Set; + +public class MockMultiplierProcessFunction + implements OneInputStreamProcessFunction<Integer, Integer> { + + private final ValueStateDeclaration<Integer> valueStateDeclaration = + StateDeclarations.valueState("value-state", TypeDescriptors.INT); + + @Override + public Set<StateDeclaration> usesStates() { + return new HashSet<>(Collections.singletonList(valueStateDeclaration)); + } + + @Override + public void processRecord(Integer record, Collector<Integer> output, PartitionedContext ctx) + throws Exception { + Optional<ValueState<Integer>> stateOptional = + ctx.getStateManager().getState(valueStateDeclaration); + if (!stateOptional.isPresent()) { + throw new RuntimeException("State is not available"); + } + ValueState<Integer> state = stateOptional.get(); + Integer res = (state.value() == null ? 1 : state.value()) * record; + state.update(res); + output.collect(state.value()); + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockRecudingMultiplierProcessFunction.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockRecudingMultiplierProcessFunction.java new file mode 100644 index 00000000000..f777c8a09e0 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockRecudingMultiplierProcessFunction.java @@ -0,0 +1,68 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.datastream.impl.operators; + +import org.apache.flink.api.common.functions.ReduceFunction; +import org.apache.flink.api.common.state.ReducingState; +import org.apache.flink.api.common.state.ReducingStateDeclaration; +import org.apache.flink.api.common.state.StateDeclaration; +import org.apache.flink.api.common.state.StateDeclarations; +import org.apache.flink.api.common.typeinfo.TypeDescriptors; +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.PartitionedContext; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; + +import java.util.Collections; +import java.util.HashSet; +import java.util.Optional; +import java.util.Set; + +public class MockRecudingMultiplierProcessFunction + implements OneInputStreamProcessFunction<Integer, Integer> { + + private final ReducingStateDeclaration<Integer> reducingStateDeclaration = + StateDeclarations.reducingState( + "reducing-state", + TypeDescriptors.INT, + new ReduceFunction<Integer>() { + @Override + public Integer reduce(Integer value1, Integer value2) throws Exception { + return value2 * value1; + } + }); + + @Override + public Set<StateDeclaration> usesStates() { + return new HashSet<>(Collections.singletonList(reducingStateDeclaration)); + } + + @Override + public void processRecord(Integer record, Collector<Integer> output, PartitionedContext ctx) + throws Exception { + Optional<ReducingState<Integer>> stateOptional = + ctx.getStateManager().getState(reducingStateDeclaration); + if (!stateOptional.isPresent()) { + throw new RuntimeException("State is not available"); + } + ReducingState<Integer> state = stateOptional.get(); + state.add(record); + + output.collect(state.get()); + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockSumAggregateProcessFunction.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockSumAggregateProcessFunction.java new file mode 100644 index 00000000000..7447a000809 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/MockSumAggregateProcessFunction.java @@ -0,0 +1,83 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.datastream.impl.operators; + +import org.apache.flink.api.common.functions.AggregateFunction; +import org.apache.flink.api.common.state.AggregatingState; +import org.apache.flink.api.common.state.AggregatingStateDeclaration; +import org.apache.flink.api.common.state.StateDeclaration; +import org.apache.flink.api.common.state.StateDeclarations; +import org.apache.flink.api.common.typeinfo.TypeDescriptors; +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.PartitionedContext; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; + +import java.util.Collections; +import java.util.HashSet; +import java.util.Optional; +import java.util.Set; + +public class MockSumAggregateProcessFunction + implements OneInputStreamProcessFunction<Integer, Integer> { + + private final AggregatingStateDeclaration<Integer, Integer, Integer> + aggregatingStateDeclaration = + StateDeclarations.aggregatingState( + "agg-state", + TypeDescriptors.INT, + new AggregateFunction<Integer, Integer, Integer>() { + @Override + public Integer createAccumulator() { + return 0; + } + + @Override + public Integer add(Integer value, Integer accumulator) { + return value + accumulator; + } + + @Override + public Integer getResult(Integer accumulator) { + return accumulator; + } + + @Override + public Integer merge(Integer a, Integer b) { + return a + b; + } + }); + + @Override + public Set<StateDeclaration> usesStates() { + return new HashSet<>(Collections.singletonList(aggregatingStateDeclaration)); + } + + @Override + public void processRecord(Integer record, Collector<Integer> output, PartitionedContext ctx) + throws Exception { + Optional<AggregatingState<Integer, Integer>> stateOptional = + ctx.getStateManager().getState(aggregatingStateDeclaration); + if (!stateOptional.isPresent()) { + throw new RuntimeException("State is not available"); + } + AggregatingState<Integer, Integer> state = stateOptional.get(); + state.add(record); + output.collect(state.get()); + } +}
