make BytesMap#getEntryIterator copy-aware

Author: shuo.cs

flink-table/flink-table-runtime-blink/src/main/java/org/apache/flink/table/runtime/operators/aggregate/window/SlicingWindowAggOperatorBuilder.java

@@ -138,13 +138,9 @@ public SlicingWindowAggOperatorBuilder countStarIndex(int indexOfCountStart) {
         if (isGlobalAgg) {
             combinerFactory =
                     new GlobalAggCombiner.Factory(
-                            localGeneratedAggregateFunction,
-                            globalGeneratedAggregateFunction,
-                            keySerializer);
+                            localGeneratedAggregateFunction, globalGeneratedAggregateFunction);
         } else {
-            combinerFactory =
-                    new AggCombiner.Factory(
-                            generatedAggregateFunction, keySerializer, inputSerializer);
+            combinerFactory = new AggCombiner.Factory(generatedAggregateFunction);
         }
         final WindowBuffer.Factory bufferFactory =
                 new RecordsWindowBuffer.Factory(keySerializer, inputSerializer, combinerFactory);

flink-table/flink-table-runtime-blink/src/main/java/org/apache/flink/table/runtime/operators/aggregate/window/buffers/RecordsWindowBuffer.java

@@ -38,6 +38,7 @@
 import java.time.ZoneId;
 import java.util.Iterator;
 
+import static org.apache.flink.table.runtime.util.StateConfigUtil.isStateImmutableInStateBackend;
 import static org.apache.flink.table.runtime.util.TimeWindowUtil.isWindowFired;
 
 /**
@@ -51,6 +52,9 @@ public final class RecordsWindowBuffer implements WindowBuffer {
     private final WindowKey reuseWindowKey;
     private final AbstractRowDataSerializer<RowData> recordSerializer;
     private final ZoneId shiftTimeZone;
+    // copy key and input record if necessary(e.g., heap state backend),
+    // because key and record are reused.
+    private final boolean requiresCopy;
 
     private long minSliceEnd = Long.MAX_VALUE;
 
@@ -61,13 +65,15 @@ public RecordsWindowBuffer(
             RecordsCombiner combineFunction,
             PagedTypeSerializer<RowData> keySer,
             AbstractRowDataSerializer<RowData> inputSer,
+            boolean requiresCopy,
             ZoneId shiftTimeZone) {
         this.combineFunction = combineFunction;
         this.recordsBuffer =
                 new WindowBytesMultiMap(
                         operatorOwner, memoryManager, memorySize, keySer, inputSer.getArity());
         this.recordSerializer = inputSer;
         this.reuseWindowKey = new WindowKeySerializer(keySer).createInstance();
+        this.requiresCopy = requiresCopy;
         this.shiftTimeZone = shiftTimeZone;
     }
 
@@ -102,7 +108,7 @@ public void advanceProgress(long progress) throws Exception {
     public void flush() throws Exception {
         if (recordsBuffer.getNumKeys() > 0) {
             KeyValueIterator<WindowKey, Iterator<RowData>> entryIterator =
-                    recordsBuffer.getEntryIterator();
+                    recordsBuffer.getEntryIterator(requiresCopy);
             while (entryIterator.advanceNext()) {
                 combineFunction.combine(entryIterator.getKey(), entryIterator.getValue());
             }
@@ -155,13 +161,15 @@ public WindowBuffer create(
             RecordsCombiner combiner =
                     factory.createRecordsCombiner(
                             runtimeContext, timerService, stateBackend, windowState, isEventTime);
+            boolean requiresCopy = !isStateImmutableInStateBackend(stateBackend);
             return new RecordsWindowBuffer(
                     operatorOwner,
                     memoryManager,
                     memorySize,
                     combiner,
                     keySer,
                     inputSer,
+                    requiresCopy,
                     shiftTimeZone);
         }
     }
@@ -202,6 +210,7 @@ public WindowBuffer create(
                     combiner,
                     keySer,
                     inputSer,
+                    false,
                     shiftTimeZone);
         }
     }

flink-table/flink-table-runtime-blink/src/main/java/org/apache/flink/table/runtime/operators/aggregate/window/combines/AggCombiner.java

@@ -19,7 +19,6 @@
 package org.apache.flink.table.runtime.operators.aggregate.window.combines;
 
 import org.apache.flink.api.common.functions.RuntimeContext;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
 import org.apache.flink.api.common.typeutils.base.LongSerializer;
 import org.apache.flink.runtime.state.KeyedStateBackend;
 import org.apache.flink.table.data.RowData;
@@ -37,7 +36,6 @@
 import java.util.Iterator;
 
 import static org.apache.flink.table.data.util.RowDataUtil.isAccumulateMsg;
-import static org.apache.flink.table.runtime.util.StateConfigUtil.isStateImmutableInStateBackend;
 import static org.apache.flink.table.runtime.util.TimeWindowUtil.isWindowFired;
 
 /**
@@ -58,15 +56,6 @@ public class AggCombiner implements RecordsCombiner {
     /** Function used to handle all aggregates. */
     private final NamespaceAggsHandleFunction<Long> aggregator;
 
-    /** Whether to copy key and input record, because key and record are reused. */
-    private final boolean requiresCopy;
-
-    /** Serializer to copy key if required. */
-    private final TypeSerializer<RowData> keySerializer;
-
-    /** Serializer to copy record if required. */
-    private final TypeSerializer<RowData> recordSerializer;
-
     /** Whether the operator works in event-time mode, used to indicate registering which timer. */
     private final boolean isEventTime;
 
@@ -75,31 +64,18 @@ public AggCombiner(
             StateKeyContext keyContext,
             WindowValueState<Long> accState,
             NamespaceAggsHandleFunction<Long> aggregator,
-            boolean requiresCopy,
-            TypeSerializer<RowData> keySerializer,
-            TypeSerializer<RowData> valueSerializer,
             boolean isEventTime) {
         this.timerService = timerService;
         this.keyContext = keyContext;
         this.accState = accState;
         this.aggregator = aggregator;
-        this.requiresCopy = requiresCopy;
-        this.keySerializer = keySerializer;
-        this.recordSerializer = valueSerializer;
         this.isEventTime = isEventTime;
     }
 
     @Override
     public void combine(WindowKey windowKey, Iterator<RowData> records) throws Exception {
         // step 0: set current key for states and timers
-        final RowData key;
-        if (requiresCopy) {
-            // the incoming key is reused, we should copy it if state backend doesn't copy it
-            key = keySerializer.copy(windowKey.getKey());
-        } else {
-            key = windowKey.getKey();
-        }
-        keyContext.setCurrentKey(key);
+        keyContext.setCurrentKey(windowKey.getKey());
 
         // step 1: get the accumulator for the current key and window
         Long window = windowKey.getWindow();
@@ -114,10 +90,6 @@ public void combine(WindowKey windowKey, Iterator<RowData> records) throws Excep
         // step 3: do accumulate
         while (records.hasNext()) {
             RowData record = records.next();
-            if (requiresCopy) {
-                // the incoming record is reused, we should copy it if state backend doesn't copy it
-                record = recordSerializer.copy(record);
-            }
             if (isAccumulateMsg(record)) {
                 aggregator.accumulate(record);
             } else {
@@ -156,16 +128,9 @@ public static final class Factory implements RecordsCombiner.Factory {
         private static final long serialVersionUID = 1L;
 
         private final GeneratedNamespaceAggsHandleFunction<Long> genAggsHandler;
-        private final TypeSerializer<RowData> keySerializer;
-        private final TypeSerializer<RowData> valueSerializer;
 
-        public Factory(
-                GeneratedNamespaceAggsHandleFunction<Long> genAggsHandler,
-                TypeSerializer<RowData> keySerializer,
-                TypeSerializer<RowData> valueSerializer) {
+        public Factory(GeneratedNamespaceAggsHandleFunction<Long> genAggsHandler) {
             this.genAggsHandler = genAggsHandler;
-            this.keySerializer = keySerializer;
-            this.valueSerializer = valueSerializer;
         }
 
         @Override
@@ -181,16 +146,12 @@ public RecordsCombiner createRecordsCombiner(
             aggregator.open(
                     new PerWindowStateDataViewStore(
                             stateBackend, LongSerializer.INSTANCE, runtimeContext));
-            boolean requiresCopy = !isStateImmutableInStateBackend(stateBackend);
             WindowValueState<Long> windowValueState = (WindowValueState<Long>) windowState;
             return new AggCombiner(
                     timerService,
                     stateBackend::setCurrentKey,
                     windowValueState,
                     aggregator,
-                    requiresCopy,
-                    keySerializer,
-                    valueSerializer,
                     isEventTime);
         }
     }

flink-table/flink-table-runtime-blink/src/main/java/org/apache/flink/table/runtime/operators/aggregate/window/combines/GlobalAggCombiner.java

@@ -19,7 +19,6 @@
 package org.apache.flink.table.runtime.operators.aggregate.window.combines;
 
 import org.apache.flink.api.common.functions.RuntimeContext;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
 import org.apache.flink.api.common.typeutils.base.LongSerializer;
 import org.apache.flink.runtime.state.KeyedStateBackend;
 import org.apache.flink.table.data.RowData;
@@ -36,7 +35,6 @@
 import java.time.ZoneId;
 import java.util.Iterator;
 
-import static org.apache.flink.table.runtime.util.StateConfigUtil.isStateImmutableInStateBackend;
 import static org.apache.flink.table.runtime.util.TimeWindowUtil.isWindowFired;
 
 /**
@@ -62,27 +60,17 @@ public class GlobalAggCombiner implements RecordsCombiner {
     /** Global aggregate function to handle global accumulator rows. */
     private final NamespaceAggsHandleFunction<Long> globalAggregator;
 
-    /** Whether to copy key and input record, because key and record are reused. */
-    private final boolean requiresCopy;
-
-    /** Serializer to copy key if required. */
-    private final TypeSerializer<RowData> keySerializer;
-
     public GlobalAggCombiner(
             WindowTimerService<Long> timerService,
             StateKeyContext keyContext,
             WindowValueState<Long> accState,
             NamespaceAggsHandleFunction<Long> localAggregator,
-            NamespaceAggsHandleFunction<Long> globalAggregator,
-            boolean requiresCopy,
-            TypeSerializer<RowData> keySerializer) {
+            NamespaceAggsHandleFunction<Long> globalAggregator) {
         this.timerService = timerService;
         this.keyContext = keyContext;
         this.accState = accState;
         this.localAggregator = localAggregator;
         this.globalAggregator = globalAggregator;
-        this.requiresCopy = requiresCopy;
-        this.keySerializer = keySerializer;
     }
 
     @Override
@@ -99,14 +87,7 @@ public void combine(WindowKey windowKey, Iterator<RowData> localAccs) throws Exc
 
     private void combineAccumulator(WindowKey windowKey, RowData acc) throws Exception {
         // step 1: set current key for states and timers
-        final RowData key;
-        if (requiresCopy) {
-            // the incoming key is reused, we should copy it if state backend doesn't copy it
-            key = keySerializer.copy(windowKey.getKey());
-        } else {
-            key = windowKey.getKey();
-        }
-        keyContext.setCurrentKey(key);
+        keyContext.setCurrentKey(windowKey.getKey());
         Long window = windowKey.getWindow();
 
         // step2: merge acc into state
@@ -145,15 +126,12 @@ public static final class Factory implements RecordsCombiner.Factory {
 
         private final GeneratedNamespaceAggsHandleFunction<Long> genLocalAggsHandler;
         private final GeneratedNamespaceAggsHandleFunction<Long> genGlobalAggsHandler;
-        private final TypeSerializer<RowData> keySerializer;
 
         public Factory(
                 GeneratedNamespaceAggsHandleFunction<Long> genLocalAggsHandler,
-                GeneratedNamespaceAggsHandleFunction<Long> genGlobalAggsHandler,
-                TypeSerializer<RowData> keySerializer) {
+                GeneratedNamespaceAggsHandleFunction<Long> genGlobalAggsHandler) {
             this.genLocalAggsHandler = genLocalAggsHandler;
             this.genGlobalAggsHandler = genGlobalAggsHandler;
-            this.keySerializer = keySerializer;
         }
 
         @Override
@@ -174,16 +152,13 @@ public RecordsCombiner createRecordsCombiner(
             globalAggregator.open(
                     new PerWindowStateDataViewStore(
                             stateBackend, LongSerializer.INSTANCE, runtimeContext));
-            boolean requiresCopy = !isStateImmutableInStateBackend(stateBackend);
             WindowValueState<Long> windowValueState = (WindowValueState<Long>) windowState;
             return new GlobalAggCombiner(
                     timerService,
                     stateBackend::setCurrentKey,
                     windowValueState,
                     localAggregator,
-                    globalAggregator,
-                    requiresCopy,
-                    keySerializer);
+                    globalAggregator);
         }
     }
 }

flink-table/flink-table-runtime-blink/src/main/java/org/apache/flink/table/runtime/operators/rank/window/WindowRankOperatorBuilder.java

@@ -138,11 +138,7 @@ public WindowRankOperatorBuilder windowEndIndex(int windowEndIndex) {
                         "Illegal window end index %s, it should not be negative!", windowEndIndex));
         final RecordsCombiner.Factory combinerFactory =
                 new TopNRecordsCombiner.Factory(
-                        generatedSortKeyComparator,
-                        sortKeySelector,
-                        keySerializer,
-                        inputSerializer,
-                        rankEnd);
+                        generatedSortKeyComparator, sortKeySelector, inputSerializer, rankEnd);
         final WindowBuffer.Factory bufferFactory =
                 new RecordsWindowBuffer.Factory(keySerializer, inputSerializer, combinerFactory);
         final SlicingWindowProcessor<Long> windowProcessor =

flink-table/flink-table-runtime-blink/src/main/java/org/apache/flink/table/runtime/operators/rank/window/combines/TopNRecordsCombiner.java

@@ -41,7 +41,6 @@
 import java.util.Map;
 
 import static org.apache.flink.table.data.util.RowDataUtil.isAccumulateMsg;
-import static org.apache.flink.table.runtime.util.StateConfigUtil.isStateImmutableInStateBackend;
 
 /**
  * An implementation of {@link RecordsCombiner} that save topN records of incremental input records
@@ -67,12 +66,6 @@ public final class TopNRecordsCombiner implements RecordsCombiner {
     /** TopN size. */
     private final long topN;
 
-    /** Whether to copy input key, because key is reused. */
-    private final boolean requiresCopyKey;
-
-    /** Serializer to copy key if required. */
-    private final TypeSerializer<RowData> keySerializer;
-
     /** Serializer to copy record if required. */
     private final TypeSerializer<RowData> recordSerializer;
 
@@ -86,8 +79,6 @@ public TopNRecordsCombiner(
             Comparator<RowData> sortKeyComparator,
             KeySelector<RowData, RowData> sortKeySelector,
             long topN,
-            boolean requiresCopyKey,
-            TypeSerializer<RowData> keySerializer,
             TypeSerializer<RowData> recordSerializer,
             boolean isEventTime) {
         this.timerService = timerService;
@@ -96,8 +87,6 @@ public TopNRecordsCombiner(
         this.sortKeyComparator = sortKeyComparator;
         this.sortKeySelector = sortKeySelector;
         this.topN = topN;
-        this.requiresCopyKey = requiresCopyKey;
-        this.keySerializer = keySerializer;
         this.recordSerializer = recordSerializer;
         this.isEventTime = isEventTime;
     }
@@ -123,14 +112,7 @@ public void combine(WindowKey windowKey, Iterator<RowData> records) throws Excep
 
         // step 2: flush data in TopNBuffer into state
         Iterator<Map.Entry<RowData, Collection<RowData>>> bufferItr = buffer.entrySet().iterator();
-        final RowData key;
-        if (requiresCopyKey) {
-            // the incoming key is reused, we should copy it if state backend doesn't copy it
-            key = keySerializer.copy(windowKey.getKey());
-        } else {
-            key = windowKey.getKey();
-        }
-        keyContext.setCurrentKey(key);
+        keyContext.setCurrentKey(windowKey.getKey());
         Long window = windowKey.getWindow();
         while (bufferItr.hasNext()) {
             Map.Entry<RowData, Collection<RowData>> entry = bufferItr.next();
@@ -165,19 +147,16 @@ public static final class Factory implements RecordsCombiner.Factory {
         // The util to compare two sortKey equals to each other.
         private final GeneratedRecordComparator generatedSortKeyComparator;
         private final KeySelector<RowData, RowData> sortKeySelector;
-        private final TypeSerializer<RowData> keySerializer;
         private final TypeSerializer<RowData> recordSerializer;
         private final long topN;
 
         public Factory(
                 GeneratedRecordComparator genSortKeyComparator,
                 RowDataKeySelector sortKeySelector,
-                TypeSerializer<RowData> keySerializer,
                 TypeSerializer<RowData> recordSerializer,
                 long topN) {
             this.generatedSortKeyComparator = genSortKeyComparator;
             this.sortKeySelector = sortKeySelector;
-            this.keySerializer = keySerializer;
             this.recordSerializer = recordSerializer;
             this.topN = topN;
         }
@@ -192,7 +171,6 @@ public RecordsCombiner createRecordsCombiner(
                 throws Exception {
             final Comparator<RowData> sortKeyComparator =
                     generatedSortKeyComparator.newInstance(runtimeContext.getUserCodeClassLoader());
-            boolean requiresCopyKey = !isStateImmutableInStateBackend(stateBackend);
             WindowMapState<Long, List<RowData>> windowMapState =
                     (WindowMapState<Long, List<RowData>>) windowState;
             return new TopNRecordsCombiner(
@@ -202,8 +180,6 @@ public RecordsCombiner createRecordsCombiner(
                     sortKeyComparator,
                     sortKeySelector,
                     topN,
-                    requiresCopyKey,
-                    keySerializer,
                     recordSerializer,
                     isEventTime);
         }