This is an automated email from the ASF dual-hosted git repository. pnowojski pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/flink.git
commit f91e6224b83fd0250ce1bedfd9e2a0e7db7c637b Author: Piotr Nowojski <[email protected]> AuthorDate: Wed Oct 28 19:20:34 2026 +0100 [FLINK-41256][runtime] Fail on operators changing the current key out of order in BATCH execution BatchExecutionKeyedStateBackend processes keys one at a time in sorted order: every key change fires the outgoing key's timers and clears its state. Nothing prevented an operator from switching back to a key it had already left via setCurrentKey, for example to process records it buffered for several keys. That silently fired timers early and lost state: an operator that is correct in STREAMING produced wrong results in BATCH. Now that the keyed input path of key-sorted inputs changes keys via moveToNextKey, make the batch backend's setCurrentKey throw for any key other than the current one. Operators that buffer records across keys can override moveToNextKey to flush before the key changes. Async state context switches (setCurrentKeyAndKeyGroup) are not checked yet. The Python operators register the timers of the bundle they finish before the key changes right away, while the previous key is still current, instead of switching back to it later. Generated-by: Claude Code (Claude Opus 5.5) --- .../python/AbstractPythonFunctionOperator.java | 3 + .../state/BatchExecutionKeyedStateBackend.java | 19 ++- .../sorted/state/BatchExecutionKeyOrderTest.java | 108 +++++++++++++++++ .../benchmark/StateBackendBenchmarkUtilsTest.java | 2 +- .../datastream/DataStreamBatchExecutionITCase.java | 129 +++++++++++++++++++++ 5 files changed, 258 insertions(+), 3 deletions(-) diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractPythonFunctionOperator.java b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractPythonFunctionOperator.java index 809a622ef35..42007c865fa 100644 --- a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractPythonFunctionOperator.java +++ b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractPythonFunctionOperator.java @@ -234,6 +234,9 @@ public abstract class AbstractPythonFunctionOperator<OUT> extends AbstractStream while (!isBundleFinished()) { try { invokeFinishBundle(); + // register the timers of the bundle while the current key is still set, as + // they must not switch the key back once it has moved on + drainUnregisteredTimers(); fireAllRegisteredTimers(newKey); } catch (Exception e) { throw new WrappingRuntimeException(e); diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionKeyedStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionKeyedStateBackend.java index bfb127e0b9d..6661b85e215 100644 --- a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionKeyedStateBackend.java +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionKeyedStateBackend.java @@ -104,9 +104,23 @@ public class BatchExecutionKeyedStateBackend<K> implements CheckpointableKeyedSt this.executionConfig = executionConfig; } + /** + * Only allowed for the current key. Keys are processed one at a time, in sorted order, and may + * only advance via {@link #moveToNextKey}: switching to another key here would fire the + * outgoing key's timers early and drop its state, or revisit a key whose state is already gone. + */ @Override public void setCurrentKey(K newKey) { - moveToNextKey(newKey); + if (!Objects.equals(newKey, currentKey)) { + throw new IllegalStateException( + String.format( + "Cannot switch the current key from %s to %s: BATCH execution processes" + + " keys one at a time in sorted order, so an operator must not" + + " change the current key itself (e.g. to flush records it" + + " buffered for other keys). Keys only advance with the sorted" + + " input.", + currentKey, newKey)); + } } @Override @@ -132,7 +146,8 @@ public class BatchExecutionKeyedStateBackend<K> implements CheckpointableKeyedSt @Override public void setCurrentKeyAndKeyGroup(K newKey, int newKeyGroupIndex) { - setCurrentKey(newKey); + // Used by the async state context switches, which are not checked for key order yet. + moveToNextKey(newKey); } @Override diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionKeyOrderTest.java b/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionKeyOrderTest.java new file mode 100644 index 00000000000..137242ecdaf --- /dev/null +++ b/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionKeyOrderTest.java @@ -0,0 +1,108 @@ +/* + * 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.streaming.api.operators.sorted.state; + +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.common.typeutils.base.IntSerializer; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.streaming.api.graph.StreamConfig; +import org.apache.flink.streaming.api.operators.AbstractStreamOperator; +import org.apache.flink.streaming.api.operators.OneInputStreamOperator; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness; + +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * The single-key-at-a-time {@link BatchExecutionKeyedStateBackend} only moves forward through the + * sorted keys. Switching back to an earlier key would silently fire its timers early and lose its + * state, so it must fail loudly. + */ +class BatchExecutionKeyOrderTest { + + @Test + void testInputPathMovesThroughKeys() throws Exception { + try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> harness = + createBatchHarness(new KeyRevisitingOperator(false))) { + harness.open(); + for (int key = 1; key <= 3; key++) { + harness.processElement(new StreamRecord<>(key)); + harness.processElement(new StreamRecord<>(key)); + } + assertThat(harness.extractOutputValues()).containsExactly(1, 1, 2, 2, 3, 3); + } + } + + @Test + void testOperatorSwitchingBackToPreviousKeyFails() throws Exception { + try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> harness = + createBatchHarness(new KeyRevisitingOperator(true))) { + harness.open(); + harness.processElement(new StreamRecord<>(1)); + assertThatThrownBy(() -> harness.processElement(new StreamRecord<>(2))) + .hasStackTraceContaining("Cannot switch the current key from 2 to 1"); + } + } + + private static KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> + createBatchHarness(OneInputStreamOperator<Integer, Integer> operator) throws Exception { + KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> harness = + new KeyedOneInputStreamOperatorTestHarness<>( + operator, (KeySelector<Integer, Integer>) value -> value, Types.INT); + harness.setStateBackend(new BatchExecutionStateBackend()); + // As in BATCH execution, where keyed inputs are sorted by key. + StreamConfig config = harness.getStreamConfig(); + config.setInputs( + new StreamConfig.NetworkInputConfig( + IntSerializer.INSTANCE, 0, StreamConfig.InputRequirement.SORTED)); + config.serializeAllConfigs(); + return harness; + } + + /** + * Emits each record. With {@code revisit}, it first switches back to the previous record's key + * (as an operator flushing a cross-key buffer would) and then to the current one. + */ + private static class KeyRevisitingOperator extends AbstractStreamOperator<Integer> + implements OneInputStreamOperator<Integer, Integer> { + + private final boolean revisit; + private Integer previousKey; + + KeyRevisitingOperator(boolean revisit) { + this.revisit = revisit; + } + + @Override + public void processElement(StreamRecord<Integer> element) { + Integer key = (Integer) getCurrentKey(); + // Same-key calls are always fine. + setCurrentKey(key); + if (revisit && previousKey != null && !previousKey.equals(key)) { + setCurrentKey(previousKey); + setCurrentKey(key); + } + previousKey = key; + output.collect(element); + } + } +} diff --git a/flink-test-utils-parent/flink-test-utils/src/test/java/org/apache/flink/state/benchmark/StateBackendBenchmarkUtilsTest.java b/flink-test-utils-parent/flink-test-utils/src/test/java/org/apache/flink/state/benchmark/StateBackendBenchmarkUtilsTest.java index 0bdc8e48f1f..93f7a1a0ed2 100644 --- a/flink-test-utils-parent/flink-test-utils/src/test/java/org/apache/flink/state/benchmark/StateBackendBenchmarkUtilsTest.java +++ b/flink-test-utils-parent/flink-test-utils/src/test/java/org/apache/flink/state/benchmark/StateBackendBenchmarkUtilsTest.java @@ -123,7 +123,7 @@ public class StateBackendBenchmarkUtilsTest { KeyedStateBackend<Long> backend = createKeyedStateBackend(backendType); ListState<Long> listState = getListState(backend, listStateDescriptor); for (long i = 0; i < 10; i++) { - backend.setCurrentKey(i); + backend.moveToNextKey(i); listState.add(i); } if (backend instanceof RocksDBKeyedStateBackend) { diff --git a/flink-tests/src/test/java/org/apache/flink/api/datastream/DataStreamBatchExecutionITCase.java b/flink-tests/src/test/java/org/apache/flink/api/datastream/DataStreamBatchExecutionITCase.java index 67c08688b31..94f6ba7753b 100644 --- a/flink-tests/src/test/java/org/apache/flink/api/datastream/DataStreamBatchExecutionITCase.java +++ b/flink-tests/src/test/java/org/apache/flink/api/datastream/DataStreamBatchExecutionITCase.java @@ -26,8 +26,10 @@ import org.apache.flink.api.common.state.ListState; import org.apache.flink.api.common.state.ListStateDescriptor; import org.apache.flink.api.common.state.MapStateDescriptor; import org.apache.flink.api.common.state.ReadOnlyBroadcastState; +import org.apache.flink.api.common.state.ValueState; import org.apache.flink.api.common.state.ValueStateDescriptor; import org.apache.flink.api.common.typeinfo.BasicTypeInfo; +import org.apache.flink.api.common.typeinfo.Types; import org.apache.flink.api.common.typeutils.base.StringSerializer; import org.apache.flink.api.java.tuple.Tuple2; import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration; @@ -44,8 +46,10 @@ import org.apache.flink.streaming.api.operators.AbstractInput; import org.apache.flink.streaming.api.operators.AbstractStreamOperator; import org.apache.flink.streaming.api.operators.AbstractStreamOperatorFactory; import org.apache.flink.streaming.api.operators.AbstractStreamOperatorV2; +import org.apache.flink.streaming.api.operators.BoundedOneInput; import org.apache.flink.streaming.api.operators.Input; import org.apache.flink.streaming.api.operators.MultipleInputStreamOperator; +import org.apache.flink.streaming.api.operators.OneInputStreamOperator; import org.apache.flink.streaming.api.operators.StreamOperator; import org.apache.flink.streaming.api.operators.StreamOperatorParameters; import org.apache.flink.streaming.api.operators.TwoInputStreamOperator; @@ -58,15 +62,21 @@ import org.apache.flink.test.util.MiniClusterWithClientResource; import org.apache.flink.util.CloseableIterator; import org.apache.flink.util.CollectionUtil; import org.apache.flink.util.Collector; +import org.apache.flink.util.FlinkRuntimeException; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; +import java.util.ArrayList; import java.util.Arrays; +import java.util.HashMap; import java.util.List; +import java.util.Map; +import java.util.Objects; import static org.apache.flink.util.CollectionUtil.iteratorToList; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** * Integration test for {@link RuntimeExecutionMode#BATCH} execution on the DataStream API. @@ -530,6 +540,60 @@ class DataStreamBatchExecutionITCase { } } + /** + * An operator that buffers records of several keys and later processes them key by key, via + * {@code setCurrentKey}, works in STREAMING. In BATCH execution the keyed state backend keeps + * only the current key's state, so going back to a key that the sorted input already left would + * silently lose that key's state: the job must fail instead. + */ + @Test + void batchOperatorGoingBackToAPreviousKeyFails() throws Exception { + assertThat(runBufferingSum(RuntimeExecutionMode.STREAMING, false)) + .containsExactlyInAnyOrderEntriesOf(EXPECTED_BUFFERING_SUMS); + assertThatThrownBy(() -> runBufferingSum(RuntimeExecutionMode.BATCH, false)) + .hasStackTraceContaining("Cannot switch the current key from b to a"); + } + + /** The same operator works in BATCH if it flushes its buffer before the next key. */ + @Test + void batchOperatorFlushingBeforeTheNextKey() throws Exception { + assertThat(runBufferingSum(RuntimeExecutionMode.STREAMING, true)) + .containsExactlyInAnyOrderEntriesOf(EXPECTED_BUFFERING_SUMS); + assertThat(runBufferingSum(RuntimeExecutionMode.BATCH, true)) + .containsExactlyInAnyOrderEntriesOf(EXPECTED_BUFFERING_SUMS); + } + + private static final Map<String, Integer> EXPECTED_BUFFERING_SUMS = Map.of("a", 3, "b", 1); + + /** Returns the last sum emitted for each key. */ + private static Map<String, Integer> runBufferingSum( + RuntimeExecutionMode mode, boolean flushBeforeNextKey) throws Exception { + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setRuntimeMode(mode); + env.setParallelism(1); + RestartStrategyUtils.configureNoRestartStrategy(env); + + // With a buffer of two records, a1 and a2 are summed into a's state, a3 stays buffered + // when the input moves on to b, and the flush of [a3, b1] goes back to a. + DataStream<Tuple2<String, Integer>> sums = + env.fromData( + Tuple2.of("a", 1), + Tuple2.of("a", 1), + Tuple2.of("a", 1), + Tuple2.of("b", 1)) + .keyBy(in -> in.f0) + .transform( + "buffering-sum", + Types.TUPLE(Types.STRING, Types.INT), + new BufferingSumOperator(2, flushBeforeNextKey)); + + Map<String, Integer> lastSums = new HashMap<>(); + try (CloseableIterator<Tuple2<String, Integer>> result = sums.executeAndCollect()) { + result.forEachRemaining(sum -> lastSums.put(sum.f0, sum.f1)); + } + return lastSums; + } + private StreamExecutionEnvironment getExecutionEnvironment() { final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); env.setRuntimeMode(RuntimeExecutionMode.BATCH); @@ -739,4 +803,69 @@ class DataStreamBatchExecutionITCase { }); } } + + /** + * Buffers up to {@code bufferSize} records of any keys and then sums them per key into keyed + * state, switching to each record's key with {@code setCurrentKey}, like operators that batch + * records across keys to save state accesses. Emits the updated sum of each processed record. + * Optionally flushes the buffer before the input moves on to the next key, which BATCH + * execution requires. + */ + private static class BufferingSumOperator + extends AbstractStreamOperator<Tuple2<String, Integer>> + implements OneInputStreamOperator<Tuple2<String, Integer>, Tuple2<String, Integer>>, + BoundedOneInput { + + private final int bufferSize; + private final boolean flushBeforeNextKey; + private final List<Tuple2<String, Integer>> buffer = new ArrayList<>(); + private transient ValueState<Integer> sum; + + BufferingSumOperator(int bufferSize, boolean flushBeforeNextKey) { + this.bufferSize = bufferSize; + this.flushBeforeNextKey = flushBeforeNextKey; + } + + @Override + public void moveToNextKey(Object key) { + if (flushBeforeNextKey && !Objects.equals(key, getCurrentKey())) { + try { + flush(); + } catch (Exception e) { + throw new FlinkRuntimeException(e); + } + } + super.moveToNextKey(key); + } + + @Override + public void open() throws Exception { + super.open(); + sum = getRuntimeContext().getState(new ValueStateDescriptor<>("sum", Types.INT)); + } + + @Override + public void processElement(StreamRecord<Tuple2<String, Integer>> element) throws Exception { + buffer.add(element.getValue()); + if (buffer.size() >= bufferSize) { + flush(); + } + } + + @Override + public void endInput() throws Exception { + flush(); + } + + private void flush() throws Exception { + for (Tuple2<String, Integer> record : buffer) { + setCurrentKey(record.f0); + Integer current = sum.value(); + int updated = (current == null ? 0 : current) + record.f1; + sum.update(updated); + output.collect(new StreamRecord<>(Tuple2.of(record.f0, updated))); + } + buffer.clear(); + } + } }
