[flink] Flink sink support hash by bucket key for PrimaryKey Table

Author: swuferhong

fluss-common/src/main/java/com/alibaba/fluss/bucketing/BucketingFunction.java

@@ -20,8 +20,10 @@
 
 import javax.annotation.Nullable;
 
+import java.io.Serializable;
+
 /** An interface to assign a bucket according to the bucket key byte array. */
-public interface BucketingFunction {
+public interface BucketingFunction extends Serializable {
 
     /**
      * Assign a bucket according to the bucket key byte array.

fluss-common/src/main/java/com/alibaba/fluss/memory/MemorySegment.java

@@ -23,6 +23,7 @@
 import java.io.DataInput;
 import java.io.DataOutput;
 import java.io.IOException;
+import java.io.Serializable;
 import java.nio.BufferOverflowException;
 import java.nio.BufferUnderflowException;
 import java.nio.ByteBuffer;
@@ -64,7 +65,7 @@
  * implementations on invocations of abstract methods.
  */
 @Internal
-public final class MemorySegment {
+public final class MemorySegment implements Serializable {
 
     /** The unsafe handle for transparent memory copied (heap / off-heap). */
     private static final sun.misc.Unsafe UNSAFE = MemoryUtils.UNSAFE;

fluss-common/src/main/java/com/alibaba/fluss/row/compacted/CompactedRowWriter.java

@@ -64,7 +64,7 @@
  * the positive integer is doubled. We assume that the probability of general integers being
  * positive is higher, so sacrifice the negative number to promote the positive number.
  */
-public class CompactedRowWriter {
+public class CompactedRowWriter implements Serializable {
 
     private final int headerSizeInBytes;
 

fluss-common/src/main/java/com/alibaba/fluss/row/encode/KeyEncoder.java

@@ -23,10 +23,11 @@
 
 import javax.annotation.Nullable;
 
+import java.io.Serializable;
 import java.util.List;
 
 /** An interface for encoding key of row into bytes. */
-public interface KeyEncoder {
+public interface KeyEncoder extends Serializable {
 
     /** Encode the key of given row to byte array. */
     byte[] encodeKey(InternalRow row);

fluss-connectors/fluss-connector-flink/src/main/java/com/alibaba/fluss/connector/flink/FlinkConnectorOptions.java

@@ -106,6 +106,15 @@ public class FlinkConnectorOptions {
                     .defaultValue(false)
                     .withDescription("Whether to ignore retract（-U/-D) record.");
 
+    public static final ConfigOption<Boolean> SINK_PRE_HASH_BY_BUCKET_KEY =
+            ConfigOptions.key("sink.pre-hash-by-bucket-key")
+                    .booleanType()
+                    .defaultValue(false)
+                    .withDescription(
+                            "Whether to pre hash the record by bucket key before write to sink. Hash the data with the "
+                                    + "same bucket key to be processed by the same task can improve the efficiency"
+                                    + " of client processing and reduce resource consumption");
+
     // --------------------------------------------------------------------------------------------
     // table storage specific options
     // --------------------------------------------------------------------------------------------

fluss-connectors/fluss-connector-flink/src/main/java/com/alibaba/fluss/connector/flink/catalog/FlinkTableFactory.java

@@ -52,6 +52,7 @@
 import java.util.Set;
 
 import static com.alibaba.fluss.connector.flink.catalog.FlinkCatalog.LAKE_TABLE_SPLITTER;
+import static com.alibaba.fluss.connector.flink.utils.FlinkConnectorOptionsUtils.getBucketKeys;
 import static com.alibaba.fluss.connector.flink.utils.FlinkConversions.toFlinkOption;
 
 /** Factory to create table source and table sink for Fluss. */
@@ -151,7 +152,10 @@ public DynamicTableSink createDynamicTableSink(Context context) {
                 context.getPrimaryKeyIndexes(),
                 isStreamingMode,
                 tableOptions.get(toFlinkOption(ConfigOptions.TABLE_MERGE_ENGINE)),
-                tableOptions.get(FlinkConnectorOptions.SINK_IGNORE_DELETE));
+                tableOptions.get(FlinkConnectorOptions.SINK_IGNORE_DELETE),
+                tableOptions.get(FlinkConnectorOptions.BUCKET_NUMBER),
+                getBucketKeys(tableOptions),
+                tableOptions.get(FlinkConnectorOptions.SINK_PRE_HASH_BY_BUCKET_KEY));
     }
 
     @Override
@@ -176,6 +180,7 @@ public Set<ConfigOption<?>> optionalOptions() {
                                 FlinkConnectorOptions.SCAN_PARTITION_DISCOVERY_INTERVAL,
                                 FlinkConnectorOptions.LOOKUP_ASYNC,
                                 FlinkConnectorOptions.SINK_IGNORE_DELETE,
+                                FlinkConnectorOptions.SINK_PRE_HASH_BY_BUCKET_KEY,
                                 LookupOptions.MAX_RETRIES,
                                 LookupOptions.CACHE_TYPE,
                                 LookupOptions.PARTIAL_CACHE_EXPIRE_AFTER_ACCESS,

fluss-connectors/fluss-connector-flink/src/main/java/com/alibaba/fluss/connector/flink/sink/FlinkSink.java

@@ -22,11 +22,14 @@
 import com.alibaba.fluss.connector.flink.sink.writer.FlinkSinkWriter;
 import com.alibaba.fluss.connector.flink.sink.writer.UpsertSinkWriter;
 import com.alibaba.fluss.metadata.TablePath;
+import com.alibaba.fluss.row.encode.KeyEncoder;
 
 import org.apache.flink.api.connector.sink2.Sink;
 import org.apache.flink.api.connector.sink2.SinkWriter;
 import org.apache.flink.api.connector.sink2.WriterInitContext;
 import org.apache.flink.runtime.metrics.groups.InternalSinkWriterMetricGroup;
+import org.apache.flink.streaming.api.connector.sink2.WithPreWriteTopology;
+import org.apache.flink.streaming.api.datastream.DataStream;
 import org.apache.flink.table.data.RowData;
 import org.apache.flink.table.types.logical.RowType;
 
@@ -35,8 +38,12 @@
 import java.io.IOException;
 import java.io.Serializable;
 
-/** Flink sink for Fluss. */
-class FlinkSink implements Sink<RowData> {
+/**
+ * Flink sink for Fluss.
+ *
+ * <p>TODO: WithPreWriteTopology need to be changed to supportsPreWriteTopology in Flink 1.20
+ */
+class FlinkSink implements Sink<RowData>, WithPreWriteTopology<RowData> {
 
     private static final long serialVersionUID = 1L;
 
@@ -60,9 +67,16 @@ public SinkWriter<RowData> createWriter(WriterInitContext context) throws IOExce
         return flinkSinkWriter;
     }
 
+    @Override
+    public DataStream<RowData> addPreWriteTopology(DataStream<RowData> input) {
+        return builder.addPreWriteTopology(input);
+    }
+
     @Internal
     interface SinkWriterBuilder<W extends FlinkSinkWriter> extends Serializable {
         W createWriter();
+
+        DataStream<RowData> addPreWriteTopology(DataStream<RowData> input);
     }
 
     @Internal
@@ -90,6 +104,11 @@ public AppendSinkWriterBuilder(
         public AppendSinkWriter createWriter() {
             return new AppendSinkWriter(tablePath, flussConfig, tableRowType, ignoreDelete);
         }
+
+        @Override
+        public DataStream<RowData> addPreWriteTopology(DataStream<RowData> input) {
+            return input;
+        }
     }
 
     @Internal
@@ -102,24 +121,42 @@ static class UpsertSinkWriterBuilder implements SinkWriterBuilder<UpsertSinkWrit
         private final RowType tableRowType;
         private final @Nullable int[] targetColumnIndexes;
         private final boolean ignoreDelete;
+        private final int numBucket;
+        private final KeyEncoder bucketKeyEncoder;
+        private final boolean sinkReHash;
 
         UpsertSinkWriterBuilder(
                 TablePath tablePath,
                 Configuration flussConfig,
                 RowType tableRowType,
                 @Nullable int[] targetColumnIndexes,
-                boolean ignoreDelete) {
+                boolean ignoreDelete,
+                int numBucket,
+                KeyEncoder bucketKeyEncoder,
+                boolean sinkReHash) {
             this.tablePath = tablePath;
             this.flussConfig = flussConfig;
             this.tableRowType = tableRowType;
             this.targetColumnIndexes = targetColumnIndexes;
             this.ignoreDelete = ignoreDelete;
+            this.numBucket = numBucket;
+            this.bucketKeyEncoder = bucketKeyEncoder;
+            this.sinkReHash = sinkReHash;
         }
 
         @Override
         public UpsertSinkWriter createWriter() {
             return new UpsertSinkWriter(
                     tablePath, flussConfig, tableRowType, targetColumnIndexes, ignoreDelete);
         }
+
+        @Override
+        public DataStream<RowData> addPreWriteTopology(DataStream<RowData> input) {
+            return sinkReHash
+                    ? input.partitionCustom(
+                            (bucketId, numPartitions) -> bucketId % numPartitions,
+                            new RowDataKeySelector(bucketKeyEncoder, numBucket))
+                    : input;
+        }
     }
 }

fluss-connectors/fluss-connector-flink/src/main/java/com/alibaba/fluss/connector/flink/sink/FlinkTableSink.java

@@ -24,6 +24,7 @@
 import com.alibaba.fluss.metadata.MergeEngineType;
 import com.alibaba.fluss.metadata.TablePath;
 import com.alibaba.fluss.row.GenericRow;
+import com.alibaba.fluss.row.encode.KeyEncoder;
 
 import org.apache.flink.table.api.ValidationException;
 import org.apache.flink.table.catalog.Column;
@@ -52,6 +53,8 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import static com.alibaba.fluss.connector.flink.utils.FlinkConversions.toFlussRowType;
+
 /** A Flink {@link DynamicTableSink}. */
 public class FlinkTableSink
         implements DynamicTableSink,
@@ -67,6 +70,9 @@ public class FlinkTableSink
     private final boolean streaming;
     @Nullable private final MergeEngineType mergeEngineType;
     private final boolean ignoreDelete;
+    private final int numBucket;
+    private final List<String> bucketKeys;
+    private final boolean sinkReHash;
 
     private boolean appliedUpdates = false;
     @Nullable private GenericRow deleteRow;
@@ -78,14 +84,20 @@ public FlinkTableSink(
             int[] primaryKeyIndexes,
             boolean streaming,
             @Nullable MergeEngineType mergeEngineType,
-            boolean ignoreDelete) {
+            boolean ignoreDelete,
+            int numBucket,
+            List<String> bucketKeys,
+            boolean sinkReHash) {
         this.tablePath = tablePath;
         this.flussConfig = flussConfig;
         this.tableRowType = tableRowType;
         this.primaryKeyIndexes = primaryKeyIndexes;
         this.streaming = streaming;
         this.mergeEngineType = mergeEngineType;
         this.ignoreDelete = ignoreDelete;
+        this.numBucket = numBucket;
+        this.bucketKeys = bucketKeys;
+        this.sinkReHash = sinkReHash;
     }
 
     @Override
@@ -168,7 +180,10 @@ public SinkRuntimeProvider getSinkRuntimeProvider(Context context) {
                                 flussConfig,
                                 tableRowType,
                                 targetColumnIndexes,
-                                ignoreDelete)
+                                ignoreDelete,
+                                numBucket,
+                                KeyEncoder.of(toFlussRowType(tableRowType), bucketKeys, null),
+                                sinkReHash)
                         : new FlinkSink.AppendSinkWriterBuilder(
                                 tablePath, flussConfig, tableRowType, ignoreDelete);
 
@@ -195,7 +210,10 @@ public DynamicTableSink copy() {
                         primaryKeyIndexes,
                         streaming,
                         mergeEngineType,
-                        ignoreDelete);
+                        ignoreDelete,
+                        numBucket,
+                        bucketKeys,
+                        sinkReHash);
         sink.appliedUpdates = appliedUpdates;
         sink.deleteRow = deleteRow;
         return sink;

fluss-connectors/fluss-connector-flink/src/main/java/com/alibaba/fluss/connector/flink/sink/RowDataKeySelector.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) 2025 Alibaba Group Holding Ltd.
+ *
+ * Licensed 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 com.alibaba.fluss.connector.flink.sink;
+
+import com.alibaba.fluss.bucketing.FlussBucketingFunction;
+import com.alibaba.fluss.connector.flink.row.FlinkAsFlussRow;
+import com.alibaba.fluss.row.encode.KeyEncoder;
+
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.table.data.RowData;
+
+/** {@link KeySelector} to get bucket id. */
+public class RowDataKeySelector implements KeySelector<RowData, Integer> {
+
+    private final FlussBucketingFunction flussBucketingFunction;
+    private final KeyEncoder bucketKeyEncoder;
+    private final int numBucket;
+
+    public RowDataKeySelector(KeyEncoder bucketKeyEncoder, int numBucket) {
+        this.bucketKeyEncoder = bucketKeyEncoder;
+        this.flussBucketingFunction = new FlussBucketingFunction();
+        this.numBucket = numBucket;
+    }
+
+    @Override
+    public Integer getKey(RowData rowData) throws Exception {
+        FlinkAsFlussRow row = new FlinkAsFlussRow().replace(rowData);
+        return flussBucketingFunction.bucketing(bucketKeyEncoder.encodeKey(row), numBucket);
+    }
+}

fluss-connectors/fluss-connector-flink/src/main/java/com/alibaba/fluss/connector/flink/utils/FlinkConnectorOptionsUtils.java

@@ -27,7 +27,11 @@
 import java.time.LocalDateTime;
 import java.time.ZoneId;
 import java.time.format.DateTimeFormatter;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
 import java.util.Optional;
+import java.util.stream.Collectors;
 
 import static com.alibaba.fluss.connector.flink.FlinkConnectorOptions.SCAN_STARTUP_MODE;
 import static com.alibaba.fluss.connector.flink.FlinkConnectorOptions.SCAN_STARTUP_TIMESTAMP;
@@ -63,6 +67,17 @@ public static StartupOptions getStartupOptions(ReadableConfig tableOptions, Zone
         return options;
     }
 
+    public static List<String> getBucketKeys(ReadableConfig tableOptions) {
+        Optional<String> bucketKey = tableOptions.getOptional(FlinkConnectorOptions.BUCKET_KEY);
+        if (!bucketKey.isPresent()) {
+            // log tables don't have bucket key by default
+            return new ArrayList<>();
+        }
+
+        String[] keys = bucketKey.get().split(",");
+        return Arrays.stream(keys).collect(Collectors.toList());
+    }
+
     public static int[] getBucketKeyIndexes(ReadableConfig tableOptions, RowType schema) {
         Optional<String> bucketKey = tableOptions.getOptional(FlinkConnectorOptions.BUCKET_KEY);
         if (!bucketKey.isPresent()) {