From 1a461a9b806d1865722a771c87f0fbbef9b10e13 Mon Sep 17 00:00:00 2001 From: Hongshun Wang Date: Sun, 20 Sep 2026 20:56:48 +0800 Subject: [PATCH 1/2] [flink] Support BUCKET_LOAD_BALANCE shuffle strategy for table with bucket kyes. --- .../fluss/flink/FlinkConnectorOptions.java | 1 + .../BucketLoadBalanceChannelComputer.java | 167 ++++++++++++++++ .../fluss/flink/sink/ChannelComputer.java | 51 +++++ .../apache/fluss/flink/sink/FlinkSink.java | 30 +++ .../flink/sink/shuffle/DistributionMode.java | 21 +- .../lookup/FlussLookupInputPartitioner.java | 53 +---- .../utils/FlinkConnectorOptionsUtils.java | 5 +- .../BucketLoadBalanceChannelComputerTest.java | 186 ++++++++++++++++++ .../flink/sink/FlinkTableSinkITCase.java | 10 +- .../flink/sink/FlussSinkBuilderTest.java | 4 + website/docs/engine-flink/options.md | 10 +- 11 files changed, 488 insertions(+), 50 deletions(-) create mode 100644 fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/BucketLoadBalanceChannelComputer.java create mode 100644 fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/BucketLoadBalanceChannelComputerTest.java diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/FlinkConnectorOptions.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/FlinkConnectorOptions.java index bc7529fe23c..b22715d10a2 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/FlinkConnectorOptions.java +++ b/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.\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" diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/BucketLoadBalanceChannelComputer.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/BucketLoadBalanceChannelComputer.java new file mode 100644 index 00000000000..edd70dfe1bd --- /dev/null +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/BucketLoadBalanceChannelComputer.java @@ -0,0 +1,167 @@ +/* + * 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.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 org.apache.fluss.utils.MathUtils; + +import javax.annotation.Nullable; + +import java.util.Arrays; +import java.util.List; + +/** + * {@link ChannelComputer} for {@link + * org.apache.fluss.flink.sink.shuffle.DistributionMode#BUCKET_LOAD_BALANCE}. + * + *

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. + * + *

The core routing algorithm ({@link #routeByBucket}) is shared between the sink (this class) + * and the source-side lookup-join partitioner ({@code FlussLookupInputPartitioner}). + * + *

Records with the same bucket key always route to the same channel. + * + * @param the type of records + */ +@Internal +public class BucketLoadBalanceChannelComputer implements ChannelComputer { + + private static final long serialVersionUID = 1L; + + private final @Nullable DataLakeFormat lakeFormat; + private final int numBucket; + private final RowType flussRowType; + private final List bucketKeys; + private final FlussSerializationSchema serializationSchema; + + private transient int numChannels; + private transient BucketingFunction bucketingFunction; + private transient KeyEncoder bucketKeyEncoder; + + public BucketLoadBalanceChannelComputer( + RowType flussRowType, + List bucketKeys, + @Nullable DataLakeFormat lakeFormat, + int numBucket, + FlussSerializationSchema 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.bucketingFunction = BucketingFunction.of(lakeFormat); + this.bucketKeyEncoder = KeyEncoder.ofBucketKeyEncoder(flussRowType, bucketKeys, lakeFormat); + + 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(); + byte[] bucketKeyBytes = bucketKeyEncoder.encodeKey(row); + return routeByBucket( + bucketKeyBytes, bucketKeyBytes, numBucket, numChannels, bucketingFunction); + } 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"; + } + + // ------------------------------------------------------------------------------------------ + // Shared routing logic + // ------------------------------------------------------------------------------------------ + + /** + * Routes a record to a downstream channel using bucket-ID + LCM-based slot assignment. + * + *

This is the core routing algorithm used by both the sink-side {@code + * BucketLoadBalanceChannelComputer} and the source-side lookup-join partitioner ({@code + * FlussLookupInputPartitioner}). Callers provide two separate byte arrays so that different + * hash keys can be used for the intra-bucket slot selection without affecting bucket-ID + * computation. + * + *

+ * + *

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 original {@code BUCKET} distribution mode. + * + * @param bucketKeyBytes encoded bucket key (determines bucket ID) + * @param hashKeyBytes encoded key used for the intra-bucket slot hash + * @param numBuckets total number of buckets + * @param numChannels number of downstream channels (subtask parallelism) + * @param bucketingFunction the bucketing function for computing the bucket ID + * @return channel index in {@code [0, numChannels)} + */ + public static int routeByBucket( + byte[] bucketKeyBytes, + byte[] hashKeyBytes, + int numBuckets, + int numChannels, + BucketingFunction bucketingFunction) { + 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 (numBuckets >= numChannels && numBuckets % numChannels == 0) { + return bucketId % numChannels; + } + + int recordHash = MathUtils.murmurHash(Arrays.hashCode(hashKeyBytes)); + return ChannelComputer.select(bucketId, numBuckets, numChannels, recordHash); + } +} diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/ChannelComputer.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/ChannelComputer.java index a52189dc694..65a911dc104 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/ChannelComputer.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/ChannelComputer.java @@ -62,4 +62,55 @@ static int select(String partitionName, int bucket, int numChannels) { static int select(int bucket, int numChannels) { return bucket % numChannels; } + + /** + * Selects a channel using LCM-based logical slot assignment, ensuring even load distribution + * across all channels regardless of whether the bucket count and channel count are evenly + * divisible. + * + *

Three cases are handled: + * + *

+ * + * @param bucketId the bucket ID computed by {@code BucketingFunction} + * @param numBuckets total number of buckets in the table + * @param numChannels number of downstream channels (subtask parallelism) + * @param recordHash hash of the record (e.g. murmur hash of the bucket key bytes) + * @return the target channel index in {@code [0, numChannels)} + */ + static int select(int bucketId, int numBuckets, int numChannels, int recordHash) { + if (numBuckets >= numChannels && numBuckets % numChannels == 0) { + return bucketId % numChannels; + } + + 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); + } + + static int greatestCommonDivisor(int first, int second) { + while (second != 0) { + int remainder = first % second; + first = second; + second = remainder; + } + return first; + } } diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/FlinkSink.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/FlinkSink.java index 32b1d6414cf..897fad96557 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/FlinkSink.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/FlinkSink.java @@ -173,6 +173,12 @@ public DataStream addPreWriteTopology(DataStream 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 +240,18 @@ private DataStream bucketShuffle(DataStream input) { flussSerializationSchema), input.getParallelism()); } + + private DataStream bucketLoadBalanceShuffle(DataStream input) { + return partition( + input, + new BucketLoadBalanceChannelComputer<>( + toFlussRowType(tableRowType), + bucketKeys, + lakeFormat, + numBucket, + flussSerializationSchema), + input.getParallelism()); + } } @Internal @@ -336,6 +354,18 @@ public DataStream addPreWriteTopology(DataStream 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( diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/shuffle/DistributionMode.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/shuffle/DistributionMode.java index 63883533bec..32838d8629d 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/shuffle/DistributionMode.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/shuffle/DistributionMode.java @@ -65,5 +65,24 @@ public enum DistributionMode { *

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. + * + *

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. + * + *

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, making this mode safe for + * Primary Key tables (merge semantics preserved). + * + *

Characteristics: + * + *

Requires 'bucket.key' to be defined. Suitable for tables where intra-bucket ordering is + * not required but even load distribution matters. + */ + BUCKET_LOAD_BALANCE } diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/lookup/FlussLookupInputPartitioner.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/lookup/FlussLookupInputPartitioner.java index 7aef52476a2..7ad47e40d92 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/lookup/FlussLookupInputPartitioner.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/lookup/FlussLookupInputPartitioner.java @@ -25,7 +25,6 @@ import org.apache.fluss.row.InternalRow; import org.apache.fluss.row.encode.CompactedKeyEncoder; import org.apache.fluss.row.encode.KeyEncoder; -import org.apache.fluss.utils.MathUtils; import org.apache.flink.table.data.RowData; import org.apache.flink.table.types.logical.RowType; @@ -33,7 +32,6 @@ import javax.annotation.Nullable; import java.util.ArrayList; -import java.util.Arrays; import java.util.List; import static org.apache.fluss.utils.Preconditions.checkArgument; @@ -44,11 +42,10 @@ * bucketing used by {@code PrimaryKeyLookuper}/{@code PrefixKeyLookuper} (bucket-key encoding + * {@link BucketingFunction}). * - *

Partitioned and non-partitioned tables use the same strategy. When the bucket and subtask - * counts do not evenly divide each other, weighted logical slots balance the expected load across - * subtasks. This keeps every lookup key on a stable subtask while bounding each bucket's RPC - * fan-out. Existing bucket-affinity mappings are preserved when either count is an exact multiple - * of the other. + *

The routing delegates to {@link + * org.apache.fluss.flink.sink.BucketLoadBalanceChannelComputer#routeByBucket} which uses LCM-based + * logical slot assignment to balance load evenly across all subtasks while keeping every lookup key + * on a stable subtask. */ public class FlussLookupInputPartitioner implements InputDataPartitionerAdapter { @@ -128,47 +125,17 @@ public int partition(RowData joinKeys, int numPartitions) { RowData normalizedKey = normalizer.normalizeLookupKey(joinKeys); InternalRow flussKeyRow = reuseRow.replace(normalizedKey); byte[] bucketKeyBytes = bucketKeyEncoder.encodeKey(flussKeyRow); - // BucketingFunction always returns a non-negative bucket id. - int bucketId = bucketingFunction.bucketing(bucketKeyBytes, numBuckets); - if (numBuckets >= numPartitions && numBuckets % numPartitions == 0) { - return bucketId % numPartitions; - } - byte[] lookupKeyBytes = lookupKeyEncoder.encodeKey(flussKeyRow); - // Do not derive this hash from the bucket hash. The low bits of that hash determine the - // bucket id, so reusing it can make some logical slots unreachable. - int lookupKeyHash = MathUtils.murmurHash(Arrays.hashCode(lookupKeyBytes)); - if (numPartitions > numBuckets && numPartitions % numBuckets == 0) { - // Preserve the original disjoint round-robin assignment when every bucket owns the - // same number of subtasks. - int candidateCount = numPartitions / numBuckets; - return (lookupKeyHash % candidateCount) * numBuckets + bucketId; - } - - // Represent the assignment with LCM(numBuckets, numPartitions) logical slots without - // materializing them. Each bucket owns numPartitions / gcd consecutive slots and each - // subtask owns numBuckets / gcd consecutive slots. A uniform hash within a bucket therefore - // gives every subtask the same expected number of slots, while a bucket can only fan out to - // the subtasks whose slot ranges overlap its own range. - int gcd = greatestCommonDivisor(numBuckets, numPartitions); - int slotsPerBucket = numPartitions / gcd; - int slotsPerSubtask = numBuckets / gcd; - int slotWithinBucket = lookupKeyHash % slotsPerBucket; - long logicalSlot = (long) bucketId * slotsPerBucket + slotWithinBucket; - return (int) (logicalSlot / slotsPerSubtask); + // Delegate to the shared bucket-load-balance routing algorithm. + // bucketKeyBytes determines the bucket; lookupKeyBytes provides a separate hash + // domain so that different lookup keys in the same bucket are spread across the + // bucket's subtask range while the same lookup key always lands on the same subtask. + return org.apache.fluss.flink.sink.BucketLoadBalanceChannelComputer.routeByBucket( + bucketKeyBytes, lookupKeyBytes, numBuckets, numPartitions, bucketingFunction); } @Override public boolean isDeterministic() { return true; } - - private static int greatestCommonDivisor(int first, int second) { - while (second != 0) { - int remainder = first % second; - first = second; - second = remainder; - } - return first; - } } diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/FlinkConnectorOptionsUtils.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/FlinkConnectorOptionsUtils.java index 48c8a739c6b..ec03c777a33 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/FlinkConnectorOptionsUtils.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/FlinkConnectorOptionsUtils.java @@ -91,12 +91,13 @@ public static void validateDistributionModeForMergeEngine( @Nullable MergeEngineType mergeEngineType, DistributionMode distributionMode) { if (mergeEngineType != null && distributionMode != DistributionMode.BUCKET + && distributionMode != DistributionMode.BUCKET_LOAD_BALANCE && distributionMode != DistributionMode.AUTO) { throw new IllegalArgumentException( String.format( "For primary key tables with merge engine ('%s'), " - + "'sink.distribution-mode' must be 'bucket' or 'auto' (default). " - + "Disabling shuffle breaks merge semantics because records with the same key " + + "'sink.distribution-mode' must be 'bucket', 'bucket_load_balance' or 'auto' (default). " + + "Disabling keyed shuffle breaks merge semantics because records with the same key " + "must be processed by the same task. Current mode: %s", mergeEngineType, distributionMode)); } diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/BucketLoadBalanceChannelComputerTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/BucketLoadBalanceChannelComputerTest.java new file mode 100644 index 00000000000..12942820de9 --- /dev/null +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/BucketLoadBalanceChannelComputerTest.java @@ -0,0 +1,186 @@ +/* + * 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.flink.sink.serializer.FlussSerializationSchema; +import org.apache.fluss.flink.sink.serializer.RowDataSerializationSchema; +import org.apache.fluss.flink.sink.serializer.SerializerInitContextImpl; + +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.StringData; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; + +import static org.apache.fluss.record.TestData.DATA1_ROW_TYPE; +import static org.assertj.core.api.Assertions.assertThat; + +/** Test for {@link BucketLoadBalanceChannelComputer}. */ +class BucketLoadBalanceChannelComputerTest { + + private static final FlussSerializationSchema serializationSchema = + new RowDataSerializationSchema(false, false); + + @BeforeAll + static void init() throws Exception { + serializationSchema.open(new SerializerInitContextImpl(DATA1_ROW_TYPE, false)); + } + + @Test + void testDeterministicRouting() { + int numBucket = 10; + BucketLoadBalanceChannelComputer channelComputer = + new BucketLoadBalanceChannelComputer<>( + DATA1_ROW_TYPE, + Collections.singletonList("a"), + null, + numBucket, + serializationSchema); + + // Same bucket key always routes to the same channel + for (int numChannel = 1; numChannel <= 10; numChannel++) { + channelComputer.setup(numChannel); + for (int i = 0; i < 100; i++) { + int expectedChannel = -1; + for (int retry = 0; retry < 5; retry++) { + GenericRowData row = GenericRowData.of(i, StringData.fromString("a1")); + int channel = channelComputer.channel(row); + if (expectedChannel < 0) { + expectedChannel = channel; + } else { + assertThat(channel) + .as( + "Same key should route to same channel for numChannels=%d, key=%d", + numChannel, i) + .isEqualTo(expectedChannel); + } + assertThat(channel).isLessThan(numChannel); + } + } + } + } + + @Test + void testAllChannelsReceiveData() { + int numBucket = 3; + int numChannels = 7; + BucketLoadBalanceChannelComputer channelComputer = + new BucketLoadBalanceChannelComputer<>( + DATA1_ROW_TYPE, + Collections.singletonList("a"), + null, + numBucket, + serializationSchema); + channelComputer.setup(numChannels); + + Set usedChannels = new HashSet<>(); + // Generate many different bucket key values to cover all channels + for (int i = 0; i < 500; i++) { + GenericRowData row = GenericRowData.of(i, StringData.fromString("val" + i)); + int channel = channelComputer.channel(row); + usedChannels.add(channel); + } + + // Unlike BUCKET mode where some channels may be idle when numBucket < numChannels, + // BUCKET_LOAD_BALANCE should use all channels + assertThat(usedChannels).hasSize(numChannels); + } + + @Test + void testEvenDistribution() { + int numBucket = 5; + int numChannels = 8; + BucketLoadBalanceChannelComputer channelComputer = + new BucketLoadBalanceChannelComputer<>( + DATA1_ROW_TYPE, + Collections.singletonList("a"), + null, + numBucket, + serializationSchema); + channelComputer.setup(numChannels); + + Map channelCounts = new HashMap<>(); + int totalRecords = 10000; + for (int i = 0; i < totalRecords; i++) { + GenericRowData row = GenericRowData.of(i, StringData.fromString("key" + i)); + int channel = channelComputer.channel(row); + channelCounts.merge(channel, 1, Integer::sum); + } + + // Each channel should have roughly totalRecords / numChannels records + double expectedPerChannel = (double) totalRecords / numChannels; + for (int c = 0; c < numChannels; c++) { + int count = channelCounts.getOrDefault(c, 0); + // Allow 20% deviation from expected + assertThat((double) count) + .as("Channel %d has %d records, expected ~%.0f", c, count, expectedPerChannel) + .isBetween(expectedPerChannel * 0.8, expectedPerChannel * 1.2); + } + } + + @Test + void testExactDivisibleCase() { + int numBucket = 6; + int numChannels = 3; + BucketLoadBalanceChannelComputer channelComputer = + new BucketLoadBalanceChannelComputer<>( + DATA1_ROW_TYPE, + Collections.singletonList("a"), + null, + numBucket, + serializationSchema); + channelComputer.setup(numChannels); + + // When numBucket % numChannels == 0, should behave like BUCKET mode + // (each channel gets numBucket / numChannels buckets) + for (int i = 0; i < 100; i++) { + GenericRowData row = GenericRowData.of(i, StringData.fromString("v" + i)); + int channel = channelComputer.channel(row); + assertThat(channel).isBetween(0, numChannels - 1); + } + } + + @Test + void testMoreChannelsThanBuckets() { + int numBucket = 3; + int numChannels = 6; + BucketLoadBalanceChannelComputer channelComputer = + new BucketLoadBalanceChannelComputer<>( + DATA1_ROW_TYPE, + Collections.singletonList("a"), + null, + numBucket, + serializationSchema); + channelComputer.setup(numChannels); + + Set usedChannels = new HashSet<>(); + for (int i = 0; i < 500; i++) { + GenericRowData row = GenericRowData.of(i, StringData.fromString("k" + i)); + int channel = channelComputer.channel(row); + usedChannels.add(channel); + } + + assertThat(usedChannels).hasSize(numChannels); + } +} diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlinkTableSinkITCase.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlinkTableSinkITCase.java index 09a59de57ff..a35e030cd3d 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlinkTableSinkITCase.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlinkTableSinkITCase.java @@ -247,6 +247,8 @@ void testAppendLogWithBucketKey(DistributionMode distributionMode) throws Except String insertPlan = tEnv.explainSql(insertSql, ExplainDetail.JSON_EXECUTION_PLAN); if (distributionMode == DistributionMode.BUCKET) { assertThat(insertPlan).contains("\"ship_strategy\" : \"BUCKET\""); + } else if (distributionMode == DistributionMode.BUCKET_LOAD_BALANCE) { + assertThat(insertPlan).contains("\"ship_strategy\" : \"BUCKET_LOAD_BALANCE\""); } else { assertThat(insertPlan).contains("\"ship_strategy\" : \"FORWARD\""); } @@ -390,10 +392,10 @@ void testAppendLogPartitionTable(DistributionMode distributionMode) throws Excep + "(11, 3511, 'stave'), " + "(12, 3512, 'Tim')"; - if (distributionMode == DistributionMode.BUCKET) { + if (distributionMode == DistributionMode.BUCKET + || distributionMode == DistributionMode.BUCKET_LOAD_BALANCE) { assertThatThrownBy(() -> tEnv.explainSql(insertSql, ExplainDetail.JSON_EXECUTION_PLAN)) - .hasMessageContaining( - "BUCKET mode is only supported for log tables with bucket keys"); + .hasMessageContaining("mode is only supported for log tables with bucket keys"); return; } @@ -457,6 +459,8 @@ void testPut(DistributionMode distributionMode) throws Exception { String insertPlan = tEnv.explainSql(insertSql, ExplainDetail.JSON_EXECUTION_PLAN); if (distributionMode == DistributionMode.BUCKET) { assertThat(insertPlan).contains("\"ship_strategy\" : \"BUCKET\""); + } else if (distributionMode == DistributionMode.BUCKET_LOAD_BALANCE) { + assertThat(insertPlan).contains("\"ship_strategy\" : \"BUCKET_LOAD_BALANCE\""); } else if (distributionMode == DistributionMode.AUTO || distributionMode == DistributionMode.NONE) { assertThat(insertPlan).contains("\"ship_strategy\" : \"FORWARD\""); diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlussSinkBuilderTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlussSinkBuilderTest.java index a6c848faa66..2bf0a18b2e8 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlussSinkBuilderTest.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlussSinkBuilderTest.java @@ -160,6 +160,10 @@ void testSinkShuffle() throws Exception { builder.setDistributionMode(DistributionMode.PARTITION_DYNAMIC); shuffleMode = getFieldValue(builder, "distributionMode"); assertThat(shuffleMode).isEqualTo(DistributionMode.PARTITION_DYNAMIC); + + builder.setDistributionMode(DistributionMode.BUCKET_LOAD_BALANCE); + shuffleMode = getFieldValue(builder, "distributionMode"); + assertThat(shuffleMode).isEqualTo(DistributionMode.BUCKET_LOAD_BALANCE); } @Test diff --git a/website/docs/engine-flink/options.md b/website/docs/engine-flink/options.md index 65e8fc3ae81..66b3151fa7f 100644 --- a/website/docs/engine-flink/options.md +++ b/website/docs/engine-flink/options.md @@ -148,7 +148,7 @@ See more details about [ALTER TABLE ... SET](engine-flink/ddl.md#set-properties) |-----------------------------------------------------|------------|-------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| | sink.ignore-delete | Boolean | false | If set to true, the sink will ignore DELETE and UPDATE_BEFORE changelog events. | | sink.bucket-shuffle | Boolean | true | Whether to shuffle by bucket id before write to sink. Shuffling the data with the same bucket id to be processed by the same task can improve the efficiency of client processing and reduce resource consumption. For Log Table, bucket shuffle will only take effect when the 'bucket.key' is defined. For Primary Key table, it is enabled by default. This option is deprecated. Please use `sink.distribution-mode` instead, which provides more flexible distribution strategies.| -| sink.distribution-mode | Enum | AUTO | Defines the distribution mode for shuffling data to the sink. Available options are `AUTO`, `NONE`, `BUCKET`, and `PARTITION_DYNAMIC`. See [Distribution Modes](#distribution-modes) for details about each option. | +| sink.distribution-mode | Enum | AUTO | Defines the distribution mode for shuffling data to the sink. Available options are `AUTO`, `NONE`, `BUCKET`, `BUCKET_LOAD_BALANCE`, and `PARTITION_DYNAMIC`. See [Distribution Modes](#distribution-modes) for details about each option. | | client.writer.buffer.memory-size | MemorySize | 64mb | The total bytes of memory the writer can use to buffer internal rows. | | client.writer.buffer.page-size | MemorySize | 128kb | Size of every page in memory buffers (`client.writer.buffer.memory-size`). | | client.writer.buffer.per-request-memory-size | MemorySize | 16mb | The minimum number of bytes that will be allocated by the writer rounded down to the closest multiple of client.writer.buffer.page-size. It must be greater than or equal to client.writer.buffer.page-size. This option allows to allocate memory in batches to have better CPU-cached friendliness due to contiguous segments. | @@ -186,6 +186,14 @@ Shuffle data by bucket ID before writing to sink. This groups data with the same - For Log Tables, bucket shuffle only takes effect when the `bucket.key` is defined - **Note:** When sink parallelism exceeds the number of buckets, some sink tasks may remain idle without receiving data +#### BUCKET_LOAD_BALANCE +Shuffle data by bucket key before writing to sink. Unlike **BUCKET**, this mode uses an LCM-based logical slot assignment to ensure every subtask receives data, even when the bucket count and sink parallelism do not evenly divide each other. + +**Characteristics:** +- Records with the same bucket key always route to the same subtask, so it is safe for Primary Key tables and merge-engine tables +- For Log Tables, it only takes effect when the `bucket.key` is defined +- Intra-bucket ordering is not guaranteed; use this mode when load balancing matters more than per-bucket ordering + #### PARTITION_DYNAMIC Dynamically adjusts shuffle strategy based on partition key traffic patterns. This mode monitors data distribution and adjusts the shuffle behavior to balance the load. From 949c17c4484b5df7b86d237059bf36e108959276 Mon Sep 17 00:00:00 2001 From: Hongshun Wang Date: Wed, 23 Sep 2026 14:51:24 +0800 Subject: [PATCH 2/2] modified based on CR --- .../fluss/flink/FlinkConnectorOptions.java | 2 +- .../shuffle/BucketLoadBalanceRouter.java | 134 ++++++++++++++++++ .../BucketLoadBalanceChannelComputer.java | 73 ++-------- .../fluss/flink/sink/ChannelComputer.java | 51 ------- .../apache/fluss/flink/sink/FlinkSink.java | 6 + .../flink/sink/shuffle/DistributionMode.java | 7 +- .../lookup/FlussLookupInputPartitioner.java | 41 ++---- .../utils/FlinkConnectorOptionsUtils.java | 42 ++++++ .../flink/catalog/FlinkTableFactoryTest.java | 56 ++++++++ .../flink/sink/FlinkTableSinkITCase.java | 19 +++ .../flink/sink/FlussSinkBuilderTest.java | 37 ++++- website/docs/engine-flink/options.md | 2 +- 12 files changed, 325 insertions(+), 145 deletions(-) create mode 100644 fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/shuffle/BucketLoadBalanceRouter.java diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/FlinkConnectorOptions.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/FlinkConnectorOptions.java index b22715d10a2..fd4fd4d8022 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/FlinkConnectorOptions.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/FlinkConnectorOptions.java @@ -232,7 +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.\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" diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/shuffle/BucketLoadBalanceRouter.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/shuffle/BucketLoadBalanceRouter.java new file mode 100644 index 00000000000..b11389ec232 --- /dev/null +++ b/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. + * + *

Two callers share this router and must stay consistent: + * + *

    + *
  • 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. + *
  • 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. + *
+ * + *

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. + * + *

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; + } +} diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/BucketLoadBalanceChannelComputer.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/BucketLoadBalanceChannelComputer.java index edd70dfe1bd..dc9386aab4c 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/BucketLoadBalanceChannelComputer.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/BucketLoadBalanceChannelComputer.java @@ -21,17 +21,16 @@ 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 org.apache.fluss.utils.MathUtils; import javax.annotation.Nullable; -import java.util.Arrays; import java.util.List; /** @@ -43,8 +42,8 @@ * computer uses an LCM-based logical slot assignment so that every channel receives traffic from * every bucket. * - *

The core routing algorithm ({@link #routeByBucket}) is shared between the sink (this class) - * and the source-side lookup-join partitioner ({@code FlussLookupInputPartitioner}). + *

The routing algorithm is shared with the source-side lookup-join partitioner ({@code + * FlussLookupInputPartitioner}); see {@link BucketLoadBalanceRouter}. * *

Records with the same bucket key always route to the same channel. * @@ -62,8 +61,7 @@ public class BucketLoadBalanceChannelComputer implements ChannelComputer private final FlussSerializationSchema serializationSchema; private transient int numChannels; - private transient BucketingFunction bucketingFunction; - private transient KeyEncoder bucketKeyEncoder; + private transient BucketLoadBalanceRouter bucketLoadBalanceRouter; public BucketLoadBalanceChannelComputer( RowType flussRowType, @@ -81,8 +79,12 @@ public BucketLoadBalanceChannelComputer( @Override public void setup(int numChannels) { this.numChannels = numChannels; - this.bucketingFunction = BucketingFunction.of(lakeFormat); - this.bucketKeyEncoder = KeyEncoder.ofBucketKeyEncoder(flussRowType, bucketKeys, lakeFormat); + this.bucketLoadBalanceRouter = + new BucketLoadBalanceRouter( + KeyEncoder.ofBucketKeyEncoder(flussRowType, bucketKeys, lakeFormat), + null, + BucketingFunction.of(lakeFormat), + numBucket); try { this.serializationSchema.open(new SerializerInitContextImpl(flussRowType, false)); @@ -96,9 +98,7 @@ public int channel(InputT record) { try { RowWithOp rowWithOp = serializationSchema.serialize(record); InternalRow row = rowWithOp.getRow(); - byte[] bucketKeyBytes = bucketKeyEncoder.encodeKey(row); - return routeByBucket( - bucketKeyBytes, bucketKeyBytes, numBucket, numChannels, bucketingFunction); + return bucketLoadBalanceRouter.route(row, numChannels); } catch (Exception e) { throw new FlussRuntimeException( String.format( @@ -113,55 +113,4 @@ record != null ? record.getClass().getName() : "null", e.getMessage()), public String toString() { return "BUCKET_LOAD_BALANCE"; } - - // ------------------------------------------------------------------------------------------ - // Shared routing logic - // ------------------------------------------------------------------------------------------ - - /** - * Routes a record to a downstream channel using bucket-ID + LCM-based slot assignment. - * - *

This is the core routing algorithm used by both the sink-side {@code - * BucketLoadBalanceChannelComputer} and the source-side lookup-join partitioner ({@code - * FlussLookupInputPartitioner}). Callers provide two separate byte arrays so that different - * hash keys can be used for the intra-bucket slot selection without affecting bucket-ID - * computation. - * - *

    - *
  • Sink usage: both {@code bucketKeyBytes} and {@code hashKeyBytes} are the encoded bucket - * key — records with the same bucket key always land on the same subtask. - *
  • Lookup-join usage: {@code bucketKeyBytes} comes from the bucket-key encoder while - * {@code hashKeyBytes} comes from the full lookup-key encoder, giving a different hash - * domain within each bucket so that different lookup keys in the same bucket are spread - * across the bucket's subtask range. - *
- * - *

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 original {@code BUCKET} distribution mode. - * - * @param bucketKeyBytes encoded bucket key (determines bucket ID) - * @param hashKeyBytes encoded key used for the intra-bucket slot hash - * @param numBuckets total number of buckets - * @param numChannels number of downstream channels (subtask parallelism) - * @param bucketingFunction the bucketing function for computing the bucket ID - * @return channel index in {@code [0, numChannels)} - */ - public static int routeByBucket( - byte[] bucketKeyBytes, - byte[] hashKeyBytes, - int numBuckets, - int numChannels, - BucketingFunction bucketingFunction) { - 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 (numBuckets >= numChannels && numBuckets % numChannels == 0) { - return bucketId % numChannels; - } - - int recordHash = MathUtils.murmurHash(Arrays.hashCode(hashKeyBytes)); - return ChannelComputer.select(bucketId, numBuckets, numChannels, recordHash); - } } diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/ChannelComputer.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/ChannelComputer.java index 65a911dc104..a52189dc694 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/ChannelComputer.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/ChannelComputer.java @@ -62,55 +62,4 @@ static int select(String partitionName, int bucket, int numChannels) { static int select(int bucket, int numChannels) { return bucket % numChannels; } - - /** - * Selects a channel using LCM-based logical slot assignment, ensuring even load distribution - * across all channels regardless of whether the bucket count and channel count are evenly - * divisible. - * - *

Three cases are handled: - * - *

    - *
  • {@code numBuckets >= numChannels && numBuckets % numChannels == 0}: simple {@code - * bucketId % numChannels}. - *
  • {@code numChannels > numBuckets && numChannels % numBuckets == 0}: disjoint - * round-robin, preserving the original assignment shape. - *
  • Otherwise: LCM-based logical slots — each bucket owns {@code numChannels / gcd} - * consecutive slots and each channel owns {@code numBuckets / gcd} consecutive slots. A - * hash within the bucket distributes records evenly across channels whose slot ranges - * overlap the bucket's range. - *
- * - * @param bucketId the bucket ID computed by {@code BucketingFunction} - * @param numBuckets total number of buckets in the table - * @param numChannels number of downstream channels (subtask parallelism) - * @param recordHash hash of the record (e.g. murmur hash of the bucket key bytes) - * @return the target channel index in {@code [0, numChannels)} - */ - static int select(int bucketId, int numBuckets, int numChannels, int recordHash) { - if (numBuckets >= numChannels && numBuckets % numChannels == 0) { - return bucketId % numChannels; - } - - 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); - } - - static int greatestCommonDivisor(int first, int second) { - while (second != 0) { - int remainder = first % second; - first = second; - second = remainder; - } - return first; - } } diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/FlinkSink.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/FlinkSink.java index 897fad96557..c7b9060716e 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/FlinkSink.java +++ b/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; @@ -335,6 +336,11 @@ public UpsertSinkWriter createWriter( @Override public DataStream addPreWriteTopology(DataStream 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 stream; switch (distributionMode) { case NONE: diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/shuffle/DistributionMode.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/shuffle/DistributionMode.java index 32838d8629d..d2fb7752ec6 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/shuffle/DistributionMode.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/shuffle/DistributionMode.java @@ -76,8 +76,11 @@ public enum DistributionMode { * *

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, making this mode safe for - * Primary Key tables (merge semantics preserved). + * 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. * *

Characteristics: * diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/lookup/FlussLookupInputPartitioner.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/lookup/FlussLookupInputPartitioner.java index 7ad47e40d92..a90d2f4ef6a 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/lookup/FlussLookupInputPartitioner.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/lookup/FlussLookupInputPartitioner.java @@ -20,6 +20,7 @@ import org.apache.fluss.bucketing.BucketingFunction; import org.apache.fluss.flink.adapter.SupportsLookupCustomShuffleAdapter.InputDataPartitionerAdapter; import org.apache.fluss.flink.row.FlinkAsFlussRow; +import org.apache.fluss.flink.shuffle.BucketLoadBalanceRouter; import org.apache.fluss.flink.utils.FlinkConversions; import org.apache.fluss.metadata.DataLakeFormat; import org.apache.fluss.row.InternalRow; @@ -42,10 +43,9 @@ * bucketing used by {@code PrimaryKeyLookuper}/{@code PrefixKeyLookuper} (bucket-key encoding + * {@link BucketingFunction}). * - *

The routing delegates to {@link - * org.apache.fluss.flink.sink.BucketLoadBalanceChannelComputer#routeByBucket} which uses LCM-based - * logical slot assignment to balance load evenly across all subtasks while keeping every lookup key - * on a stable subtask. + *

The routing delegates to {@link BucketLoadBalanceRouter} which uses LCM-based logical slot + * assignment to balance load evenly across all subtasks while keeping every lookup key on a stable + * subtask. */ public class FlussLookupInputPartitioner implements InputDataPartitionerAdapter { @@ -60,9 +60,7 @@ public class FlussLookupInputPartitioner implements InputDataPartitionerAdapter @Nullable private final DataLakeFormat lakeFormat; private final int numBuckets; - private transient KeyEncoder bucketKeyEncoder; - private transient KeyEncoder lookupKeyEncoder; - private transient BucketingFunction bucketingFunction; + private transient BucketLoadBalanceRouter bucketLoadBalanceRouter; private transient FlinkAsFlussRow reuseRow; /** @@ -96,16 +94,17 @@ public FlussLookupInputPartitioner( } private void ensureInitialized() { - if (bucketKeyEncoder == null) { + if (bucketLoadBalanceRouter == null) { org.apache.fluss.types.RowType flussKeyType = FlinkConversions.toFlussRowType(keyFlinkRowType); // bucketing uses the bucket-key encoder consistent with the client's bucket routing - bucketKeyEncoder = - KeyEncoder.ofBucketKeyEncoder(flussKeyType, bucketKeyNames, lakeFormat); - lookupKeyEncoder = - CompactedKeyEncoder.createKeyEncoder( - flussKeyType, keyFlinkRowType.getFieldNames()); - bucketingFunction = BucketingFunction.of(lakeFormat); + bucketLoadBalanceRouter = + new BucketLoadBalanceRouter( + KeyEncoder.ofBucketKeyEncoder(flussKeyType, bucketKeyNames, lakeFormat), + CompactedKeyEncoder.createKeyEncoder( + flussKeyType, keyFlinkRowType.getFieldNames()), + BucketingFunction.of(lakeFormat), + numBuckets); reuseRow = new FlinkAsFlussRow(); } } @@ -124,18 +123,6 @@ public int partition(RowData joinKeys, int numPartitions) { // normalize the projected join keys into the Fluss key order RowData normalizedKey = normalizer.normalizeLookupKey(joinKeys); InternalRow flussKeyRow = reuseRow.replace(normalizedKey); - byte[] bucketKeyBytes = bucketKeyEncoder.encodeKey(flussKeyRow); - byte[] lookupKeyBytes = lookupKeyEncoder.encodeKey(flussKeyRow); - // Delegate to the shared bucket-load-balance routing algorithm. - // bucketKeyBytes determines the bucket; lookupKeyBytes provides a separate hash - // domain so that different lookup keys in the same bucket are spread across the - // bucket's subtask range while the same lookup key always lands on the same subtask. - return org.apache.fluss.flink.sink.BucketLoadBalanceChannelComputer.routeByBucket( - bucketKeyBytes, lookupKeyBytes, numBuckets, numPartitions, bucketingFunction); - } - - @Override - public boolean isDeterministic() { - return true; + return bucketLoadBalanceRouter.route(flussKeyRow, numPartitions); } } diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/FlinkConnectorOptionsUtils.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/FlinkConnectorOptionsUtils.java index ec03c777a33..520eeb03875 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/FlinkConnectorOptionsUtils.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/utils/FlinkConnectorOptionsUtils.java @@ -83,6 +83,10 @@ public static void validateTableSourceOptions(ReadableConfig tableOptions) { *

  • FIRST_ROW/VERSIONED: merge semantics require same-key records go to same task * * + *

    Additionally, the AGGREGATION merge engine enables Undo Recovery on the sink, whose + * one-writer-per-bucket assumption further excludes BUCKET_LOAD_BALANCE; see {@link + * #validateDistributionModeForUndoRecovery}. + * * @param mergeEngineType the merge engine type (can be null for non-merge-engine tables) * @param distributionMode the distribution mode configured for the sink * @throws IllegalArgumentException if distribution mode is incompatible with merge engine @@ -101,6 +105,44 @@ public static void validateDistributionModeForMergeEngine( + "must be processed by the same task. Current mode: %s", mergeEngineType, distributionMode)); } + // The AGGREGATION merge engine enables Undo Recovery on the sink; the + // one-writer-per-bucket constraint of Undo Recovery further excludes BUCKET_LOAD_BALANCE. + validateDistributionModeForUndoRecovery( + mergeEngineType == MergeEngineType.AGGREGATION, distributionMode); + } + + /** + * Validates that the distribution mode is safe for the sink's Undo Recovery. + * + *

    Undo Recovery, which is enabled automatically for aggregation tables, assumes each bucket + * is written by exactly one sink subtask: the writer state tracks the last written offset per + * bucket and fails on conflicting offsets during recovery. BUCKET_LOAD_BALANCE and NONE may fan + * one bucket out to several subtasks and break that assumption. + * + *

    This is the single guard shared by the entry-point validation ({@link + * #validateDistributionModeForMergeEngine}) and the sink builder's defense-in-depth check, so + * the two can never drift apart. + * + * @param enableUndoRecovery whether Undo Recovery is enabled for the sink + * @param distributionMode the distribution mode configured for the sink + * @throws IllegalArgumentException if the distribution mode may fan one bucket out to several + * subtasks while Undo Recovery is enabled + */ + public static void validateDistributionModeForUndoRecovery( + boolean enableUndoRecovery, DistributionMode distributionMode) { + if (enableUndoRecovery + && distributionMode != DistributionMode.AUTO + && distributionMode != DistributionMode.BUCKET) { + throw new IllegalArgumentException( + String.format( + "'sink.distribution-mode' = '%s' is not supported when Undo Recovery " + + "is enabled (automatic for the aggregation merge engine): " + + "Undo Recovery assumes each bucket is written by exactly one " + + "sink subtask, but this distribution mode may fan one bucket " + + "out to several subtasks. Please use 'bucket' or 'auto' " + + "(default) instead.", + distributionMode)); + } } public static StartupOptions getStartupOptions(ReadableConfig tableOptions, ZoneId timeZone) { diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/catalog/FlinkTableFactoryTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/catalog/FlinkTableFactoryTest.java index 37575097c24..54d3fa49273 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/catalog/FlinkTableFactoryTest.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/catalog/FlinkTableFactoryTest.java @@ -21,6 +21,7 @@ import org.apache.fluss.flink.FlinkConnectorOptions; import org.apache.fluss.flink.adapter.CatalogTableAdapter; import org.apache.fluss.flink.sink.FlinkTableSink; +import org.apache.fluss.flink.sink.shuffle.DistributionMode; import org.apache.fluss.flink.source.BinlogFlinkTableSource; import org.apache.fluss.flink.source.ChangelogFlinkTableSource; import org.apache.fluss.flink.source.FlinkTableSource; @@ -321,6 +322,61 @@ void testSink() { createTableSink(schema, properties); } + @Test + void testSinkDistributionModeForMergeEngine() { + ResolvedSchema schema = createBasicSchema(); + + // keyed modes are accepted for every merge engine + for (String mergeEngine : new String[] {"first_row", "versioned", "aggregation"}) { + for (String mode : new String[] {"auto", "bucket"}) { + Map properties = getBasicOptionsWithBucketKey(); + properties.put(ConfigOptions.TABLE_MERGE_ENGINE.key(), mergeEngine); + properties.put(FlinkConnectorOptions.SINK_DISTRIBUTION_MODE.key(), mode); + createTableSink(schema, properties); + } + } + + // unkeyed modes are rejected for every merge engine + for (String mergeEngine : new String[] {"first_row", "versioned", "aggregation"}) { + for (String mode : new String[] {"none", "partition_dynamic"}) { + Map properties = getBasicOptionsWithBucketKey(); + properties.put(ConfigOptions.TABLE_MERGE_ENGINE.key(), mergeEngine); + properties.put(FlinkConnectorOptions.SINK_DISTRIBUTION_MODE.key(), mode); + assertThatThrownBy(() -> createTableSink(schema, properties)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining( + "'sink.distribution-mode' must be 'bucket', 'bucket_load_balance' or 'auto'"); + } + } + + // bucket_load_balance is accepted for merge engines without Undo Recovery ... + for (String mergeEngine : new String[] {"first_row", "versioned"}) { + Map properties = getBasicOptionsWithBucketKey(); + properties.put(ConfigOptions.TABLE_MERGE_ENGINE.key(), mergeEngine); + properties.put( + FlinkConnectorOptions.SINK_DISTRIBUTION_MODE.key(), "bucket_load_balance"); + createTableSink(schema, properties); + } + + // ... but rejected for aggregation, whose Undo Recovery assumes one writer per bucket + Map aggregationProperties = getBasicOptionsWithBucketKey(); + aggregationProperties.put(ConfigOptions.TABLE_MERGE_ENGINE.key(), "aggregation"); + aggregationProperties.put( + FlinkConnectorOptions.SINK_DISTRIBUTION_MODE.key(), "bucket_load_balance"); + assertThatThrownBy(() -> createTableSink(schema, aggregationProperties)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("is not supported when Undo Recovery is enabled") + .hasMessageContaining("BUCKET_LOAD_BALANCE") + .hasMessageContaining("Please use 'bucket' or 'auto'"); + + // without a merge engine, all modes stay accepted (no Undo Recovery is involved) + for (DistributionMode mode : DistributionMode.values()) { + Map properties = getBasicOptionsWithBucketKey(); + properties.put(FlinkConnectorOptions.SINK_DISTRIBUTION_MODE.key(), mode.name()); + createTableSink(schema, properties); + } + } + private ResolvedSchema createBasicSchema() { return new ResolvedSchema( Arrays.asList( diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlinkTableSinkITCase.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlinkTableSinkITCase.java index a35e030cd3d..4ac752a980a 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlinkTableSinkITCase.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlinkTableSinkITCase.java @@ -495,6 +495,25 @@ void testPut(DistributionMode distributionMode) throws Exception { assertResultsIgnoreOrder(rowIter, expectedRows, true); } + @Test + @MultiVersionTest + void testAggregationMergeEngineRejectsBucketLoadBalance() { + // Undo Recovery for aggregation tables assumes one writer per bucket, while + // bucket_load_balance may fan one bucket out to several subtasks, so the combination + // must be rejected when the sink is planned. + tEnv.executeSql( + "create table agg_load_balance_sink (a int not null primary key not enforced, " + + "b int) " + + "with('table.merge-engine' = 'aggregation', " + + "'fields.b.agg' = 'sum', " + + "'sink.distribution-mode' = 'bucket_load_balance')"); + + String insertSql = "INSERT INTO agg_load_balance_sink VALUES (1, 1)"; + assertThatThrownBy(() -> tEnv.explainSql(insertSql, ExplainDetail.JSON_EXECUTION_PLAN)) + .hasStackTraceContaining("is not supported when Undo Recovery is enabled") + .hasStackTraceContaining("Please use 'bucket' or 'auto'"); + } + @Test void testPutDuringAddColumn() throws Exception { tEnv.executeSql( diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlussSinkBuilderTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlussSinkBuilderTest.java index 2bf0a18b2e8..d095d117175 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlussSinkBuilderTest.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlussSinkBuilderTest.java @@ -210,7 +210,7 @@ void testUndoRecoverySinkCannotBeAddedToMultipleTopologies() { Collections.emptyList(), Collections.emptyList(), null, - DistributionMode.NONE, + DistributionMode.BUCKET, new OrderSerializationSchema(), true, null); @@ -222,6 +222,41 @@ void testUndoRecoverySinkCannotBeAddedToMultipleTopologies() { .hasMessageContaining("multiple topologies"); } + @Test + void testUndoRecoveryRejectsUnsafeDistributionModes() { + StreamExecutionEnvironment environment = + StreamExecutionEnvironment.getExecutionEnvironment(); + DataStream input = environment.fromElements(new Order()); + // Undo Recovery assumes each bucket is written by exactly one subtask; modes that may fan + // one bucket out to several subtasks must be rejected even when the builder is + // constructed directly, bypassing the entry-point validation. + for (DistributionMode mode : + new DistributionMode[] { + DistributionMode.NONE, DistributionMode.BUCKET_LOAD_BALANCE + }) { + FlinkSink.UpsertSinkWriterBuilder writerBuilder = + new FlinkSink.UpsertSinkWriterBuilder<>( + TablePath.of(databaseName, tableName), + new Configuration(), + RowType.of(new IntType()), + null, + 1, + Collections.emptyList(), + Collections.emptyList(), + null, + mode, + new OrderSerializationSchema(), + true, + null); + + assertThatThrownBy(() -> writerBuilder.addPreWriteTopology(input)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("is not supported when Undo Recovery is enabled") + .hasMessageContaining("'" + mode.name() + "'") + .hasMessageContaining("Please use 'bucket' or 'auto'"); + } + } + @Test void testComputeTargetColumnIndexesFullUpdate() { int[] result = diff --git a/website/docs/engine-flink/options.md b/website/docs/engine-flink/options.md index 66b3151fa7f..5a8bbc92053 100644 --- a/website/docs/engine-flink/options.md +++ b/website/docs/engine-flink/options.md @@ -190,7 +190,7 @@ Shuffle data by bucket ID before writing to sink. This groups data with the same Shuffle data by bucket key before writing to sink. Unlike **BUCKET**, this mode uses an LCM-based logical slot assignment to ensure every subtask receives data, even when the bucket count and sink parallelism do not evenly divide each other. **Characteristics:** -- Records with the same bucket key always route to the same subtask, so it is safe for Primary Key tables and merge-engine tables +- Records with the same bucket key always route to the same subtask, so merge semantics are preserved at runtime; however, this mode is **not supported for tables with the `aggregation` merge engine**, because Undo Recovery requires each bucket to be written by exactly one subtask, while bucket_load_balance may fan one bucket out to several subtasks - For Log Tables, it only takes effect when the `bucket.key` is defined - Intra-bucket ordering is not guaranteed; use this mode when load balancing matters more than per-bucket ordering