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 2a6ee6488cb4d8f3c72635aae62d5f445ef270fa
Author: Jeyhun Karimov <[email protected]>
AuthorDate: Wed Jun 5 22:15:14 2024 +0200

    [FLINK-34977][API] Support compile-time checks for invalid states
---
 .../state/IllegalRedistributionModeException.java  |  31 ++++++
 .../impl/stream/BroadcastStreamImpl.java           |  13 +++
 .../datastream/impl/stream/GlobalStreamImpl.java   |  24 +++++
 .../impl/stream/KeyedPartitionStreamImpl.java      |  41 ++++++++
 .../impl/stream/NonKeyedPartitionStreamImpl.java   |  32 ++++++
 .../flink/datastream/impl/utils/StreamUtils.java   |  18 ++++
 .../impl/stream/BroadcastStreamImplTest.java       |  58 +++++++++++
 .../impl/stream/GlobalStreamImplTest.java          |  79 +++++++++++++++
 .../impl/stream/KeyedPartitionStreamImplTest.java  |  85 ++++++++++++++++
 .../stream/NonKeyedPartitionStreamImplTest.java    | 110 +++++++++++++++++++++
 .../datastream/impl/stream/StreamTestUtils.java    |  67 +++++++++++++
 11 files changed, 558 insertions(+)

diff --git 
a/flink-core-api/src/main/java/org/apache/flink/api/common/state/IllegalRedistributionModeException.java
 
b/flink-core-api/src/main/java/org/apache/flink/api/common/state/IllegalRedistributionModeException.java
new file mode 100644
index 00000000000..4b832f09923
--- /dev/null
+++ 
b/flink-core-api/src/main/java/org/apache/flink/api/common/state/IllegalRedistributionModeException.java
@@ -0,0 +1,31 @@
+/*
+ * 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.api.common.state;
+
+import org.apache.flink.annotation.Experimental;
+
+/** A special {@link IllegalStateException} indicating a mismatch in {@link 
StateDeclaration}. */
+@Experimental
+public class IllegalRedistributionModeException extends IllegalStateException {
+    private static final long serialVersionUID = 1L;
+
+    public 
IllegalRedistributionModeException(StateDeclaration.RedistributionMode mode) {
+        super(String.format("Unexpected Redistribution Mode %s", mode));
+    }
+}
diff --git 
a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/BroadcastStreamImpl.java
 
b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/BroadcastStreamImpl.java
index 670bed82a96..defd870be68 100644
--- 
a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/BroadcastStreamImpl.java
+++ 
b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/BroadcastStreamImpl.java
@@ -18,6 +18,7 @@
 
 package org.apache.flink.datastream.impl.stream;
 
+import org.apache.flink.api.common.state.StateDeclaration;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.dag.Transformation;
 import org.apache.flink.api.java.functions.KeySelector;
@@ -35,6 +36,11 @@ import org.apache.flink.datastream.impl.utils.StreamUtils;
 import org.apache.flink.streaming.api.transformations.PartitionTransformation;
 import org.apache.flink.streaming.runtime.partitioner.BroadcastPartitioner;
 
+import java.util.Collections;
+import java.util.HashSet;
+
+import static 
org.apache.flink.datastream.impl.utils.StreamUtils.validateStates;
+
 /** The implementation of {@link BroadcastStream}. */
 public class BroadcastStreamImpl<T> extends AbstractDataStream<T> implements 
BroadcastStream<T> {
     public BroadcastStreamImpl(
@@ -53,6 +59,9 @@ public class BroadcastStreamImpl<T> extends 
AbstractDataStream<T> implements Bro
     public <K, T_OTHER, OUT> 
ProcessConfigurableAndNonKeyedPartitionStream<OUT> connectAndProcess(
             KeyedPartitionStream<K, T_OTHER> other,
             TwoInputBroadcastStreamProcessFunction<T_OTHER, T, OUT> 
processFunction) {
+        // no state redistribution mode check is required here, since all 
redistribution modes are
+        // acceptable
+
         TypeInformation<OUT> outTypeInfo =
                 StreamUtils.getOutputTypeForTwoInputBroadcastProcessFunction(
                         processFunction,
@@ -77,6 +86,10 @@ public class BroadcastStreamImpl<T> extends 
AbstractDataStream<T> implements Bro
     public <T_OTHER, OUT> ProcessConfigurableAndNonKeyedPartitionStream<OUT> 
connectAndProcess(
             NonKeyedPartitionStream<T_OTHER> other,
             TwoInputBroadcastStreamProcessFunction<T_OTHER, T, OUT> 
processFunction) {
+        validateStates(
+                processFunction.usesStates(),
+                new 
HashSet<>(Collections.singletonList(StateDeclaration.RedistributionMode.NONE)));
+
         TypeInformation<OUT> outTypeInfo =
                 StreamUtils.getOutputTypeForTwoInputBroadcastProcessFunction(
                         processFunction,
diff --git 
a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/GlobalStreamImpl.java
 
b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/GlobalStreamImpl.java
index f2d6a038496..8dddc935318 100644
--- 
a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/GlobalStreamImpl.java
+++ 
b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/GlobalStreamImpl.java
@@ -18,6 +18,7 @@
 
 package org.apache.flink.datastream.impl.stream;
 
+import org.apache.flink.api.common.state.StateDeclaration;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.connector.dsv2.Sink;
 import org.apache.flink.api.dag.Transformation;
@@ -44,6 +45,12 @@ import 
org.apache.flink.streaming.api.transformations.PartitionTransformation;
 import org.apache.flink.streaming.runtime.partitioner.ShufflePartitioner;
 import org.apache.flink.util.OutputTag;
 
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+
+import static 
org.apache.flink.datastream.impl.utils.StreamUtils.validateStates;
+
 /** The implementation of {@link GlobalStream}. */
 public class GlobalStreamImpl<T> extends AbstractDataStream<T> implements 
GlobalStream<T> {
     public GlobalStreamImpl(
@@ -54,6 +61,11 @@ public class GlobalStreamImpl<T> extends 
AbstractDataStream<T> implements Global
     @Override
     public <OUT> ProcessConfigurableAndGlobalStream<OUT> process(
             OneInputStreamProcessFunction<T, OUT> processFunction) {
+        validateStates(
+                processFunction.usesStates(),
+                new HashSet<>(
+                        
Collections.singletonList(StateDeclaration.RedistributionMode.IDENTICAL)));
+
         TypeInformation<OUT> outType =
                 
StreamUtils.getOutputTypeForOneInputProcessFunction(processFunction, getType());
         ProcessOperator<T, OUT> operator = new 
ProcessOperator<>(processFunction);
@@ -63,6 +75,11 @@ public class GlobalStreamImpl<T> extends 
AbstractDataStream<T> implements Global
     @Override
     public <OUT1, OUT2> TwoGlobalStreams<OUT1, OUT2> process(
             TwoOutputStreamProcessFunction<T, OUT1, OUT2> processFunction) {
+        validateStates(
+                processFunction.usesStates(),
+                new HashSet<>(
+                        
Collections.singletonList(StateDeclaration.RedistributionMode.IDENTICAL)));
+
         Tuple2<TypeInformation<OUT1>, TypeInformation<OUT2>> twoOutputType =
                 
StreamUtils.getOutputTypesForTwoOutputProcessFunction(processFunction, 
getType());
         TypeInformation<OUT1> firstOutputType = twoOutputType.f0;
@@ -83,6 +100,13 @@ public class GlobalStreamImpl<T> extends 
AbstractDataStream<T> implements Global
     public <T_OTHER, OUT> ProcessConfigurableAndGlobalStream<OUT> 
connectAndProcess(
             GlobalStream<T_OTHER> other,
             TwoInputNonBroadcastStreamProcessFunction<T, T_OTHER, OUT> 
processFunction) {
+        validateStates(
+                processFunction.usesStates(),
+                new HashSet<>(
+                        Arrays.asList(
+                                StateDeclaration.RedistributionMode.NONE,
+                                
StateDeclaration.RedistributionMode.IDENTICAL)));
+
         TypeInformation<OUT> outTypeInfo =
                 
StreamUtils.getOutputTypeForTwoInputNonBroadcastProcessFunction(
                         processFunction, getType(), 
((GlobalStreamImpl<T_OTHER>) other).getType());
diff --git 
a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/KeyedPartitionStreamImpl.java
 
b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/KeyedPartitionStreamImpl.java
index ae477b7e65c..85b6c1531cd 100644
--- 
a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/KeyedPartitionStreamImpl.java
+++ 
b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/KeyedPartitionStreamImpl.java
@@ -18,6 +18,7 @@
 
 package org.apache.flink.datastream.impl.stream;
 
+import org.apache.flink.api.common.state.StateDeclaration;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.connector.dsv2.Sink;
 import org.apache.flink.api.dag.Transformation;
@@ -49,6 +50,10 @@ import 
org.apache.flink.streaming.runtime.partitioner.KeyGroupStreamPartitioner;
 import org.apache.flink.streaming.runtime.partitioner.ShufflePartitioner;
 import org.apache.flink.util.OutputTag;
 
+import java.util.Collections;
+import java.util.HashSet;
+
+import static 
org.apache.flink.datastream.impl.utils.StreamUtils.validateStates;
 import static org.apache.flink.util.Preconditions.checkNotNull;
 
 /** The implementation of {@link KeyedPartitionStream}. */
@@ -102,6 +107,11 @@ public class KeyedPartitionStreamImpl<K, V> extends 
AbstractDataStream<V>
     @Override
     public <OUT> ProcessConfigurableAndNonKeyedPartitionStream<OUT> process(
             OneInputStreamProcessFunction<V, OUT> processFunction) {
+        validateStates(
+                processFunction.usesStates(),
+                new HashSet<>(
+                        
Collections.singletonList(StateDeclaration.RedistributionMode.IDENTICAL)));
+
         TypeInformation<OUT> outType;
         outType = 
StreamUtils.getOutputTypeForOneInputProcessFunction(processFunction, getType());
 
@@ -118,6 +128,11 @@ public class KeyedPartitionStreamImpl<K, V> extends 
AbstractDataStream<V>
     public <OUT> ProcessConfigurableAndKeyedPartitionStream<K, OUT> process(
             OneInputStreamProcessFunction<V, OUT> processFunction,
             KeySelector<OUT, K> newKeySelector) {
+        validateStates(
+                processFunction.usesStates(),
+                new HashSet<>(
+                        
Collections.singletonList(StateDeclaration.RedistributionMode.IDENTICAL)));
+
         TypeInformation<OUT> outType =
                 
StreamUtils.getOutputTypeForOneInputProcessFunction(processFunction, getType());
         KeyedProcessOperator<K, V, OUT> operator =
@@ -142,6 +157,11 @@ public class KeyedPartitionStreamImpl<K, V> extends 
AbstractDataStream<V>
             TwoOutputStreamProcessFunction<V, OUT1, OUT2> processFunction,
             KeySelector<OUT1, K> keySelector1,
             KeySelector<OUT2, K> keySelector2) {
+        validateStates(
+                processFunction.usesStates(),
+                new HashSet<>(
+                        
Collections.singletonList(StateDeclaration.RedistributionMode.IDENTICAL)));
+
         Tuple2<TypeInformation<OUT1>, TypeInformation<OUT2>> twoOutputType =
                 
StreamUtils.getOutputTypesForTwoOutputProcessFunction(processFunction, 
getType());
         TypeInformation<OUT1> firstOutputType = twoOutputType.f0;
@@ -188,6 +208,11 @@ public class KeyedPartitionStreamImpl<K, V> extends 
AbstractDataStream<V>
     @Override
     public <OUT1, OUT2> TwoNonKeyedPartitionStreams<OUT1, OUT2> process(
             TwoOutputStreamProcessFunction<V, OUT1, OUT2> processFunction) {
+        validateStates(
+                processFunction.usesStates(),
+                new HashSet<>(
+                        
Collections.singletonList(StateDeclaration.RedistributionMode.IDENTICAL)));
+
         Tuple2<TypeInformation<OUT1>, TypeInformation<OUT2>> twoOutputType =
                 
StreamUtils.getOutputTypesForTwoOutputProcessFunction(processFunction, 
getType());
         TypeInformation<OUT1> firstOutputType = twoOutputType.f0;
@@ -217,6 +242,11 @@ public class KeyedPartitionStreamImpl<K, V> extends 
AbstractDataStream<V>
     public <T_OTHER, OUT> ProcessConfigurableAndNonKeyedPartitionStream<OUT> 
connectAndProcess(
             KeyedPartitionStream<K, T_OTHER> other,
             TwoInputNonBroadcastStreamProcessFunction<V, T_OTHER, OUT> 
processFunction) {
+        validateStates(
+                processFunction.usesStates(),
+                new HashSet<>(
+                        
Collections.singletonList(StateDeclaration.RedistributionMode.IDENTICAL)));
+
         TypeInformation<OUT> outTypeInfo =
                 
StreamUtils.getOutputTypeForTwoInputNonBroadcastProcessFunction(
                         processFunction,
@@ -242,6 +272,11 @@ public class KeyedPartitionStreamImpl<K, V> extends 
AbstractDataStream<V>
             KeyedPartitionStream<K, T_OTHER> other,
             TwoInputNonBroadcastStreamProcessFunction<V, T_OTHER, OUT> 
processFunction,
             KeySelector<OUT, K> newKeySelector) {
+        validateStates(
+                processFunction.usesStates(),
+                new HashSet<>(
+                        
Collections.singletonList(StateDeclaration.RedistributionMode.IDENTICAL)));
+
         TypeInformation<OUT> outTypeInfo =
                 
StreamUtils.getOutputTypeForTwoInputNonBroadcastProcessFunction(
                         processFunction,
@@ -274,6 +309,9 @@ public class KeyedPartitionStreamImpl<K, V> extends 
AbstractDataStream<V>
     public <T_OTHER, OUT> ProcessConfigurableAndNonKeyedPartitionStream<OUT> 
connectAndProcess(
             BroadcastStream<T_OTHER> other,
             TwoInputBroadcastStreamProcessFunction<V, T_OTHER, OUT> 
processFunction) {
+        // no state redistribution mode check is required here, since all 
redistribution modes are
+        // acceptable
+
         TypeInformation<OUT> outTypeInfo =
                 StreamUtils.getOutputTypeForTwoInputBroadcastProcessFunction(
                         processFunction,
@@ -299,6 +337,9 @@ public class KeyedPartitionStreamImpl<K, V> extends 
AbstractDataStream<V>
             BroadcastStream<T_OTHER> other,
             TwoInputBroadcastStreamProcessFunction<V, T_OTHER, OUT> 
processFunction,
             KeySelector<OUT, K> newKeySelector) {
+        // no state redistribution mode check is required here, since all 
redistribution modes are
+        // acceptable
+
         TypeInformation<OUT> outTypeInfo =
                 StreamUtils.getOutputTypeForTwoInputBroadcastProcessFunction(
                         processFunction,
diff --git 
a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/NonKeyedPartitionStreamImpl.java
 
b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/NonKeyedPartitionStreamImpl.java
index 4c2b4f13bee..dda90d290be 100644
--- 
a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/NonKeyedPartitionStreamImpl.java
+++ 
b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/NonKeyedPartitionStreamImpl.java
@@ -18,6 +18,7 @@
 
 package org.apache.flink.datastream.impl.stream;
 
+import org.apache.flink.api.common.state.StateDeclaration;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.connector.dsv2.Sink;
 import org.apache.flink.api.dag.Transformation;
@@ -45,6 +46,12 @@ import 
org.apache.flink.streaming.runtime.partitioner.GlobalPartitioner;
 import org.apache.flink.streaming.runtime.partitioner.ShufflePartitioner;
 import org.apache.flink.util.OutputTag;
 
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+
+import static 
org.apache.flink.datastream.impl.utils.StreamUtils.validateStates;
+
 /** The implementation of {@link NonKeyedPartitionStream}. */
 public class NonKeyedPartitionStreamImpl<T> extends AbstractDataStream<T>
         implements NonKeyedPartitionStream<T> {
@@ -56,6 +63,13 @@ public class NonKeyedPartitionStreamImpl<T> extends 
AbstractDataStream<T>
     @Override
     public <OUT> ProcessConfigurableAndNonKeyedPartitionStream<OUT> process(
             OneInputStreamProcessFunction<T, OUT> processFunction) {
+        validateStates(
+                processFunction.usesStates(),
+                new HashSet<>(
+                        Arrays.asList(
+                                StateDeclaration.RedistributionMode.NONE,
+                                
StateDeclaration.RedistributionMode.IDENTICAL)));
+
         TypeInformation<OUT> outType =
                 
StreamUtils.getOutputTypeForOneInputProcessFunction(processFunction, getType());
         ProcessOperator<T, OUT> operator = new 
ProcessOperator<>(processFunction);
@@ -69,6 +83,13 @@ public class NonKeyedPartitionStreamImpl<T> extends 
AbstractDataStream<T>
     @Override
     public <OUT1, OUT2> TwoNonKeyedPartitionStreams<OUT1, OUT2> process(
             TwoOutputStreamProcessFunction<T, OUT1, OUT2> processFunction) {
+        validateStates(
+                processFunction.usesStates(),
+                new HashSet<>(
+                        Arrays.asList(
+                                StateDeclaration.RedistributionMode.NONE,
+                                
StateDeclaration.RedistributionMode.IDENTICAL)));
+
         Tuple2<TypeInformation<OUT1>, TypeInformation<OUT2>> twoOutputType =
                 
StreamUtils.getOutputTypesForTwoOutputProcessFunction(processFunction, 
getType());
         TypeInformation<OUT1> firstOutputType = twoOutputType.f0;
@@ -93,6 +114,13 @@ public class NonKeyedPartitionStreamImpl<T> extends 
AbstractDataStream<T>
     public <T_OTHER, OUT> ProcessConfigurableAndNonKeyedPartitionStream<OUT> 
connectAndProcess(
             NonKeyedPartitionStream<T_OTHER> other,
             TwoInputNonBroadcastStreamProcessFunction<T, T_OTHER, OUT> 
processFunction) {
+        validateStates(
+                processFunction.usesStates(),
+                new HashSet<>(
+                        Arrays.asList(
+                                StateDeclaration.RedistributionMode.NONE,
+                                
StateDeclaration.RedistributionMode.IDENTICAL)));
+
         TypeInformation<OUT> outTypeInfo =
                 
StreamUtils.getOutputTypeForTwoInputNonBroadcastProcessFunction(
                         processFunction,
@@ -117,6 +145,10 @@ public class NonKeyedPartitionStreamImpl<T> extends 
AbstractDataStream<T>
     public <T_OTHER, OUT> ProcessConfigurableAndNonKeyedPartitionStream<OUT> 
connectAndProcess(
             BroadcastStream<T_OTHER> other,
             TwoInputBroadcastStreamProcessFunction<T, T_OTHER, OUT> 
processFunction) {
+        validateStates(
+                processFunction.usesStates(),
+                new 
HashSet<>(Collections.singletonList(StateDeclaration.RedistributionMode.NONE)));
+
         TypeInformation<OUT> outTypeInfo =
                 StreamUtils.getOutputTypeForTwoInputBroadcastProcessFunction(
                         processFunction,
diff --git 
a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/utils/StreamUtils.java
 
b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/utils/StreamUtils.java
index fb14f527b22..8a7ee569c9e 100644
--- 
a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/utils/StreamUtils.java
+++ 
b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/utils/StreamUtils.java
@@ -18,6 +18,8 @@
 
 package org.apache.flink.datastream.impl.utils;
 
+import org.apache.flink.api.common.state.IllegalRedistributionModeException;
+import org.apache.flink.api.common.state.StateDeclaration;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.connector.dsv2.Sink;
 import org.apache.flink.api.connector.dsv2.WrappedSink;
@@ -48,6 +50,8 @@ import 
org.apache.flink.streaming.api.transformations.DataStreamV2SinkTransforma
 import org.apache.flink.streaming.api.transformations.OneInputTransformation;
 import org.apache.flink.streaming.api.transformations.TwoInputTransformation;
 
+import java.util.Set;
+
 /**
  * This class encapsulates the common logic for all type of streams. It can be 
used to handle things
  * like extract type information, create a new transformation and so on for 
AbstractDataStream.
@@ -310,4 +314,18 @@ public final class StreamUtils {
             GlobalStreamImpl<T> stream) {
         return new ProcessConfigurableAndGlobalStreamImpl<>(stream);
     }
+
+    /** Wrap a {@link GlobalStreamImpl} with configure handle. */
+    public static void validateStates(
+            Set<StateDeclaration> inputStateDeclarations,
+            Set<StateDeclaration.RedistributionMode> invalidStateDeclarations) 
{
+        inputStateDeclarations.stream()
+                .map(StateDeclaration::getRedistributionMode)
+                .forEach(
+                        mode -> {
+                            if (invalidStateDeclarations.contains(mode)) {
+                                throw new 
IllegalRedistributionModeException(mode);
+                            }
+                        });
+    }
 }
diff --git 
a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/BroadcastStreamImplTest.java
 
b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/BroadcastStreamImplTest.java
index 0bd867276bc..6f49a780a2c 100644
--- 
a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/BroadcastStreamImplTest.java
+++ 
b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/BroadcastStreamImplTest.java
@@ -18,6 +18,10 @@
 
 package org.apache.flink.datastream.impl.stream;
 
+import org.apache.flink.api.common.state.IllegalRedistributionModeException;
+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.api.common.typeinfo.Types;
 import org.apache.flink.api.dag.Transformation;
 import org.apache.flink.datastream.api.stream.KeyedPartitionStream;
@@ -25,15 +29,25 @@ import 
org.apache.flink.datastream.impl.ExecutionEnvironmentImpl;
 import org.apache.flink.datastream.impl.TestingTransformation;
 import org.apache.flink.streaming.api.transformations.TwoInputTransformation;
 
+import org.assertj.core.api.Assertions;
 import org.junit.jupiter.api.Test;
 
+import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 
 import static 
org.apache.flink.datastream.impl.stream.StreamTestUtils.assertProcessType;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /** Tests for {@link BroadcastStreamImpl}. */
 class BroadcastStreamImplTest {
+
+    private final StateDeclaration modeNoneStateDeclaration =
+            StateDeclarations.listStateBuilder("list-state-none", 
TypeDescriptors.INT)
+                    
.redistributeWithMode(StateDeclaration.RedistributionMode.NONE)
+                    .build();
+
     @Test
     void testConnectNonKeyedStream() throws Exception {
         ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
@@ -49,6 +63,29 @@ class BroadcastStreamImplTest {
         assertProcessType(transformations.get(0), 
TwoInputTransformation.class, Types.LONG);
     }
 
+    @Test
+    void testStateErrorWithConnectNonKeyedStream() throws Exception {
+        ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
+        BroadcastStreamImpl<Integer> stream =
+                new BroadcastStreamImpl<>(env, new 
TestingTransformation<>("t1", Types.INT, 1));
+        NonKeyedPartitionStreamImpl<Long> nonKeyedStream =
+                new NonKeyedPartitionStreamImpl<>(
+                        env, new TestingTransformation<>("t2", Types.LONG, 2));
+        stream.connectAndProcess(
+                nonKeyedStream, new 
StreamTestUtils.NoOpTwoInputBroadcastStreamProcessFunction());
+
+        assertThatThrownBy(
+                        () ->
+                                stream.connectAndProcess(
+                                        nonKeyedStream,
+                                        new StreamTestUtils
+                                                
.NoOpTwoInputBroadcastStreamProcessFunction(
+                                                new HashSet<>(
+                                                        
Collections.singletonList(
+                                                                
modeNoneStateDeclaration)))))
+                .isInstanceOf(IllegalRedistributionModeException.class);
+    }
+
     @Test
     void testConnectKeyedStream() throws Exception {
         ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
@@ -65,6 +102,27 @@ class BroadcastStreamImplTest {
         assertProcessType(transformations.get(0), 
TwoInputTransformation.class, Types.LONG);
     }
 
+    @Test
+    void testStateErrorWithConnectKeyedStream() throws Exception {
+        ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
+        BroadcastStreamImpl<Integer> stream =
+                new BroadcastStreamImpl<>(env, new 
TestingTransformation<>("t1", Types.INT, 1));
+        NonKeyedPartitionStreamImpl<Long> nonKeyedStream =
+                new NonKeyedPartitionStreamImpl<>(
+                        env, new TestingTransformation<>("t2", Types.LONG, 2));
+
+        Assertions.assertThatCode(
+                        () ->
+                                stream.connectAndProcess(
+                                        nonKeyedStream.keyBy(x -> x),
+                                        new StreamTestUtils
+                                                
.NoOpTwoInputBroadcastStreamProcessFunction(
+                                                new HashSet<>(
+                                                        
Collections.singletonList(
+                                                                
modeNoneStateDeclaration)))))
+                .doesNotThrowAnyException();
+    }
+
     @Test
     void testConnectKeyedStreamWithOutputKeySelector() throws Exception {
         ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
diff --git 
a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/GlobalStreamImplTest.java
 
b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/GlobalStreamImplTest.java
index 0b57135a818..a2cf59088ed 100644
--- 
a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/GlobalStreamImplTest.java
+++ 
b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/GlobalStreamImplTest.java
@@ -19,6 +19,10 @@
 package org.apache.flink.datastream.impl.stream;
 
 import org.apache.flink.api.common.operators.SlotSharingGroup;
+import org.apache.flink.api.common.state.IllegalRedistributionModeException;
+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.api.common.typeinfo.Types;
 import org.apache.flink.api.connector.dsv2.DataStreamV2SinkUtils;
 import org.apache.flink.api.dag.Transformation;
@@ -35,6 +39,9 @@ import 
org.apache.flink.streaming.api.transformations.PartitionTransformation;
 
 import org.junit.jupiter.api.Test;
 
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -42,6 +49,17 @@ import static 
org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /** Tests for {@link GlobalStreamImpl}. */
 class GlobalStreamImplTest {
+
+    private final StateDeclaration modeIdenticalStateDeclaration =
+            StateDeclarations.listStateBuilder("list-state-identical", 
TypeDescriptors.INT)
+                    
.redistributeWithMode(StateDeclaration.RedistributionMode.IDENTICAL)
+                    .build();
+
+    private final StateDeclaration modeNoneStateDeclaration =
+            StateDeclarations.listStateBuilder("list-state-none", 
TypeDescriptors.INT)
+                    
.redistributeWithMode(StateDeclaration.RedistributionMode.NONE)
+                    .build();
+
     @Test
     void testParallelism() throws Exception {
         ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
@@ -59,6 +77,38 @@ class GlobalStreamImplTest {
         assertThat(transformations.get(2).getParallelism()).isOne();
     }
 
+    @Test
+    void testStateErrorWithOneInputStream() throws Exception {
+        ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
+        GlobalStreamImpl<Integer> stream =
+                new GlobalStreamImpl<>(env, new TestingTransformation<>("t1", 
Types.INT, 1));
+
+        assertThatThrownBy(
+                        () ->
+                                stream.process(
+                                        new NoOpOneInputStreamProcessFunction(
+                                                new HashSet<>(
+                                                        
Collections.singletonList(
+                                                                
modeIdenticalStateDeclaration)))))
+                .isInstanceOf(IllegalRedistributionModeException.class);
+    }
+
+    @Test
+    void testStateErrorWithTwoOutputStream() throws Exception {
+        ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
+        GlobalStreamImpl<Integer> stream =
+                new GlobalStreamImpl<>(env, new TestingTransformation<>("t1", 
Types.INT, 1));
+
+        assertThatThrownBy(
+                        () ->
+                                stream.process(
+                                        new NoOpTwoOutputStreamProcessFunction(
+                                                new HashSet<>(
+                                                        
Collections.singletonList(
+                                                                
modeIdenticalStateDeclaration)))))
+                .isInstanceOf(IllegalRedistributionModeException.class);
+    }
+
     @Test
     void testPartitioning() throws Exception {
         ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
@@ -80,6 +130,35 @@ class GlobalStreamImplTest {
         assertThat(transformations.get(2).getParallelism()).isOne();
     }
 
+    @Test
+    void testStateErrorWithTwoInputStream() throws Exception {
+        ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
+        GlobalStreamImpl<Integer> stream =
+                new GlobalStreamImpl<>(env, new TestingTransformation<>("t1", 
Types.INT, 1));
+
+        List<StateDeclaration> stateDeclarations =
+                Arrays.asList(modeNoneStateDeclaration, 
modeIdenticalStateDeclaration);
+
+        for (StateDeclaration stateDeclaration1 : stateDeclarations) {
+            for (StateDeclaration stateDeclaration2 : stateDeclarations) {
+                assertThatThrownBy(
+                                () ->
+                                        stream.connectAndProcess(
+                                                new GlobalStreamImpl<>(
+                                                        env,
+                                                        new 
TestingTransformation<>(
+                                                                "t2", 
Types.LONG, 1)),
+                                                new StreamTestUtils
+                                                        
.NoOpTwoInputNonBroadcastStreamProcessFunction(
+                                                        new HashSet<>(
+                                                                Arrays.asList(
+                                                                        
stateDeclaration1,
+                                                                        
stateDeclaration2)))))
+                        
.isInstanceOf(IllegalRedistributionModeException.class);
+            }
+        }
+    }
+
     @Test
     void testToSink() throws Exception {
         ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
diff --git 
a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/KeyedPartitionStreamImplTest.java
 
b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/KeyedPartitionStreamImplTest.java
index ac68d72ad8e..db1d3ffeefe 100644
--- 
a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/KeyedPartitionStreamImplTest.java
+++ 
b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/KeyedPartitionStreamImplTest.java
@@ -19,6 +19,10 @@
 package org.apache.flink.datastream.impl.stream;
 
 import org.apache.flink.api.common.operators.SlotSharingGroup;
+import org.apache.flink.api.common.state.IllegalRedistributionModeException;
+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.api.common.typeinfo.Types;
 import org.apache.flink.api.connector.dsv2.DataStreamV2SinkUtils;
 import org.apache.flink.api.dag.Transformation;
@@ -38,13 +42,23 @@ import 
org.apache.flink.streaming.api.transformations.TwoInputTransformation;
 
 import org.junit.jupiter.api.Test;
 
+import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 
 import static 
org.apache.flink.datastream.impl.stream.StreamTestUtils.assertProcessType;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /** Tests for {@link KeyedPartitionStreamImpl}. */
 public class KeyedPartitionStreamImplTest {
+
+    private final StateDeclaration modeIdenticalStateDeclaration =
+            StateDeclarations.listStateBuilder("list-state", 
TypeDescriptors.INT)
+                    
.redistributeWithMode(StateDeclaration.RedistributionMode.IDENTICAL)
+                    .build();
+
     @Test
     void testPartitioning() throws Exception {
         ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
@@ -65,6 +79,22 @@ public class KeyedPartitionStreamImplTest {
         assertThat(transformations.get(2).getParallelism()).isOne();
     }
 
+    @Test
+    void testStateErrorWithOneInputStream() throws Exception {
+        ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
+        KeyedPartitionStream<Integer, Integer> stream = createKeyedStream(env);
+
+        assertThatThrownBy(
+                        () ->
+                                stream.keyBy((data) -> 1)
+                                        .process(
+                                                new 
NoOpOneInputStreamProcessFunction(
+                                                        new HashSet<>(
+                                                                
Collections.singletonList(
+                                                                        
modeIdenticalStateDeclaration)))))
+                .isInstanceOf(IllegalRedistributionModeException.class);
+    }
+
     @Test
     void testProcess() throws Exception {
         ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
@@ -97,6 +127,21 @@ public class KeyedPartitionStreamImplTest {
         assertProcessType(transformations.get(1), 
OneInputTransformation.class, Types.INT);
     }
 
+    @Test
+    void testStateErrorWithProcessTwoOutput() throws Exception {
+        ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
+        KeyedPartitionStream<Integer, Integer> stream = createKeyedStream(env);
+
+        assertThatThrownBy(
+                        () ->
+                                stream.process(
+                                        new NoOpTwoOutputStreamProcessFunction(
+                                                new HashSet<>(
+                                                        
Collections.singletonList(
+                                                                
modeIdenticalStateDeclaration)))))
+                .isInstanceOf(IllegalRedistributionModeException.class);
+    }
+
     @Test
     void testConnectKeyedStream() throws Exception {
         ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
@@ -121,6 +166,26 @@ public class KeyedPartitionStreamImplTest {
         assertProcessType(transformations.get(1), 
TwoInputTransformation.class, Types.LONG);
     }
 
+    @Test
+    void testStateErrorWithConnectKeyedStream() throws Exception {
+        ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
+        KeyedPartitionStream<Integer, Integer> stream = createKeyedStream(env);
+
+        assertThatThrownBy(
+                        () ->
+                                stream.connectAndProcess(
+                                        createKeyedStream(
+                                                env,
+                                                new 
TestingTransformation<>("t2", Types.LONG, 1),
+                                                (KeySelector<Long, Integer>) 
Math::toIntExact),
+                                        new StreamTestUtils
+                                                
.NoOpTwoInputNonBroadcastStreamProcessFunction(
+                                                new HashSet<>(
+                                                        
Collections.singletonList(
+                                                                
modeIdenticalStateDeclaration)))))
+                .isInstanceOf(IllegalRedistributionModeException.class);
+    }
+
     @Test
     void testConnectBroadcastStream() throws Exception {
         ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
@@ -138,6 +203,26 @@ public class KeyedPartitionStreamImplTest {
         assertProcessType(transformations.get(1), 
TwoInputTransformation.class, Types.LONG);
     }
 
+    @Test
+    void testStateErrorWithConnectBroadcastStream() throws Exception {
+        ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
+        KeyedPartitionStream<Long, Long> stream =
+                createKeyedStream(env, new TestingTransformation<>("t1", 
Types.LONG, 1), x -> x);
+        BroadcastStreamImpl<Integer> s =
+                new BroadcastStreamImpl<>(env, new 
TestingTransformation<>("t2", Types.INT, 1));
+
+        assertThatCode(
+                        () ->
+                                stream.connectAndProcess(
+                                        s,
+                                        new StreamTestUtils
+                                                
.NoOpTwoInputBroadcastStreamProcessFunction(
+                                                new HashSet<>(
+                                                        
Collections.singletonList(
+                                                                
modeIdenticalStateDeclaration)))))
+                .doesNotThrowAnyException();
+    }
+
     @Test
     void testToSink() throws Exception {
         ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
diff --git 
a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/NonKeyedPartitionStreamImplTest.java
 
b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/NonKeyedPartitionStreamImplTest.java
index db7a7b5d2c4..4280f003b0d 100644
--- 
a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/NonKeyedPartitionStreamImplTest.java
+++ 
b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/NonKeyedPartitionStreamImplTest.java
@@ -19,6 +19,10 @@
 package org.apache.flink.datastream.impl.stream;
 
 import org.apache.flink.api.common.operators.SlotSharingGroup;
+import org.apache.flink.api.common.state.IllegalRedistributionModeException;
+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.api.common.typeinfo.Types;
 import org.apache.flink.api.connector.dsv2.DataStreamV2SinkUtils;
 import org.apache.flink.api.dag.Transformation;
@@ -34,13 +38,28 @@ import 
org.apache.flink.streaming.api.transformations.TwoInputTransformation;
 
 import org.junit.jupiter.api.Test;
 
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 
 import static 
org.apache.flink.datastream.impl.stream.StreamTestUtils.assertProcessType;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /** Tests for {@link NonKeyedPartitionStreamImpl}. */
 class NonKeyedPartitionStreamImplTest {
+
+    private final StateDeclaration modeIdenticalStateDeclaration =
+            StateDeclarations.listStateBuilder("list-state-identical", 
TypeDescriptors.INT)
+                    
.redistributeWithMode(StateDeclaration.RedistributionMode.IDENTICAL)
+                    .build();
+
+    private final StateDeclaration modeNoneStateDeclaration =
+            StateDeclarations.listStateBuilder("list-state-none", 
TypeDescriptors.INT)
+                    
.redistributeWithMode(StateDeclaration.RedistributionMode.NONE)
+                    .build();
+
     @Test
     void testPartitioning() throws Exception {
         ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
@@ -63,6 +82,46 @@ class NonKeyedPartitionStreamImplTest {
         assertThat(transformations.get(2).getParallelism()).isOne();
     }
 
+    @Test
+    void testStateErrorWithOneInputStream() throws Exception {
+        ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
+        NonKeyedPartitionStreamImpl<Integer> stream =
+                new NonKeyedPartitionStreamImpl<>(
+                        env, new TestingTransformation<>("t1", Types.INT, 1));
+
+        for (StateDeclaration stateDeclaration :
+                Arrays.asList(modeNoneStateDeclaration, 
modeIdenticalStateDeclaration)) {
+            assertThatThrownBy(
+                            () ->
+                                    stream.process(
+                                            new 
StreamTestUtils.NoOpOneInputStreamProcessFunction(
+                                                    new HashSet<>(
+                                                            
Collections.singletonList(
+                                                                    
stateDeclaration)))))
+                    .isInstanceOf(IllegalRedistributionModeException.class);
+        }
+    }
+
+    @Test
+    void testStateErrorWithTwoOutputStream() throws Exception {
+        ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
+        NonKeyedPartitionStreamImpl<Integer> stream =
+                new NonKeyedPartitionStreamImpl<>(
+                        env, new TestingTransformation<>("t1", Types.INT, 1));
+
+        for (StateDeclaration stateDeclaration :
+                Arrays.asList(modeNoneStateDeclaration, 
modeIdenticalStateDeclaration)) {
+            assertThatThrownBy(
+                            () ->
+                                    stream.process(
+                                            new 
StreamTestUtils.NoOpTwoOutputStreamProcessFunction(
+                                                    new HashSet<>(
+                                                            
Collections.singletonList(
+                                                                    
stateDeclaration)))))
+                    .isInstanceOf(IllegalRedistributionModeException.class);
+        }
+    }
+
     @Test
     void testProcess() throws Exception {
         ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
@@ -106,6 +165,36 @@ class NonKeyedPartitionStreamImplTest {
         assertProcessType(transformations.get(0), 
TwoInputTransformation.class, Types.LONG);
     }
 
+    @Test
+    void testStateErrorWithConnectNonKeyedStream() throws Exception {
+        ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
+        NonKeyedPartitionStreamImpl<Integer> stream =
+                new NonKeyedPartitionStreamImpl<>(
+                        env, new TestingTransformation<>("t1", Types.INT, 1));
+
+        for (StateDeclaration stateDeclaration1 :
+                Arrays.asList(modeNoneStateDeclaration, 
modeIdenticalStateDeclaration)) {
+            for (StateDeclaration stateDeclaration2 :
+                    Arrays.asList(modeNoneStateDeclaration, 
modeIdenticalStateDeclaration)) {
+
+                assertThatThrownBy(
+                                () ->
+                                        stream.connectAndProcess(
+                                                new 
NonKeyedPartitionStreamImpl<>(
+                                                        env,
+                                                        new 
TestingTransformation<>(
+                                                                "t2", 
Types.LONG, 1)),
+                                                new StreamTestUtils
+                                                        
.NoOpTwoInputNonBroadcastStreamProcessFunction(
+                                                        new HashSet<>(
+                                                                Arrays.asList(
+                                                                        
stateDeclaration1,
+                                                                        
stateDeclaration2)))))
+                        
.isInstanceOf(IllegalRedistributionModeException.class);
+            }
+        }
+    }
+
     @Test
     void testConnectBroadcastStream() throws Exception {
         ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
@@ -121,6 +210,27 @@ class NonKeyedPartitionStreamImplTest {
         assertProcessType(transformations.get(0), 
TwoInputTransformation.class, Types.LONG);
     }
 
+    @Test
+    void testStateErrorWithConnectBroadcastStream() throws Exception {
+        ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
+        NonKeyedPartitionStreamImpl<Long> stream =
+                new NonKeyedPartitionStreamImpl<>(
+                        env, new TestingTransformation<>("t1", Types.LONG, 1));
+
+        assertThatThrownBy(
+                        () ->
+                                stream.connectAndProcess(
+                                        new BroadcastStreamImpl<>(
+                                                env,
+                                                new 
TestingTransformation<>("t2", Types.INT, 1)),
+                                        new StreamTestUtils
+                                                
.NoOpTwoInputBroadcastStreamProcessFunction(
+                                                new HashSet<>(
+                                                        
Collections.singletonList(
+                                                                
modeNoneStateDeclaration)))))
+                .isInstanceOf(IllegalRedistributionModeException.class);
+    }
+
     @Test
     void testToSink() throws Exception {
         ExecutionEnvironmentImpl env = StreamTestUtils.getEnv();
diff --git 
a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/StreamTestUtils.java
 
b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/StreamTestUtils.java
index 85529da8bfa..75040edb158 100644
--- 
a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/StreamTestUtils.java
+++ 
b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/StreamTestUtils.java
@@ -18,6 +18,7 @@
 
 package org.apache.flink.datastream.impl.stream;
 
+import org.apache.flink.api.common.state.StateDeclaration;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.dag.Transformation;
 import org.apache.flink.datastream.api.ExecutionEnvironment;
@@ -31,6 +32,9 @@ import 
org.apache.flink.datastream.api.function.TwoOutputStreamProcessFunction;
 import org.apache.flink.datastream.impl.ExecutionEnvironmentImpl;
 import org.apache.flink.streaming.api.graph.StreamGraph;
 
+import java.util.HashSet;
+import java.util.Set;
+
 import static org.assertj.core.api.Assertions.assertThat;
 
 /** Test utils for steam. */
@@ -60,6 +64,21 @@ public final class StreamTestUtils {
     public static class NoOpOneInputStreamProcessFunction
             implements OneInputStreamProcessFunction<Integer, Long> {
 
+        private final Set<StateDeclaration> stateDeclarationSet;
+
+        public NoOpOneInputStreamProcessFunction(Set<StateDeclaration> 
stateDeclarationSet) {
+            this.stateDeclarationSet = stateDeclarationSet;
+        }
+
+        public NoOpOneInputStreamProcessFunction() {
+            this(new HashSet<>());
+        }
+
+        @Override
+        public Set<StateDeclaration> usesStates() {
+            return stateDeclarationSet;
+        }
+
         @Override
         public void processRecord(Integer record, Collector<Long> output, 
PartitionedContext ctx) {
             // do nothing.
@@ -70,6 +89,21 @@ public final class StreamTestUtils {
     public static class NoOpTwoOutputStreamProcessFunction
             implements TwoOutputStreamProcessFunction<Integer, Integer, Long> {
 
+        private final Set<StateDeclaration> stateDeclarationSet;
+
+        public NoOpTwoOutputStreamProcessFunction(Set<StateDeclaration> 
stateDeclarationSet) {
+            this.stateDeclarationSet = stateDeclarationSet;
+        }
+
+        public NoOpTwoOutputStreamProcessFunction() {
+            this(new HashSet<>());
+        }
+
+        @Override
+        public Set<StateDeclaration> usesStates() {
+            return stateDeclarationSet;
+        }
+
         @Override
         public void processRecord(
                 Integer record,
@@ -86,6 +120,22 @@ public final class StreamTestUtils {
     public static class NoOpTwoInputNonBroadcastStreamProcessFunction
             implements TwoInputNonBroadcastStreamProcessFunction<Integer, 
Long, Long> {
 
+        private final Set<StateDeclaration> stateDeclarationSet;
+
+        public NoOpTwoInputNonBroadcastStreamProcessFunction(
+                Set<StateDeclaration> stateDeclarationSet) {
+            this.stateDeclarationSet = stateDeclarationSet;
+        }
+
+        public NoOpTwoInputNonBroadcastStreamProcessFunction() {
+            this(new HashSet<>());
+        }
+
+        @Override
+        public Set<StateDeclaration> usesStates() {
+            return stateDeclarationSet;
+        }
+
         @Override
         public void processRecordFromFirstInput(
                 Integer record, Collector<Long> output, PartitionedContext 
ctx) {
@@ -104,6 +154,23 @@ public final class StreamTestUtils {
      */
     public static class NoOpTwoInputBroadcastStreamProcessFunction
             implements TwoInputBroadcastStreamProcessFunction<Long, Integer, 
Long> {
+
+        private final Set<StateDeclaration> stateDeclarationSet;
+
+        public NoOpTwoInputBroadcastStreamProcessFunction(
+                Set<StateDeclaration> stateDeclarationSet) {
+            this.stateDeclarationSet = stateDeclarationSet;
+        }
+
+        public NoOpTwoInputBroadcastStreamProcessFunction() {
+            this(new HashSet<>());
+        }
+
+        @Override
+        public Set<StateDeclaration> usesStates() {
+            return stateDeclarationSet;
+        }
+
         @Override
         public void processRecordFromNonBroadcastInput(
                 Long record, Collector<Long> output, PartitionedContext ctx) {

Reply via email to