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 f9abda3bbefae59ccd046382999ae12b20e981c0 Author: Piotr Nowojski <[email protected]> AuthorDate: Thu Oct 29 10:42:54 2026 +0100 [FLINK-41256][runtime] Don't set the current key in KeyedSortPartitionOperator KeyedSortPartitionOperator sorts records by key itself, so its input is not sorted by key, but it still runs on BatchExecutionKeyedStateBackend in BATCH execution. Setting the current key for every record therefore failed on the first key change. The operator uses neither keyed state nor timers and only needed the key to sort by, so it now takes the key from the key selector and no longer sets the current key. Generated-by: Claude Code (Claude Opus 5.5) --- .../sortpartition/KeyedSortPartitionOperator.java | 13 ++++++++-- .../KeyedSortPartitionOperatorTest.java | 29 ++++++++-------------- .../KeyedPartitionWindowedStreamITCase.java | 22 ++++++++++++++++ 3 files changed, 44 insertions(+), 20 deletions(-) diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sortpartition/KeyedSortPartitionOperator.java b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sortpartition/KeyedSortPartitionOperator.java index e986efffee3..0a9e502d5cd 100644 --- a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sortpartition/KeyedSortPartitionOperator.java +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sortpartition/KeyedSortPartitionOperator.java @@ -207,10 +207,19 @@ public class KeyedSortPartitionOperator<INPUT, KEY> extends AbstractStreamOperat } } + /** + * Does not set the current key. This operator sorts records by key itself, so its input is not + * sorted by key, and it uses neither keyed state nor timers. Switching the key of the keyed + * state backend for every record would fail in BATCH execution, which expects keys to change + * only in sorted order. + */ + @Override + public void setKeyContextElement1(StreamRecord record) {} + @Override public void processElement(StreamRecord<INPUT> element) throws Exception { - KEY currentKey = (KEY) getCurrentKey(); - recordKeySerializer.serialize(currentKey, dataOutputSerializer); + KEY key = ((KeySelector<INPUT, KEY>) stateKeySelector1).getKey(element.getValue()); + recordKeySerializer.serialize(key, dataOutputSerializer); byte[] serializedKey = dataOutputSerializer.getCopyOfBuffer(); dataOutputSerializer.clear(); if (sortFieldSelector != null) { diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/sortpartition/KeyedSortPartitionOperatorTest.java b/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/sortpartition/KeyedSortPartitionOperatorTest.java index 470381ecb46..ce70ffef386 100644 --- a/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/sortpartition/KeyedSortPartitionOperatorTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/sortpartition/KeyedSortPartitionOperatorTest.java @@ -19,13 +19,12 @@ package org.apache.flink.streaming.api.operators.sortpartition; import org.apache.flink.api.common.operators.Order; -import org.apache.flink.api.common.serialization.SerializerConfig; import org.apache.flink.api.common.typeinfo.BasicTypeInfo; import org.apache.flink.api.common.typeinfo.TypeInformation; import org.apache.flink.api.common.typeinfo.Types; import org.apache.flink.api.java.tuple.Tuple2; -import org.apache.flink.streaming.api.graph.StreamConfig; import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness; import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness; import org.apache.flink.streaming.util.TestHarnessUtil; @@ -46,10 +45,9 @@ class KeyedSortPartitionOperatorTest { KeyedSortPartitionOperator<Tuple2<Integer, String>, String> operator1 = createSortPartitionOperatorWithPositionField(); OneInputStreamOperatorTestHarness<Tuple2<Integer, String>, Tuple2<Integer, String>> - testHarness1 = new OneInputStreamOperatorTestHarness<>(operator1); - StreamConfig streamConfig1 = testHarness1.getStreamConfig(); - streamConfig1.setStateKeySerializer(Types.STRING.createSerializer((SerializerConfig) null)); - streamConfig1.serializeAllConfigs(); + testHarness1 = + new KeyedOneInputStreamOperatorTestHarness<>( + operator1, value -> value.f1, Types.STRING); Queue<Object> expectedOutput1 = new LinkedList<>(); testHarness1.open(); testHarness1.processElement(new StreamRecord<>(Tuple2.of(3, "3"))); @@ -66,10 +64,8 @@ class KeyedSortPartitionOperatorTest { KeyedSortPartitionOperator<TestPojo, String> operator2 = createSortPartitionOperatorWithStringField(); OneInputStreamOperatorTestHarness<TestPojo, TestPojo> testHarness2 = - new OneInputStreamOperatorTestHarness<>(operator2); - StreamConfig streamConfig2 = testHarness2.getStreamConfig(); - streamConfig2.setStateKeySerializer(Types.STRING.createSerializer((SerializerConfig) null)); - streamConfig2.serializeAllConfigs(); + new KeyedOneInputStreamOperatorTestHarness<>( + operator2, TestPojo::getKey, Types.STRING); Queue<Object> expectedOutput2 = new LinkedList<>(); testHarness2.open(); testHarness2.processElement(new StreamRecord<>(new TestPojo("3", 3))); @@ -86,10 +82,8 @@ class KeyedSortPartitionOperatorTest { KeyedSortPartitionOperator<TestPojo, String> operator3 = createSortPartitionOperatorWithKeySelector(); OneInputStreamOperatorTestHarness<TestPojo, TestPojo> testHarness3 = - new OneInputStreamOperatorTestHarness<>(operator3); - StreamConfig streamConfig3 = testHarness3.getStreamConfig(); - streamConfig3.setStateKeySerializer(Types.STRING.createSerializer((SerializerConfig) null)); - streamConfig3.serializeAllConfigs(); + new KeyedOneInputStreamOperatorTestHarness<>( + operator3, TestPojo::getKey, Types.STRING); Queue<Object> expectedOutput3 = new LinkedList<>(); testHarness3.open(); testHarness3.processElement(new StreamRecord<>(new TestPojo("3", 3))); @@ -109,10 +103,9 @@ class KeyedSortPartitionOperatorTest { KeyedSortPartitionOperator<Tuple2<Integer, String>, String> sortPartitionOperator = createSortPartitionOperatorWithPositionField(); OneInputStreamOperatorTestHarness<Tuple2<Integer, String>, Tuple2<Integer, String>> - testHarness = new OneInputStreamOperatorTestHarness<>(sortPartitionOperator); - StreamConfig streamConfig = testHarness.getStreamConfig(); - streamConfig.setStateKeySerializer(Types.STRING.createSerializer((SerializerConfig) null)); - streamConfig.serializeAllConfigs(); + testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + sortPartitionOperator, value -> value.f1, Types.STRING); testHarness.open(); testHarness.processElement(new StreamRecord<>(Tuple2.of(1, "1"))); testHarness.endInput(); diff --git a/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/KeyedPartitionWindowedStreamITCase.java b/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/KeyedPartitionWindowedStreamITCase.java index 78e6dbd0a5f..688a524778a 100644 --- a/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/KeyedPartitionWindowedStreamITCase.java +++ b/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/KeyedPartitionWindowedStreamITCase.java @@ -18,6 +18,7 @@ package org.apache.flink.test.streaming.runtime; +import org.apache.flink.api.common.RuntimeExecutionMode; import org.apache.flink.api.common.functions.AggregateFunction; import org.apache.flink.api.common.functions.MapFunction; import org.apache.flink.api.common.functions.MapPartitionFunction; @@ -305,6 +306,27 @@ class KeyedPartitionWindowedStreamITCase { "77 66 55 8 "); } + @Test + void testSortPartitionInBatchExecutionMode() throws Exception { + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setRuntimeMode(RuntimeExecutionMode.BATCH); + env.setParallelism(1); + CloseableIterator<String> resultIterator = + env.fromData( + Tuple2.of("key1", 7), + Tuple2.of("key2", 100), + Tuple2.of("key1", 0), + Tuple2.of("key2", 1), + Tuple2.of("key1", 3)) + .keyBy(value -> value.f0) + .fullWindowPartition() + .sortPartition(1, Order.ASCENDING) + .map(value -> value.f0 + " " + value.f1) + .executeAndCollect(); + assertThat(Lists.newArrayList(resultIterator)) + .containsExactly("key1 0", "key1 3", "key1 7", "key2 1", "key2 100"); + } + private CloseableIterator<String> sortPartitionOfTupleElementsInOrder(Order order) throws Exception { StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
