[flink] Balance MonitorFunction data shuffle for partitioned table (#…

…3180)
apache · Apr 9, 2024 · f067979 · f067979
1 parent a29c0a6
commit f067979
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/...n-flink-common/src/main/java/org/apache/paimon/flink/source/operator/MonitorFunction.java b/...n-flink-common/src/main/java/org/apache/paimon/flink/source/operator/MonitorFunction.java
@@ -19,6 +19,7 @@
 package org.apache.paimon.flink.source.operator;
 
 import org.apache.paimon.flink.utils.JavaTypeInfo;
+import org.apache.paimon.table.sink.ChannelComputer;
 import org.apache.paimon.table.source.DataSplit;
 import org.apache.paimon.table.source.EndOfScanException;
 import org.apache.paimon.table.source.ReadBuilder;
@@ -236,8 +237,12 @@ public static DataStream<RowData> buildSource(
                         new JavaTypeInfo<>(Split.class))
                 .forceNonParallel()
                 .partitionCustom(
-                        (key, numPartitions) -> key % numPartitions,
-                        split -> ((DataSplit) split).bucket())
+                        (key, numPartitions) ->
+                                ChannelComputer.select(key.f0, key.f1, numPartitions),
+                        split -> {
+                            DataSplit dataSplit = (DataSplit) split;
+                            return Tuple2.of(dataSplit.partition(), dataSplit.bucket());
+                        })
                 .transform(name + "-Reader", typeInfo, new ReadOperator(readBuilder));
     }
 }