[flink] Support BUCKET_LOAD_BALANCE shuffle strategy for table with bucket kyes.

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/FlinkConnectorOptions.java

@@ -232,6 +232,7 @@ public class FlinkConnectorOptions {
                     .defaultValue(DistributionMode.AUTO)
                     .withDescription(
                             "Defines the distribution mode for writing data to the sink. Available options are:\n"
+                                    + "- BUCKET_LOAD_BALANCE: Shuffle data by bucket key with even load distribution across all downstream subtasks using LCM-based logical slot assignment. Records with the same bucket key always route to the same subtask. Requires 'bucket.key' to be defined. Suitable for tables where intra-bucket ordering is not required but even load distribution matters. Not supported for tables with the aggregation merge engine, because Undo Recovery requires each bucket to be written by exactly one subtask.\n"
                                     + "- AUTO: Automatically chooses the best mode based on the table type. "
                                     + "Uses BUCKET mode for Primary Key Tables and Log table with bucket key to maximize throughput, "
                                     + "and NONE for Log Tables without bucket key.\n"

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/shuffle/BucketLoadBalanceRouter.java

@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.fluss.flink.shuffle;
+
+import org.apache.fluss.annotation.Internal;
+import org.apache.fluss.bucketing.BucketingFunction;
+import org.apache.fluss.row.InternalRow;
+import org.apache.fluss.row.encode.KeyEncoder;
+import org.apache.fluss.utils.MathUtils;
+
+import javax.annotation.Nullable;
+
+import java.util.Arrays;
+
+/**
+ * Shuffle data by bucket key with even load distribution across all downstream subtasks.
+ *
+ * <p>Two callers share this router and must stay consistent:
+ *
+ * <ul>
+ *   <li>The sink shuffle ({@code BucketLoadBalanceChannelComputer}) routes records to sink
+ *       subtasks. It passes the encoded bucket key as both the bucket key and the hash key, so
+ *       records with the same bucket key always land on the same subtask.
+ *   <li>The lookup-join partitioner ({@code FlussLookupInputPartitioner}) routes probe rows to
+ *       lookup subtasks. It passes the encoded bucket key for bucketing and the encoded lookup key
+ *       for the intra-bucket slot hash, spreading different lookup keys of the same bucket across
+ *       the bucket's subtask range.
+ * </ul>
+ *
+ * <p>Routing algorithm: first compute the bucket ID via {@code BucketingFunction}, then select the
+ * concrete subtask within the bucket's assigned slot range using the murmur hash of the hash key.
+ * When the bucket count is evenly divisible by the channel count the hash is skipped entirely and
+ * the result is simply {@code bucketId % numChannels}, matching the behaviour of the BUCKET
+ * distribution mode.
+ */
+@Internal
+public final class BucketLoadBalanceRouter {
+
+    private final KeyEncoder bucketKeyEncoder;
+    private final @Nullable KeyEncoder hashKeyEncoder;
+    private final BucketingFunction bucketingFunction;
+    private final int numBuckets;
+
+    /**
+     * Creates a router for the BUCKET_LOAD_BALANCE distribution strategy.
+     *
+     * @param bucketKeyEncoder encoder for the bucket key (determines bucket ID)
+     * @param hashKeyEncoder encoder for the intra-bucket slot hash; may be {@code null} when the
+     *     caller wants to use the bucket key itself as the hash key (e.g. sink shuffle)
+     * @param bucketingFunction the bucketing function for computing the bucket ID
+     * @param numBuckets total number of buckets
+     */
+    public BucketLoadBalanceRouter(
+            KeyEncoder bucketKeyEncoder,
+            @Nullable KeyEncoder hashKeyEncoder,
+            BucketingFunction bucketingFunction,
+            int numBuckets) {
+        this.bucketKeyEncoder = bucketKeyEncoder;
+        this.hashKeyEncoder = hashKeyEncoder;
+        this.bucketingFunction = bucketingFunction;
+        this.numBuckets = numBuckets;
+    }
+
+    /**
+     * Routes a record to a downstream channel using bucket-ID + LCM-based slot assignment.
+     *
+     * <p>The record is first encoded by {@code bucketKeyEncoder} to obtain the bucket key. When
+     * {@link #isEvenlyDivisible} holds, the result is simply {@code bucketId % numChannels} and the
+     * hash key is never encoded. Otherwise the intra-bucket slot hash is computed from the hash key
+     * (via {@code hashKeyEncoder} when present, otherwise the bucket key is reused).
+     *
+     * @param row the record to route
+     * @param numChannels number of downstream channels (subtask parallelism)
+     * @return channel index in {@code [0, numChannels)}
+     */
+    public int route(InternalRow row, int numChannels) {
+        byte[] bucketKeyBytes = bucketKeyEncoder.encodeKey(row);
+        int bucketId = bucketingFunction.bucketing(bucketKeyBytes, numBuckets);
+
+        // Fast-path: when bucket count is an exact multiple of channel count, every bucket
+        // maps to a single channel and no hash is needed.
+        if (isEvenlyDivisible(numBuckets, numChannels)) {
+            return bucketId % numChannels;
+        }
+
+        byte[] hashKeyBytes =
+                hashKeyEncoder != null ? hashKeyEncoder.encodeKey(row) : bucketKeyBytes;
+
+        int recordHash = MathUtils.murmurHash(Arrays.hashCode(hashKeyBytes));
+        return selectSlot(bucketId, numBuckets, numChannels, recordHash);
+    }
+
+    private static int selectSlot(int bucketId, int numBuckets, int numChannels, int recordHash) {
+        if (numChannels > numBuckets && numChannels % numBuckets == 0) {
+            int candidateCount = numChannels / numBuckets;
+            return (recordHash % candidateCount) * numBuckets + bucketId;
+        }
+
+        int gcd = greatestCommonDivisor(numBuckets, numChannels);
+        int slotsPerBucket = numChannels / gcd;
+        int slotsPerSubtask = numBuckets / gcd;
+        int slotWithinBucket = recordHash % slotsPerBucket;
+        long logicalSlot = (long) bucketId * slotsPerBucket + slotWithinBucket;
+        return (int) (logicalSlot / slotsPerSubtask);
+    }
+
+    private static int greatestCommonDivisor(int first, int second) {
+        while (second != 0) {
+            int remainder = first % second;
+            first = second;
+            second = remainder;
+        }
+        return first;
+    }
+
+    private static boolean isEvenlyDivisible(int numBuckets, int numChannels) {
+        return numBuckets >= numChannels && numBuckets % numChannels == 0;
+    }
+}

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/BucketLoadBalanceChannelComputer.java

@@ -0,0 +1,116 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.fluss.flink.sink;
+
+import org.apache.fluss.annotation.Internal;
+import org.apache.fluss.bucketing.BucketingFunction;
+import org.apache.fluss.exception.FlussRuntimeException;
+import org.apache.fluss.flink.row.RowWithOp;
+import org.apache.fluss.flink.shuffle.BucketLoadBalanceRouter;
+import org.apache.fluss.flink.sink.serializer.FlussSerializationSchema;
+import org.apache.fluss.flink.sink.serializer.SerializerInitContextImpl;
+import org.apache.fluss.metadata.DataLakeFormat;
+import org.apache.fluss.row.InternalRow;
+import org.apache.fluss.row.encode.KeyEncoder;
+import org.apache.fluss.types.RowType;
+
+import javax.annotation.Nullable;
+
+import java.util.List;
+
+/**
+ * {@link ChannelComputer} for {@link
+ * org.apache.fluss.flink.sink.shuffle.DistributionMode#BUCKET_LOAD_BALANCE}.
+ *
+ * <p>Distributes records by bucket key with even load across all downstream channels. Unlike {@link
+ * FlinkRowDataChannelComputer} (BUCKET mode) which maps each bucket to a single channel, this
+ * computer uses an LCM-based logical slot assignment so that every channel receives traffic from
+ * every bucket.
+ *
+ * <p>The routing algorithm is shared with the source-side lookup-join partitioner ({@code
+ * FlussLookupInputPartitioner}); see {@link BucketLoadBalanceRouter}.
+ *
+ * <p>Records with the same bucket key always route to the same channel.
+ *
+ * @param <InputT> the type of records
+ */
+@Internal
+public class BucketLoadBalanceChannelComputer<InputT> implements ChannelComputer<InputT> {
+
+    private static final long serialVersionUID = 1L;
+
+    private final @Nullable DataLakeFormat lakeFormat;
+    private final int numBucket;
+    private final RowType flussRowType;
+    private final List<String> bucketKeys;
+    private final FlussSerializationSchema<InputT> serializationSchema;
+
+    private transient int numChannels;
+    private transient BucketLoadBalanceRouter bucketLoadBalanceRouter;
+
+    public BucketLoadBalanceChannelComputer(
+            RowType flussRowType,
+            List<String> bucketKeys,
+            @Nullable DataLakeFormat lakeFormat,
+            int numBucket,
+            FlussSerializationSchema<InputT> serializationSchema) {
+        this.flussRowType = flussRowType;
+        this.bucketKeys = bucketKeys;
+        this.lakeFormat = lakeFormat;
+        this.numBucket = numBucket;
+        this.serializationSchema = serializationSchema;
+    }
+
+    @Override
+    public void setup(int numChannels) {
+        this.numChannels = numChannels;
+        this.bucketLoadBalanceRouter =
+                new BucketLoadBalanceRouter(
+                        KeyEncoder.ofBucketKeyEncoder(flussRowType, bucketKeys, lakeFormat),
+                        null,
+                        BucketingFunction.of(lakeFormat),
+                        numBucket);
+
+        try {
+            this.serializationSchema.open(new SerializerInitContextImpl(flussRowType, false));
+        } catch (Exception e) {
+            throw new FlussRuntimeException(e);
+        }
+    }
+
+    @Override
+    public int channel(InputT record) {
+        try {
+            RowWithOp rowWithOp = serializationSchema.serialize(record);
+            InternalRow row = rowWithOp.getRow();
+            return bucketLoadBalanceRouter.route(row, numChannels);
+        } catch (Exception e) {
+            throw new FlussRuntimeException(
+                    String.format(
+                            "Failed to serialize record of type '%s'"
+                                    + " in BucketLoadBalanceChannelComputer: %s",
+                            record != null ? record.getClass().getName() : "null", e.getMessage()),
+                    e);
+        }
+    }
+
+    @Override
+    public String toString() {
+        return "BUCKET_LOAD_BALANCE";
+    }
+}

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/FlinkSink.java

@@ -53,6 +53,7 @@
 import java.util.List;
 
 import static org.apache.fluss.flink.sink.FlinkStreamPartitioner.partition;
+import static org.apache.fluss.flink.utils.FlinkConnectorOptionsUtils.validateDistributionModeForUndoRecovery;
 import static org.apache.fluss.flink.utils.FlinkConversions.toFlussRowType;
 import static org.apache.fluss.utils.Preconditions.checkState;
 
@@ -173,6 +174,12 @@ public DataStream<InputT> addPreWriteTopology(DataStream<InputT> input) {
                     }
                     throw new UnsupportedOperationException(
                             "BUCKET mode is only supported for log tables with bucket keys");
+                case BUCKET_LOAD_BALANCE:
+                    if (!bucketKeys.isEmpty()) {
+                        return bucketLoadBalanceShuffle(input);
+                    }
+                    throw new UnsupportedOperationException(
+                            "BUCKET_LOAD_BALANCE mode is only supported for log tables with bucket keys");
                 case PARTITION_DYNAMIC:
                     if (partitionKeys.isEmpty()) {
                         throw new UnsupportedOperationException(
@@ -234,6 +241,18 @@ private DataStream<InputT> bucketShuffle(DataStream<InputT> input) {
                             flussSerializationSchema),
                     input.getParallelism());
         }
+
+        private DataStream<InputT> bucketLoadBalanceShuffle(DataStream<InputT> input) {
+            return partition(
+                    input,
+                    new BucketLoadBalanceChannelComputer<>(
+                            toFlussRowType(tableRowType),
+                            bucketKeys,
+                            lakeFormat,
+                            numBucket,
+                            flussSerializationSchema),
+                    input.getParallelism());
+        }
     }
 
     @Internal
@@ -317,6 +336,11 @@ public UpsertSinkWriter<InputT> createWriter(
 
         @Override
         public DataStream<InputT> addPreWriteTopology(DataStream<InputT> input) {
+            // Defense in depth: the entry points (FlinkTableFactory and FlussSinkBuilder)
+            // validate the distribution mode via validateDistributionModeForMergeEngine before
+            // constructing this builder, but the builder can also be constructed directly, so
+            // re-validate the Undo Recovery constraint here.
+            validateDistributionModeForUndoRecovery(enableUndoRecovery, distributionMode);
             DataStream<InputT> stream;
             switch (distributionMode) {
                 case NONE:
@@ -336,6 +360,18 @@ public DataStream<InputT> addPreWriteTopology(DataStream<InputT> input) {
                                             flussSerializationSchema),
                                     input.getParallelism());
                     break;
+                case BUCKET_LOAD_BALANCE:
+                    stream =
+                            partition(
+                                    input,
+                                    new BucketLoadBalanceChannelComputer<>(
+                                            toFlussRowType(tableRowType),
+                                            bucketKeys,
+                                            lakeFormat,
+                                            numBucket,
+                                            flussSerializationSchema),
+                                    input.getParallelism());
+                    break;
                 default:
                     throw new UnsupportedOperationException(
                             String.format(

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/shuffle/DistributionMode.java

@@ -65,5 +65,27 @@ public enum DistributionMode {
      * <p>Note: This mode has overhead costs including data statistics collection and additional
      * shuffle operations.
      */
-    PARTITION_DYNAMIC
+    PARTITION_DYNAMIC,
+
+    /**
+     * Shuffle data by bucket key with even load distribution across all downstream subtasks.
+     *
+     * <p>Unlike {@link #BUCKET} which maps each bucket to exactly one subtask (potentially leaving
+     * some subtasks idle), this mode uses an LCM-based logical slot assignment to ensure every
+     * subtask receives data.
+     *
+     * <p>Routing algorithm: first compute bucket ID via {@code BucketingFunction}, then select the
+     * concrete subtask within the bucket's assigned slot range using the bucket key's murmur hash.
+     * Records with the same bucket key always route to the same subtask, so merge semantics are
+     * preserved at runtime. However, one bucket may be written by several subtasks when the bucket
+     * count and the sink parallelism do not evenly divide each other, which breaks the
+     * one-writer-per-bucket assumption of Undo Recovery; this mode is therefore rejected for tables
+     * using the aggregation merge engine.
+     *
+     * <p>Characteristics:
+     *
+     * <p>Requires 'bucket.key' to be defined. Suitable for tables where intra-bucket ordering is
+     * not required but even load distribution matters.
+     */
+    BUCKET_LOAD_BALANCE
 }