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..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,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" 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: + * + *

+ * + *

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 new file mode 100644 index 00000000000..dc9386aab4c --- /dev/null +++ b/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}. + * + *

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 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. + * + * @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 BucketLoadBalanceRouter bucketLoadBalanceRouter; + + 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.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"; + } +} 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..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; @@ -173,6 +174,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 +241,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 @@ -317,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: @@ -336,6 +360,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..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 @@ -65,5 +65,27 @@ 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, 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: + * + *

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..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,12 +20,12 @@ 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; 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 +33,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 +43,9 @@ * 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 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 { @@ -63,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; /** @@ -99,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(); } } @@ -127,48 +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); - // 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); - } - - @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; + 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 48c8a739c6b..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 @@ -91,15 +95,54 @@ 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)); } + // 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/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..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 @@ -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\""); @@ -491,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 a6c848faa66..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 @@ -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 @@ -206,7 +210,7 @@ void testUndoRecoverySinkCannotBeAddedToMultipleTopologies() { Collections.emptyList(), Collections.emptyList(), null, - DistributionMode.NONE, + DistributionMode.BUCKET, new OrderSerializationSchema(), true, null); @@ -218,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 65e8fc3ae81..5a8bbc92053 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 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 + #### 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.