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 8d7da02879e8276f1052ddc4b9193bfde12fef00 Author: Piotr Nowojski <[email protected]> AuthorDate: Wed Oct 28 19:21:54 2026 +0100 [FLINK-41256][runtime] Use moveToNextKey on the keyed input path of key-sorted inputs The keyed input path of AbstractStreamOperator(V2) now changes the key via moveToNextKey if the operator's keyed inputs are sorted by key, i.e. in BATCH execution with sorted inputs, which is the only case where the input guarantees that keys never go back. Otherwise, it keeps calling setCurrentKey. With sorted inputs, the input path no longer calls the operator's setCurrentKey: operators that override it to act on input key changes must override moveToNextKey too. Do so for the operators that do this: - The Python operators, which finish their bundle before the key changes in BATCH execution and track the key for their timer service. - StateBootstrapWrapperOperator, which forwards key changes to the wrapped operator. Generated-by: Claude Code (Claude Opus 5.5) --- .../operators/StateBootstrapWrapperOperator.java | 5 +++++ .../python/AbstractPythonFunctionOperator.java | 6 ++++++ .../ExternalPythonKeyedCoProcessOperator.java | 9 +++++++++ .../process/ExternalPythonKeyedProcessOperator.java | 9 +++++++++ .../AbstractPythonStreamAggregateOperator.java | 5 +++++ .../api/operators/AbstractStreamOperator.java | 21 ++++++++++++++++----- .../api/operators/AbstractStreamOperatorV2.java | 21 ++++++++++++++++----- .../over/BufferDataOverWindowOperatorTest.java | 2 ++ .../over/NonBufferOverWindowOperatorTest.java | 2 ++ 9 files changed, 70 insertions(+), 10 deletions(-) diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/output/operators/StateBootstrapWrapperOperator.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/output/operators/StateBootstrapWrapperOperator.java index 6e34026bd0d..e9f2c92e5dc 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/output/operators/StateBootstrapWrapperOperator.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/output/operators/StateBootstrapWrapperOperator.java @@ -171,6 +171,11 @@ public final class StateBootstrapWrapperOperator< operator.setCurrentKey(key); } + @Override + public void moveToNextKey(Object key) { + operator.moveToNextKey(key); + } + @Override public Object getCurrentKey() { return operator.getCurrentKey(); 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 cff3f2f6b80..809a622ef35 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 @@ -219,6 +219,12 @@ public abstract class AbstractPythonFunctionOperator<OUT> extends AbstractStream super.setCurrentKey(key); } + @Override + public void moveToNextKey(Object key) { + processElementsOfCurrentKeyIfNeeded(key); + super.moveToNextKey(key); + } + private void processElementsOfCurrentKeyIfNeeded(Object newKey) { // process all the elements belonging to the current key when encountering a new key // for batch operator diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/ExternalPythonKeyedCoProcessOperator.java b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/ExternalPythonKeyedCoProcessOperator.java index dc13aafae19..a69667f97b3 100644 --- a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/ExternalPythonKeyedCoProcessOperator.java +++ b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/ExternalPythonKeyedCoProcessOperator.java @@ -227,6 +227,15 @@ public class ExternalPythonKeyedCoProcessOperator<OUT> keyForTimerService = key; } + @Override + public void moveToNextKey(Object key) { + if (inBatchExecutionMode(getKeyedStateBackend())) { + super.moveToNextKey(key); + } + + keyForTimerService = key; + } + @Override public Object getCurrentKey() { return keyForTimerService; diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/ExternalPythonKeyedProcessOperator.java b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/ExternalPythonKeyedProcessOperator.java index e92301b6bf3..cf951574b3a 100644 --- a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/ExternalPythonKeyedProcessOperator.java +++ b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/ExternalPythonKeyedProcessOperator.java @@ -242,6 +242,15 @@ public class ExternalPythonKeyedProcessOperator<OUT> keyForTimerService = key; } + @Override + public void moveToNextKey(Object key) { + if (inBatchExecutionMode(getKeyedStateBackend())) { + super.moveToNextKey(key); + } + + keyForTimerService = key; + } + @Override public Object getCurrentKey() { return keyForTimerService; diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/aggregate/AbstractPythonStreamAggregateOperator.java b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/aggregate/AbstractPythonStreamAggregateOperator.java index 0298916b6d8..cfbb85d985c 100644 --- a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/aggregate/AbstractPythonStreamAggregateOperator.java +++ b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/aggregate/AbstractPythonStreamAggregateOperator.java @@ -208,6 +208,11 @@ public abstract class AbstractPythonStreamAggregateOperator keyForTimerService = key; } + @Override + public void moveToNextKey(Object key) { + keyForTimerService = key; + } + @Override public Object getCurrentKey() { return keyForTimerService; diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java index 1e2ed86ed8a..becfdf3103f 100644 --- a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java @@ -149,6 +149,9 @@ public abstract class AbstractStreamOperator<OUT> */ protected transient KeySelector<?, ?> stateKeySelector2; + /** Whether the keyed inputs are sorted by key, so that their keys only move forward. */ + private transient boolean keyedInputsSorted; + protected transient StreamOperatorStateHandler stateHandler; protected transient InternalTimeServiceManager<?> timeServiceManager; @@ -268,6 +271,9 @@ public abstract class AbstractStreamOperator<OUT> stateKeySelector1 = config.getStatePartitioner(0, getUserCodeClassloader()); stateKeySelector2 = config.getStatePartitioner(1, getUserCodeClassloader()); + keyedInputsSorted = + Arrays.stream(config.getInputs(getUserCodeClassloader())) + .anyMatch(StreamConfig::requiresSorting); lastRecordAttributes1 = RecordAttributes.EMPTY_RECORD_ATTRIBUTES; lastRecordAttributes2 = RecordAttributes.EMPTY_RECORD_ATTRIBUTES; @@ -607,7 +613,11 @@ public abstract class AbstractStreamOperator<OUT> throws Exception { if (selector != null) { Object key = selector.getKey(record.getValue()); - setCurrentKey(key); + if (keyedInputsSorted) { + moveToNextKey(key); + } else { + setCurrentKey(key); + } } } @@ -616,10 +626,11 @@ public abstract class AbstractStreamOperator<OUT> } /** - * Sets the key of the next record of a keyed input. Unlike {@link #setCurrentKey}, which an - * operator may call to switch to any key, it never goes back to a previous key of a key-sorted - * input. Override it, not {@link #setCurrentKey}, to act before the input moves on to the next - * key, e.g. to flush records buffered for the current key. + * Sets the key of the next record of a key-sorted input, i.e. in BATCH execution with sorted + * inputs, where the input path calls it instead of {@link #setCurrentKey}. Unlike {@link + * #setCurrentKey}, which an operator may call to switch to any key, it never goes back to a + * previous key. Override it to act before a key-sorted input moves on to the next key, e.g. to + * flush records buffered for the current key. */ @Override public void moveToNextKey(Object key) { diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperatorV2.java b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperatorV2.java index 1b9e2e2df39..4bc5c6d3062 100644 --- a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperatorV2.java +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperatorV2.java @@ -104,6 +104,9 @@ public abstract class AbstractStreamOperatorV2<OUT> private final CloseableRegistry cancelables; protected final IndexedCombinedWatermarkStatus combinedWatermark; + /** Whether the keyed inputs are sorted by key, so that their keys only move forward. */ + private final boolean keyedInputsSorted; + /** Metric group for the operator. */ protected final InternalOperatorMetricGroup metrics; @@ -144,6 +147,9 @@ public abstract class AbstractStreamOperatorV2<OUT> } executionConfig = parameters.getContainingTask().getExecutionConfig(); userCodeClassLoader = parameters.getContainingTask().getUserCodeClassLoader(); + keyedInputsSorted = + Arrays.stream(config.getInputs(userCodeClassLoader)) + .anyMatch(StreamConfig::requiresSorting); cancelables = parameters.getContainingTask().getCancelables(); combinedWatermark = IndexedCombinedWatermarkStatus.forInputsCount(numberOfInputs); @@ -494,7 +500,11 @@ public abstract class AbstractStreamOperatorV2<OUT> StreamRecord<T> record, KeySelector<T, ?> selector) throws Exception { if (selector != null) { Object key = selector.getKey(record.getValue()); - setCurrentKey(key); + if (keyedInputsSorted) { + moveToNextKey(key); + } else { + setCurrentKey(key); + } } } @@ -504,10 +514,11 @@ public abstract class AbstractStreamOperatorV2<OUT> } /** - * Sets the key of the next record of a keyed input. Unlike {@link #setCurrentKey}, which an - * operator may call to switch to any key, it never goes back to a previous key of a key-sorted - * input. Override it, not {@link #setCurrentKey}, to act before the input moves on to the next - * key, e.g. to flush records buffered for the current key. + * Sets the key of the next record of a key-sorted input, i.e. in BATCH execution with sorted + * inputs, where the input path calls it instead of {@link #setCurrentKey}. Unlike {@link + * #setCurrentKey}, which an operator may call to switch to any key, it never goes back to a + * previous key. Override it to act before a key-sorted input moves on to the next key, e.g. to + * flush records buffered for the current key. */ @Override public void moveToNextKey(Object key) { diff --git a/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/over/BufferDataOverWindowOperatorTest.java b/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/over/BufferDataOverWindowOperatorTest.java index 280755d6362..20b68a68083 100644 --- a/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/over/BufferDataOverWindowOperatorTest.java +++ b/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/over/BufferDataOverWindowOperatorTest.java @@ -204,6 +204,8 @@ class BufferDataOverWindowOperatorTest { any(ClassLoader.class))) .thenReturn(0.99); when(streamConfig.getOperatorID()).thenReturn(new OperatorID()); + when(streamConfig.getInputs(any(ClassLoader.class))) + .thenReturn(new StreamConfig.InputConfig[0]); operator = new BufferDataOverWindowOperator( new StreamOperatorParameters<>( diff --git a/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/over/NonBufferOverWindowOperatorTest.java b/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/over/NonBufferOverWindowOperatorTest.java index d8fdb59d6d0..0d6d458c0d0 100644 --- a/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/over/NonBufferOverWindowOperatorTest.java +++ b/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/over/NonBufferOverWindowOperatorTest.java @@ -147,6 +147,8 @@ class NonBufferOverWindowOperatorTest { any(ClassLoader.class))) .thenReturn(0.99); when(streamConfig.getOperatorID()).thenReturn(new OperatorID()); + when(streamConfig.getInputs(any(ClassLoader.class))) + .thenReturn(new StreamConfig.InputConfig[0]); operator = new NonBufferOverWindowOperator( new StreamOperatorParameters<>(
