From c841bf7662ebe397e8f2052d619ec60c9240e35a Mon Sep 17 00:00:00 2001 From: duankaixuan <1417048384@qq.com> Date: Sat, 12 Sep 2026 20:04:07 +0800 Subject: [PATCH 1/2] [flink] Resolve per-partition bucket count in sink bucket shuffle --- .../sink/FlinkRowDataChannelComputer.java | 44 ++- .../apache/fluss/flink/sink/FlinkSink.java | 4 + .../sink/PartitionBucketCountResolver.java | 255 +++++++++++++ .../sink/FlinkRowDataChannelComputerTest.java | 10 +- .../PartitionBucketCountResolverITCase.java | 352 ++++++++++++++++++ .../PartitionBucketCountResolverTest.java | 193 ++++++++++ .../fluss/flink/sink/UndoRecoveryITCase.java | 263 +++++++++++++ 7 files changed, 1098 insertions(+), 23 deletions(-) create mode 100644 fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/PartitionBucketCountResolver.java create mode 100644 fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/PartitionBucketCountResolverITCase.java create mode 100644 fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/PartitionBucketCountResolverTest.java diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/FlinkRowDataChannelComputer.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/FlinkRowDataChannelComputer.java index 23a3e0cd043..a1aad2d5471 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/FlinkRowDataChannelComputer.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/FlinkRowDataChannelComputer.java @@ -17,14 +17,15 @@ package org.apache.fluss.flink.sink; -import org.apache.fluss.annotation.VisibleForTesting; import org.apache.fluss.bucketing.BucketingFunction; import org.apache.fluss.client.table.getter.PartitionGetter; +import org.apache.fluss.config.Configuration; 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.metadata.TablePath; import org.apache.fluss.row.InternalRow; import org.apache.fluss.row.encode.KeyEncoder; import org.apache.fluss.types.RowType; @@ -48,11 +49,11 @@ public class FlinkRowDataChannelComputer implements ChannelComputer bucketKeys; private final List partitionKeys; private final FlussSerializationSchema serializationSchema; + private final @Nullable PartitionBucketCountResolver partitionBucketCountResolver; private transient int numChannels; private transient BucketingFunction bucketingFunction; private transient KeyEncoder bucketKeyEncoder; - private transient boolean combineShuffleWithPartitionName; private transient @Nullable PartitionGetter partitionGetter; public FlinkRowDataChannelComputer( @@ -60,6 +61,8 @@ public FlinkRowDataChannelComputer( List bucketKeys, List partitionKeys, @Nullable DataLakeFormat lakeFormat, + TablePath tablePath, + Configuration flussConfig, int numBucket, FlussSerializationSchema serializationSchema) { this.flussRowType = flussRowType; @@ -67,6 +70,12 @@ public FlinkRowDataChannelComputer( this.partitionKeys = partitionKeys; this.lakeFormat = lakeFormat; this.numBucket = numBucket; + // Non-partitioned tables fix the bucket count to the table-level value; partitioned + // tables resolve the actual per-partition count at runtime. + this.partitionBucketCountResolver = + partitionKeys.isEmpty() + ? null + : new PartitionBucketCountResolver(tablePath, flussConfig, numBucket); this.serializationSchema = serializationSchema; } @@ -81,11 +90,6 @@ public void setup(int numChannels) { this.partitionGetter = new PartitionGetter(flussRowType, partitionKeys); } - // Use shared logic from ChannelComputer to determine sharding strategy - this.combineShuffleWithPartitionName = - ChannelComputer.shouldCombinePartitionInSharding( - partitionGetter != null, numBucket, numChannels); - try { // no need to read real database, thus assume to deserialize the fluss row as same as // flink table type. @@ -101,14 +105,25 @@ public int channel(InputT record) { RowWithOp rowWithOp = serializationSchema.serialize(record); InternalRow row = rowWithOp.getRow(); - int bucketId = bucketingFunction.bucketing(bucketKeyEncoder.encodeKey(row), numBucket); - if (!combineShuffleWithPartitionName) { + if (partitionBucketCountResolver == null) { + // Non-partitioned table: the bucket count is fixed to the table-level value. + int bucketId = + bucketingFunction.bucketing(bucketKeyEncoder.encodeKey(row), numBucket); return ChannelComputer.select(bucketId, numChannels); - } else { - checkNotNull(partitionGetter, "partitionGetter is null"); - String partitionName = partitionGetter.getPartition(row); + } + + checkNotNull(partitionGetter, "partitionGetter is null"); + String partitionName = partitionGetter.getPartition(row); + // Resolve the partition's actual bucket count Sharding with the stale table-level + // value would scatter the records of one bucket across multiple writer subtasks + // and eventually break sink recovery. + int bucketCount = partitionBucketCountResolver.bucketCountOf(partitionName); + int bucketId = + bucketingFunction.bucketing(bucketKeyEncoder.encodeKey(row), bucketCount); + if (ChannelComputer.shouldCombinePartitionInSharding(true, bucketCount, numChannels)) { return ChannelComputer.select(partitionName, bucketId, numChannels); } + return ChannelComputer.select(bucketId, numChannels); } catch (Exception e) { throw new FlussRuntimeException( String.format( @@ -122,9 +137,4 @@ record != null ? record.getClass().getName() : "null", e.getMessage()), public String toString() { return "BUCKET"; } - - @VisibleForTesting - boolean isCombineShuffleWithPartitionName() { - return combineShuffleWithPartitionName; - } } 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..dd89cd7d78d 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 @@ -230,6 +230,8 @@ private DataStream bucketShuffle(DataStream input) { bucketKeys, partitionKeys, lakeFormat, + tablePath, + flussConfig, numBucket, flussSerializationSchema), input.getParallelism()); @@ -332,6 +334,8 @@ public DataStream addPreWriteTopology(DataStream input) { bucketKeys, partitionKeys, lakeFormat, + tablePath, + flussConfig, numBucket, flussSerializationSchema), input.getParallelism()); diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/PartitionBucketCountResolver.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/PartitionBucketCountResolver.java new file mode 100644 index 00000000000..75cfdc73803 --- /dev/null +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/PartitionBucketCountResolver.java @@ -0,0 +1,255 @@ +/* + * 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.client.Connection; +import org.apache.fluss.client.ConnectionFactory; +import org.apache.fluss.client.admin.Admin; +import org.apache.fluss.config.Configuration; +import org.apache.fluss.exception.FlussRuntimeException; +import org.apache.fluss.metadata.PartitionInfo; +import org.apache.fluss.metadata.TablePath; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import java.io.Serializable; +import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Resolves the actual bucket count of a partition at runtime for the pre-write bucket shuffle. + * + *

After a bucket num rescale, partitions of the same table may have different bucket counts: + * existing partitions keep their counts while partitions created afterwards use the new table-level + * default. A sharding decision based on the table-level value captured at job submission time + * therefore scatters the records of one bucket across multiple writer subtasks, which eventually + * breaks sink recovery (conflicting per-bucket offsets in the WriterState). + * + *

This resolver lazily fetches the authoritative bucket count of a partition from cluster + * metadata and caches it locally. A partition's bucket count is immutable once the partition is + * created, so a cached entry is valid forever and no invalidation is ever needed. When a partition + * is not found in the metadata (a dynamically created partition whose creation is triggered by the + * downstream writer, i.e. after this sharding step), the resolver falls back to the current + * table-level bucket number, which is exactly what the partition will be created with. + */ +public class PartitionBucketCountResolver implements Serializable { + + private static final long serialVersionUID = 1L; + + private static final Logger LOG = LoggerFactory.getLogger(PartitionBucketCountResolver.class); + + private static final int MAX_RETRIES = 3; + private static final long RETRY_INTERVAL_MS = 100L; + + private final TablePath tablePath; + private final Configuration flussConfig; + private final int defaultBucketCount; + + private transient Connection connection; + private transient Admin admin; + private transient ConcurrentHashMap bucketCounts; + + /** Fetches partition metadata from the cluster. Implementations must be thread-safe. */ + interface PartitionMetadataFetcher extends Serializable { + + /** + * Returns all partitions of the table, or {@code null} if the table itself is not + * retrievable (e.g. concurrently dropped). + */ + @Nullable + List listPartitionInfos(TablePath tablePath) throws Exception; + + /** Returns the current table-level bucket number. */ + int tableBucketCount(TablePath tablePath) throws Exception; + } + + public PartitionBucketCountResolver( + TablePath tablePath, Configuration flussConfig, int defaultBucketCount) { + this.tablePath = tablePath; + this.flussConfig = flussConfig; + this.defaultBucketCount = defaultBucketCount; + } + + /** + * Returns the actual bucket count of the given partition, fetching and caching it on the first + * access. Cached entries are valid forever because a partition's bucket count never changes. + */ + public int bucketCountOf(String partitionName) { + Integer cached = bucketCounts().get(partitionName); + if (cached != null) { + return cached; + } + return resolveAndCache(partitionName); + } + + private int resolveAndCache(String partitionName) { + Integer cached = bucketCounts().get(partitionName); + if (cached != null) { + return cached; + } + + List infos = fetchPartitionInfosWithRetry(); + if (infos != null) { + // Cache every partition carried by the response: one RPC warms up the whole table. + for (PartitionInfo info : infos) { + bucketCounts().putIfAbsent(info.getPartitionName(), info.getBucketCount()); + } + } + + Integer resolved = bucketCounts().get(partitionName); + if (resolved != null) { + return resolved; + } + + // The partition doesn't exist in the metadata yet: its creation is triggered by the + // downstream writer after the sharding step, so blocking here would deadlock. Fall back + // to the current table-level bucket count, which is exactly what the partition will be + // created with; the fallback value is cached for the job's lifetime. The only problematic + // window is a rescale landing between this fallback and the actual creation: the + // partition is then created with the new table-level count while the cached value + // remains the old one, so this partition keeps sharding with a wrong count until + // restart. + int fallbackBucketCount = tableBucketCountWithRetry(); + LOG.debug( + "Partition {} not found for table {}, fall back to the table-level bucket count {}.", + partitionName, + tablePath, + fallbackBucketCount); + bucketCounts().put(partitionName, fallbackBucketCount); + return fallbackBucketCount; + } + + private List fetchPartitionInfosWithRetry() { + return fetchWithRetry( + "list partition infos", () -> fetcher().listPartitionInfos(tablePath)); + } + + private int tableBucketCountWithRetry() { + return fetchWithRetry( + "get the table-level bucket count", () -> fetcher().tableBucketCount(tablePath)); + } + + /** + * Invokes the metadata call with bounded retries: up to {@link #MAX_RETRIES} attempts with a + * short interval in between, failing fast with a {@link FlussRuntimeException} once exhausted. + */ + private T fetchWithRetry(String operation, Callable call) { + Exception lastException = null; + for (int attempt = 1; attempt <= MAX_RETRIES; attempt++) { + try { + return call.call(); + } catch (Exception e) { + lastException = e; + LOG.warn( + "Failed to {} for table {} (attempt {}/{}).", + operation, + tablePath, + attempt, + MAX_RETRIES, + e); + sleepBeforeRetry(); + } + } + throw new FlussRuntimeException( + String.format( + "Failed to %s for table %s after %d retries.", + operation, tablePath, MAX_RETRIES), + lastException); + } + + private void sleepBeforeRetry() { + try { + Thread.sleep(RETRY_INTERVAL_MS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new FlussRuntimeException("Interrupted while retrying metadata fetch.", e); + } + } + + private ConcurrentHashMap bucketCounts() { + if (bucketCounts == null) { + bucketCounts = new ConcurrentHashMap<>(); + } + return bucketCounts; + } + + protected PartitionMetadataFetcher fetcher() { + // Static sharding mode: never touch metadata, always fall back to the default count. + if (flussConfig == null) { + return new StaticPartitionMetadataFetcher(defaultBucketCount); + } + return new AdminPartitionMetadataFetcher(this); + } + + /** Production fetcher backed by a lazily created Fluss {@link Admin}. */ + private static class AdminPartitionMetadataFetcher implements PartitionMetadataFetcher { + + private static final long serialVersionUID = 1L; + + private final PartitionBucketCountResolver owner; + + private AdminPartitionMetadataFetcher(PartitionBucketCountResolver owner) { + this.owner = owner; + } + + @Nullable + @Override + public List listPartitionInfos(TablePath tablePath) throws Exception { + return owner.admin().listPartitionInfos(tablePath).get(); + } + + @Override + public int tableBucketCount(TablePath tablePath) throws Exception { + return owner.admin().getTableInfo(tablePath).get().getNumBuckets(); + } + } + + /** Fallback fetcher used when no Fluss config is provided: always reports the default count. */ + private static class StaticPartitionMetadataFetcher implements PartitionMetadataFetcher { + + private static final long serialVersionUID = 1L; + + private final int defaultBucketCount; + + private StaticPartitionMetadataFetcher(int defaultBucketCount) { + this.defaultBucketCount = defaultBucketCount; + } + + @Override + public List listPartitionInfos(TablePath tablePath) { + return null; + } + + @Override + public int tableBucketCount(TablePath tablePath) { + return defaultBucketCount; + } + } + + private Admin admin() { + if (admin == null) { + connection = ConnectionFactory.createConnection(flussConfig); + admin = connection.getAdmin(); + } + return admin; + } +} diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlinkRowDataChannelComputerTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlinkRowDataChannelComputerTest.java index 93d5254fdd1..665b7875796 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlinkRowDataChannelComputerTest.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/FlinkRowDataChannelComputerTest.java @@ -52,12 +52,13 @@ void testSelectChanel() { Collections.singletonList("a"), Collections.emptyList(), null, + null, + null, 10, serializationSchema); for (int numChannel = 1; numChannel <= 10; numChannel++) { channelComputer.setup(numChannel); - assertThat(channelComputer.isCombineShuffleWithPartitionName()).isFalse(); for (int i = 0; i < 100; i++) { int expectedChannel = -1; for (int retry = 0; retry < 5; retry++) { @@ -82,16 +83,13 @@ void testSelectChanelForPartitionedTable() { Collections.singletonList("a"), Collections.singletonList("b"), null, + null, + null, 10, serializationSchema); for (int numChannel = 1; numChannel <= 10; numChannel++) { channelComputer.setup(numChannel); - if (10 % numChannel != 0) { - assertThat(channelComputer.isCombineShuffleWithPartitionName()).isTrue(); - } else { - assertThat(channelComputer.isCombineShuffleWithPartitionName()).isFalse(); - } for (int i = 0; i < 100; i++) { int expectedChannel = -1; for (int retry = 0; retry < 5; retry++) { diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/PartitionBucketCountResolverITCase.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/PartitionBucketCountResolverITCase.java new file mode 100644 index 00000000000..64a311f690c --- /dev/null +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/PartitionBucketCountResolverITCase.java @@ -0,0 +1,352 @@ +/* + * 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.client.admin.Admin; +import org.apache.fluss.client.table.Table; +import org.apache.fluss.client.table.scanner.ScanRecord; +import org.apache.fluss.client.table.scanner.log.LogScanner; +import org.apache.fluss.client.table.scanner.log.ScanRecords; +import org.apache.fluss.flink.sink.serializer.RowDataSerializationSchema; +import org.apache.fluss.flink.utils.FlinkTestBase; +import org.apache.fluss.metadata.PartitionInfo; +import org.apache.fluss.metadata.ResolvedPartitionSpec; +import org.apache.fluss.metadata.Schema; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TableChange; +import org.apache.fluss.metadata.TableDescriptor; +import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.types.DataTypes; + +import org.apache.flink.streaming.api.datastream.DataStream; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.table.data.RowData; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Integration tests for {@link PartitionBucketCountResolver} against a real Fluss cluster. + * + *

The scenario reproduces the failure this component fixes: a partitioned table is rescaled + * (bucket.num 4 -> 8), so pre- and post-rescale partitions with different bucket counts coexist in + * one table. Before the fix, the bucket shuffle used the stale table-level count, which scattered + * the records of one bucket across multiple writer subtasks and eventually broke sink recovery + * (conflicting per-bucket offsets in the WriterState). + */ +class PartitionBucketCountResolverITCase extends FlinkTestBase { + + // Parallelism 3 does not divide either bucket count (2 % 3 != 0, 4 % 3 != 0), so both + // the pre-rescale and the post-rescale partitions exercise the combine-mode sharding + // formula (hash(partition) + bucket) % parallelism — the path that a fixed table-level + // numBuckets could not cover. + private static final int OLD_BUCKET_NUM = 2; + private static final int NEW_BUCKET_NUM = 4; + + private static final Schema schema = + Schema.newBuilder() + .column("a", DataTypes.INT()) + .column("b", DataTypes.STRING()) + .column("c", DataTypes.STRING()) + .build(); + + private static final TableDescriptor tableDescriptor = + TableDescriptor.builder() + .schema(schema) + .distributedBy(OLD_BUCKET_NUM, "a") + .partitionedBy("c") + .build(); + + private static final List PARTITION_KEYS = Collections.singletonList("c"); + + private static final int PARALLELISM = 3; + + private StreamExecutionEnvironment env; + + @BeforeEach + void setup() throws Exception { + env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(PARALLELISM); + } + + @Test + void testResolveBucketCountsAgainstRealCluster() throws Exception { + TablePath tablePath = TablePath.of(DEFAULT_DB, "resolver_it_" + System.currentTimeMillis()); + createTable(tablePath, tableDescriptor); + // Old partition, created before the rescale: 4 buckets. + createPartition(tablePath, "2024-01"); + // Rescale: partitions created afterwards use 8 buckets. + alterBucketCount(tablePath, NEW_BUCKET_NUM); + createPartition(tablePath, "2024-02"); + + // Cold start with an empty cache: the resolver must hit real cluster metadata. + PartitionBucketCountResolver resolver = + new PartitionBucketCountResolver(tablePath, clientConf, NEW_BUCKET_NUM); + + // The pre-rescale partition keeps its own count — not the rescaled table-level value. + assertThat(resolver.bucketCountOf("2024-01")).isEqualTo(OLD_BUCKET_NUM); + // The post-rescale partition resolves to the new count. + assertThat(resolver.bucketCountOf("2024-02")).isEqualTo(NEW_BUCKET_NUM); + + // A partition that doesn't exist yet falls back to the current table-level count. + assertThat(resolver.bucketCountOf("2024-03")).isEqualTo(NEW_BUCKET_NUM); + + // The fallback matches the count the partition actually gets once created. + createPartition(tablePath, "2024-03"); + assertThat(bucketCountByPartitionName(tablePath)).containsEntry("2024-03", NEW_BUCKET_NUM); + + // Cached entries stay stable across repeated lookups. + assertThat(resolver.bucketCountOf("2024-01")).isEqualTo(OLD_BUCKET_NUM); + assertThat(resolver.bucketCountOf("2024-02")).isEqualTo(NEW_BUCKET_NUM); + assertThat(resolver.bucketCountOf("2024-03")).isEqualTo(NEW_BUCKET_NUM); + } + + @Test + void testCacheAvoidsRepeatedMetadataFetches() throws Exception { + TablePath tablePath = + TablePath.of(DEFAULT_DB, "resolver_cache_it_" + System.currentTimeMillis()); + createTable(tablePath, tableDescriptor); + createPartition(tablePath, "2024-01"); + alterBucketCount(tablePath, NEW_BUCKET_NUM); + createPartition(tablePath, "2024-02"); + + CountingResolver resolver = new CountingResolver(tablePath, admin, NEW_BUCKET_NUM); + + // Miss: one listPartitionInfos call warms up every existing partition of the table. + assertThat(resolver.bucketCountOf("2024-01")).isEqualTo(OLD_BUCKET_NUM); + assertThat(resolver.bucketCountOf("2024-02")).isEqualTo(NEW_BUCKET_NUM); + assertThat(resolver.listCalls.get()).isEqualTo(1); + + // Cache hits: repeated lookups never trigger further metadata fetches. + assertThat(resolver.bucketCountOf("2024-01")).isEqualTo(OLD_BUCKET_NUM); + assertThat(resolver.bucketCountOf("2024-02")).isEqualTo(NEW_BUCKET_NUM); + assertThat(resolver.listCalls.get()).isEqualTo(1); + + // A partition that doesn't exist yet falls back to the table-level count: the miss costs + // one list call (to confirm the partition is absent) plus one table-level lookup. + assertThat(resolver.bucketCountOf("2024-03")).isEqualTo(NEW_BUCKET_NUM); + assertThat(resolver.listCalls.get()).isEqualTo(2); + assertThat(resolver.tableCalls.get()).isEqualTo(1); + + // The fallback value is cached: repeated lookups of the same partition cost nothing. + assertThat(resolver.bucketCountOf("2024-03")).isEqualTo(NEW_BUCKET_NUM); + assertThat(resolver.listCalls.get()).isEqualTo(2); + assertThat(resolver.tableCalls.get()).isEqualTo(1); + } + + /** + * Core end-to-end verification: after a rescale, writing through the real sink topology must + * aggregate one bucket into exactly one writer subtask, for pre- and post-rescale partitions + * alike. Before the fix, the stale table-level count scattered the rows of one bucket across + * multiple subtasks, which made the WriterState conflict and the sink unrecoverable. + */ + @Test + void testBucketShuffleAggregatesOneBucketToOneSubtaskAfterRescale() throws Exception { + TablePath tablePath = TablePath.of(DEFAULT_DB, "shuffle_it_" + System.currentTimeMillis()); + createTable(tablePath, tableDescriptor); + createPartition(tablePath, "2024-01"); + alterBucketCount(tablePath, NEW_BUCKET_NUM); + createPartition(tablePath, "2024-02"); + + List partitionInfos = admin.listPartitionInfos(tablePath).get(); + Map countByPartitionName = bucketCountByPartitionName(tablePath); + Map partitionNameById = partitionNameById(partitionInfos); + + // The production-shaped channel computer backed by real cluster metadata. + RowDataSerializationSchema serializationSchema = + new RowDataSerializationSchema(true, false); + FlinkRowDataChannelComputer channelComputer = + new FlinkRowDataChannelComputer<>( + schema.getRowType(), + Collections.singletonList("a"), + PARTITION_KEYS, + null, + tablePath, + clientConf, + NEW_BUCKET_NUM, + serializationSchema); + channelComputer.setup(PARALLELISM); + + // Rows for both partitions: 10 distinct bucket keys, each submitted twice, so the sharded + // topology has to aggregate duplicates of one bucket into the same writer subtask. + List rows = new ArrayList<>(); + Map channelByRowKey = new HashMap<>(); + for (String partitionName : Arrays.asList("2024-01", "2024-02")) { + for (int a = 0; a < 10; a++) { + for (int dup = 0; dup < 2; dup++) { + RowData row = + org.apache.flink.table.data.GenericRowData.of( + a, + org.apache.flink.table.data.StringData.fromString("v" + a), + org.apache.flink.table.data.StringData.fromString( + partitionName)); + rows.add(row); + channelByRowKey.put(rowKey(row), channelComputer.channel(row)); + } + } + } + + // Write through the real sink topology: the partitioner inside it calls the same + // channel logic, so the recorded channels match the actual writer assignment. + FlussSink flussSink = + FlussSink.builder() + .setBootstrapServers(bootstrapServers) + .setDatabase(DEFAULT_DB) + .setTable(tablePath.getTableName()) + .setSerializationSchema(serializationSchema) + .build(); + DataStream stream = env.fromElements(rows.toArray(new RowData[0])); + stream.sinkTo(flussSink).name("Fluss Sink"); + env.execute("test bucket shuffle aggregates one bucket to one subtask after rescale"); + + // Scan every bucket of every partition and map each row back to its sharding channel. + Table table = conn.getTable(tablePath); + LogScanner logScanner = table.newScan().createLogScanner(); + for (PartitionInfo info : partitionInfos) { + for (int b = 0; b < info.getBucketCount(); b++) { + logScanner.subscribeFromBeginning(info.getPartitionId(), b); + } + } + + Map> channelsPerBucket = new HashMap<>(); + int collected = 0; + long deadline = System.currentTimeMillis() + 60_000; + while (collected < rows.size() && System.currentTimeMillis() < deadline) { + ScanRecords scanRecords = logScanner.poll(Duration.ofSeconds(1)); + for (TableBucket bucket : scanRecords.buckets()) { + for (ScanRecord record : scanRecords.records(bucket)) { + Integer channel = channelByRowKey.get(rowKey(record.getRow())); + assertThat(channel).isNotNull(); + channelsPerBucket.computeIfAbsent(bucket, k -> new HashSet<>()).add(channel); + collected++; + } + } + } + logScanner.close(); + + // Data completeness: every submitted row landed exactly once. + assertThat(collected).isEqualTo(rows.size()); + + // Core invariant: every bucket is written by exactly one writer subtask — all rows of a + // bucket share one sharding channel, and the channel matches the sharding formula derived + // from the partition's own bucket count (4 for the pre-rescale partition, 8 for the new). + for (Map.Entry> entry : channelsPerBucket.entrySet()) { + TableBucket bucket = entry.getKey(); + String partitionName = partitionNameById.get(bucket.getPartitionId()); + int bucketCount = countByPartitionName.get(partitionName); + int expectedChannel; + if (ChannelComputer.shouldCombinePartitionInSharding(true, bucketCount, PARALLELISM)) { + expectedChannel = + ChannelComputer.select(partitionName, bucket.getBucket(), PARALLELISM); + } else { + expectedChannel = ChannelComputer.select(bucket.getBucket(), PARALLELISM); + } + assertThat(entry.getValue()) + .as("bucket %s of partition %s", bucket.getBucket(), partitionName) + .containsExactly(expectedChannel); + } + } + + // -------------------------------------------------------------------------------------------- + + private static String rowKey(org.apache.flink.table.data.RowData row) { + return row.getInt(0) + "#" + row.getString(2).toString(); + } + + private static String rowKey(org.apache.fluss.row.InternalRow row) { + return row.getInt(0) + "#" + row.getString(2).toString(); + } + + private void createPartition(TablePath tablePath, String partitionName) throws Exception { + admin.createPartition( + tablePath, + ResolvedPartitionSpec.fromPartitionName(PARTITION_KEYS, partitionName) + .toPartitionSpec(), + false) + .get(); + } + + private void alterBucketCount(TablePath tablePath, int newBucketCount) throws Exception { + admin.alterTable( + tablePath, + Collections.singletonList(TableChange.modifyBucketCount(newBucketCount)), + false) + .get(); + } + + private Map bucketCountByPartitionName(TablePath tablePath) throws Exception { + Map map = new HashMap<>(); + for (PartitionInfo info : admin.listPartitionInfos(tablePath).get()) { + map.put(info.getPartitionName(), info.getBucketCount()); + } + return map; + } + + private Map partitionNameById(List partitionInfos) { + Map map = new HashMap<>(); + for (PartitionInfo info : partitionInfos) { + map.put(info.getPartitionId(), info.getPartitionName()); + } + return map; + } + + /** Resolver variant that delegates to the test cluster's admin and counts metadata calls. */ + private static final class CountingResolver extends PartitionBucketCountResolver { + + private static final long serialVersionUID = 1L; + + private final Admin admin; + private final AtomicInteger listCalls = new AtomicInteger(); + private final AtomicInteger tableCalls = new AtomicInteger(); + + private CountingResolver(TablePath tablePath, Admin admin, int defaultBucketCount) { + super(tablePath, null, defaultBucketCount); + this.admin = admin; + } + + @Override + protected PartitionMetadataFetcher fetcher() { + return new PartitionMetadataFetcher() { + @Override + public List listPartitionInfos(TablePath tablePath) + throws Exception { + listCalls.incrementAndGet(); + return admin.listPartitionInfos(tablePath).get(); + } + + @Override + public int tableBucketCount(TablePath tablePath) throws Exception { + tableCalls.incrementAndGet(); + return admin.getTableInfo(tablePath).get().getNumBuckets(); + } + }; + } + } +} diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/PartitionBucketCountResolverTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/PartitionBucketCountResolverTest.java new file mode 100644 index 00000000000..6d2fd351c2c --- /dev/null +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/PartitionBucketCountResolverTest.java @@ -0,0 +1,193 @@ +/* + * 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.VisibleForTesting; +import org.apache.fluss.exception.FlussRuntimeException; +import org.apache.fluss.metadata.PartitionInfo; +import org.apache.fluss.metadata.ResolvedPartitionSpec; +import org.apache.fluss.metadata.TablePath; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Test for {@link PartitionBucketCountResolver}. */ +class PartitionBucketCountResolverTest { + + private static final List PARTITION_KEYS = Collections.singletonList("dt"); + + private static PartitionInfo partitionInfo(String partitionName, int bucketCount) { + return new PartitionInfo( + partitionName.hashCode(), + ResolvedPartitionSpec.fromPartitionName(PARTITION_KEYS, partitionName), + null, + bucketCount); + } + + /** Records every fetch invocation so tests can assert how often metadata is queried. */ + private static class CountingFetcher + implements PartitionBucketCountResolver.PartitionMetadataFetcher { + + private final List partitions; + private final int tableBucketCount; + private int fetchCount; + + private CountingFetcher(List partitions, int tableBucketCount) { + this.partitions = partitions; + this.tableBucketCount = tableBucketCount; + } + + @Override + public List listPartitionInfos(TablePath tablePath) { + fetchCount++; + return partitions; + } + + @Override + public int tableBucketCount(TablePath tablePath) { + return tableBucketCount; + } + } + + /** + * Test-only resolver variant that injects a fake metadata source, keeping the production class + * free of any test wiring. + */ + @VisibleForTesting + private static final class FakeSourceResolver extends PartitionBucketCountResolver { + + private static final long serialVersionUID = 1L; + + private final PartitionMetadataFetcher fetcher; + + private FakeSourceResolver(int defaultBucketCount, PartitionMetadataFetcher fetcher) { + super(null, null, defaultBucketCount); + this.fetcher = fetcher; + } + + @Override + protected PartitionMetadataFetcher fetcher() { + return fetcher; + } + } + + @Test + void testExistingPartitionResolvesAuthoritativeCountAndIsCached() { + CountingFetcher fetcher = + new CountingFetcher( + Arrays.asList(partitionInfo("2024-01", 4), partitionInfo("2024-02", 8)), 8); + PartitionBucketCountResolver resolver = new FakeSourceResolver(8, fetcher); + + // The partition exists: its authoritative count wins over the table-level default. + assertThat(resolver.bucketCountOf("2024-01")).isEqualTo(4); + assertThat(resolver.bucketCountOf("2024-02")).isEqualTo(8); + + // One RPC warms up every partition of the table; subsequent lookups never fetch again. + assertThat(fetcher.fetchCount).isEqualTo(1); + assertThat(resolver.bucketCountOf("2024-01")).isEqualTo(4); + assertThat(resolver.bucketCountOf("2024-02")).isEqualTo(8); + assertThat(fetcher.fetchCount).isEqualTo(1); + } + + @Test + void testMissingPartitionFallsBackToCurrentTableLevelCount() { + // Only the old partition exists: the new partition is not created yet (its creation is + // triggered by the downstream writer after the sharding step). + CountingFetcher fetcher = + new CountingFetcher(Collections.singletonList(partitionInfo("2024-01", 4)), 8); + PartitionBucketCountResolver resolver = new FakeSourceResolver(8, fetcher); + + // The missing partition falls back to the current table-level bucket count. + assertThat(resolver.bucketCountOf("2024-02")).isEqualTo(8); + assertThat(fetcher.fetchCount).isEqualTo(1); + + // The fallback value is cached as well, so no further RPC happens. + assertThat(resolver.bucketCountOf("2024-02")).isEqualTo(8); + assertThat(fetcher.fetchCount).isEqualTo(1); + } + + @Test + void testEveryPartitionOfTheResponseIsCached() { + List partitions = new ArrayList<>(); + for (int i = 0; i < 5; i++) { + partitions.add(partitionInfo("2024-0" + (i + 1), 4)); + } + CountingFetcher fetcher = new CountingFetcher(partitions, 8); + PartitionBucketCountResolver resolver = new FakeSourceResolver(8, fetcher); + + // Resolving one partition warms up all partitions carried by the same response. + assertThat(resolver.bucketCountOf("2024-03")).isEqualTo(4); + assertThat(fetcher.fetchCount).isEqualTo(1); + for (int i = 0; i < 5; i++) { + assertThat(resolver.bucketCountOf("2024-0" + (i + 1))).isEqualTo(4); + } + assertThat(fetcher.fetchCount).isEqualTo(1); + } + + @Test + void testMetadataFailureRetriesThenFailsFast() { + PartitionBucketCountResolver.PartitionMetadataFetcher failingFetcher = + new PartitionBucketCountResolver.PartitionMetadataFetcher() { + @Override + public List listPartitionInfos(TablePath tablePath) + throws Exception { + throw new RuntimeException("metadata rpc error"); + } + + @Override + public int tableBucketCount(TablePath tablePath) { + return 8; + } + }; + PartitionBucketCountResolver resolver = new FakeSourceResolver(8, failingFetcher); + + assertThatThrownBy(() -> resolver.bucketCountOf("2024-01")) + .isInstanceOf(FlussRuntimeException.class) + .hasMessageContaining("after 3 retries"); + } + + @Test + void testTableLookupFailureRetriesThenFailsFast() { + // The table itself is not retrievable (null response): fall through to the table-level + // fallback, whose own failure must surface after the bounded retries. + PartitionBucketCountResolver.PartitionMetadataFetcher failingFetcher = + new PartitionBucketCountResolver.PartitionMetadataFetcher() { + @Override + public List listPartitionInfos(TablePath tablePath) { + return null; + } + + @Override + public int tableBucketCount(TablePath tablePath) throws Exception { + throw new RuntimeException("table rpc error"); + } + }; + PartitionBucketCountResolver resolver = new FakeSourceResolver(8, failingFetcher); + + assertThatThrownBy(() -> resolver.bucketCountOf("2024-01")) + .isInstanceOf(FlussRuntimeException.class) + .hasMessageContaining("after 3 retries"); + } +} diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/UndoRecoveryITCase.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/UndoRecoveryITCase.java index 529aea661e8..c685f197ce7 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/UndoRecoveryITCase.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/UndoRecoveryITCase.java @@ -23,10 +23,14 @@ import org.apache.fluss.client.admin.ProducerOffsetsResult; import org.apache.fluss.client.lookup.Lookuper; import org.apache.fluss.client.table.Table; +import org.apache.fluss.client.table.scanner.batch.BatchScanUtils; +import org.apache.fluss.client.table.scanner.batch.BatchScanner; import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.flink.sink.serializer.RowDataSerializationSchema; import org.apache.fluss.flink.sink.testutils.CountingSource; import org.apache.fluss.flink.sink.testutils.FailingCountingSource; +import org.apache.fluss.metadata.ResolvedPartitionSpec; +import org.apache.fluss.metadata.TableChange; import org.apache.fluss.metadata.TablePath; import org.apache.fluss.row.InternalRow; import org.apache.fluss.server.testutils.FlussClusterExtension; @@ -34,6 +38,8 @@ import org.apache.flink.api.common.JobID; import org.apache.flink.api.common.eventtime.WatermarkStrategy; +import org.apache.flink.api.common.functions.MapFunction; +import org.apache.flink.api.common.typeinfo.TypeInformation; import org.apache.flink.configuration.CheckpointingOptions; import org.apache.flink.configuration.Configuration; import org.apache.flink.configuration.MemorySize; @@ -42,6 +48,7 @@ import org.apache.flink.core.execution.JobClient; import org.apache.flink.core.execution.SavepointFormatType; import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration; +import org.apache.flink.streaming.api.datastream.DataStream; import org.apache.flink.streaming.api.datastream.DataStreamSource; import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; import org.apache.flink.table.api.DataTypes; @@ -49,7 +56,12 @@ import org.apache.flink.table.api.Schema; import org.apache.flink.table.api.TableResult; import org.apache.flink.table.api.bridge.java.StreamTableEnvironment; +import org.apache.flink.table.data.GenericRowData; import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.runtime.typeutils.InternalTypeInfo; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.RowType; import org.apache.flink.test.util.MiniClusterWithClientResource; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -63,6 +75,9 @@ import java.io.File; import java.time.Duration; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; import java.util.concurrent.TimeUnit; import static org.apache.fluss.flink.FlinkConnectorOptions.BOOTSTRAP_SERVERS; @@ -99,6 +114,9 @@ *

    *
  • Rescale Up: {@link #testRescaleUp()} - Parallelism 1 → 2 *
  • Rescale Down: {@link #testRescaleDown()} - Parallelism 2 → 1 + *
  • Partitioned Bucket Rescale: {@link + * #testPartitionedTableRecoveryAfterBucketRescale()} - Tests recovery when the table's bucket + * num is rescaled across job submissions *
* *

Note: Rescale tests continue to use savepoints because Flink checkpoints are bound to @@ -459,6 +477,132 @@ void testRescaleDown() throws Exception { new int[] {10, 5, 3}); } + /** + * Tests that a partitioned primary key table sink recovers after the table's bucket num is + * rescaled between job submissions. + * + *

This reproduces the mixed-layout hazard the per-partition bucket count resolver fixes: + * after an ALTER bucket.num, the table-level value captured by a job submitted afterwards no + * longer matches the counts of partitions created before the rescale. Before the fix, the + * bucket shuffle sharded the pre-rescale partition with the stale table-level value, splitting + * one bucket's keys across multiple writer subtasks; the next restore then failed with + * conflicting per-bucket offsets in the WriterState, leaving the sink unrecoverable. + * + *

Pattern: + * + *

    + *
  1. Phase 1: write 10 records per key into partition 2024-01 (1 bucket), + * stop-with-savepoint + *
  2. Rescale: table-level bucket.num 1 -> 2; partition 2024-02 created with 2 buckets + *
  3. Phase 2: NEW job submission (graph captures the post-rescale table-level bucket num 2) + * restores the savepoint and writes 5 records per key into 2024-01. With the fix, the + * channel computer resolves 2024-01's actual count (1) and keeps one writer per bucket. + *
  4. Phase 3: NEW job submission restores phase 2's savepoint and writes 3 records per key + * into 2024-02. Before the fix, restoring the scattered WriterState threw conflicting + * checkpoint offsets and the job never recovered. + *
+ * + *

Parallelism 3 divides neither bucket count, so every partition shards in combine mode. + * Keys 1 and 2 resolve to bucket 0 and key 3 to bucket 1 of the pre-rescale partition (for + * bucketing mod 2), which makes the pre-fix scatter deterministic. + */ + @Test + void testPartitionedTableRecoveryAfterBucketRescale() throws Exception { + String tableName = "undo_partition_rescale_" + System.currentTimeMillis(); + TablePath tablePath = TablePath.of(DEFAULT_DB, tableName); + String producerId = "test-producer-partition-rescale-" + System.currentTimeMillis(); + final int oldBucketNum = 1; + final int newBucketNum = 2; + final int parallelism = 3; + + initTableEnvironment(null, false) + .executeSql(createPartitionedAggTableDDL(tableName, oldBucketNum)); + + // The pre-rescale partition keeps its single-bucket layout across the rescale. + createPartition(tablePath, "2024-01"); + + // Phase 1: pre-rescale job submission (graph captures table-level bucket.num = 1). + JobClient phase1Job = + startPartitionedUndoJob(tablePath, producerId, null, 10, parallelism, "2024-01"); + long expectedAfterPhase1 = 10 * VALUE_PER_RECORD; + retry( + DEFAULT_TIMEOUT, + () -> { + Map sums = scanSumsByPartitionAndKey(tablePath); + for (long key = 1; key <= 3; key++) { + assertThat(sums.get("2024-01:" + key)) + .as("Key %d sum before savepoint", key) + .isEqualTo(expectedAfterPhase1); + } + }); + + waitForCheckpoint(phase1Job.getJobID()); + String savepoint1 = + phase1Job + .stopWithSavepoint( + false, + savepointDir.getAbsolutePath(), + SavepointFormatType.CANONICAL) + .get(60, TimeUnit.SECONDS); + waitForJobTermination(phase1Job, DEFAULT_TIMEOUT); + + // Rescale: partitions created afterwards use the new table-level bucket num. + admin.alterTable( + tablePath, + Collections.singletonList(TableChange.modifyBucketCount(newBucketNum)), + false) + .get(); + createPartition(tablePath, "2024-02"); + + // Phase 2: post-rescale job submission (graph captures table-level bucket.num = 2). + JobClient phase2Job = + startPartitionedUndoJob( + tablePath, producerId, savepoint1, 5, parallelism, "2024-01"); + long expectedAfterPhase2 = 15 * VALUE_PER_RECORD; + retry( + DEFAULT_TIMEOUT, + () -> { + Map sums = scanSumsByPartitionAndKey(tablePath); + for (long key = 1; key <= 3; key++) { + assertThat(sums.get("2024-01:" + key)) + .as("Key %d sum before phase 2 savepoint", key) + .isEqualTo(expectedAfterPhase2); + } + }); + + waitForCheckpoint(phase2Job.getJobID()); + String savepoint2 = + phase2Job + .stopWithSavepoint( + false, + savepointDir.getAbsolutePath(), + SavepointFormatType.CANONICAL) + .get(60, TimeUnit.SECONDS); + waitForJobTermination(phase2Job, DEFAULT_TIMEOUT); + + // Phase 3: restore phase 2's savepoint and write into the post-rescale partition. + JobClient phase3Job = + startPartitionedUndoJob( + tablePath, producerId, savepoint2, 3, parallelism, "2024-02"); + long expectedPhase3 = 3 * VALUE_PER_RECORD; + retry( + DEFAULT_TIMEOUT, + () -> { + Map sums = scanSumsByPartitionAndKey(tablePath); + for (long key = 1; key <= 3; key++) { + assertThat(sums.get("2024-01:" + key)) + .as("Key %d sum in the pre-rescale partition after recovery", key) + .isEqualTo(expectedAfterPhase2); + assertThat(sums.get("2024-02:" + key)) + .as("Key %d sum in the post-rescale partition", key) + .isEqualTo(expectedPhase3); + } + }); + + phase3Job.cancel().get(); + waitForJobTermination(phase3Job, DEFAULT_TIMEOUT); + } + /** * Tests that undo recovery works through the SQL/Table API sink path. * @@ -782,6 +926,22 @@ private String createAggTableDDL(String tableName, int bucketNum) { DEFAULT_DB, tableName, bucketNum); } + private String createPartitionedAggTableDDL(String tableName, int bucketNum) { + return String.format( + "CREATE TABLE `%s`.`%s` (" + + " id BIGINT NOT NULL," + + " part STRING NOT NULL," + + " sum_val BIGINT," + + " PRIMARY KEY (id, part) NOT ENFORCED" + + ") PARTITIONED BY (part)" + + " WITH (" + + " 'bucket.num' = '%d'," + + " 'table.merge-engine' = 'aggregation'," + + " 'fields.sum_val.agg' = 'sum'" + + ")", + DEFAULT_DB, tableName, bucketNum); + } + /** * Unified method to start a bounded streaming job. * @@ -905,6 +1065,81 @@ private JobClient startFailoverJob( return env.executeAsync(jobName); } + /** + * Starts a bounded undo-recovery job against a partitioned primary key table. + * + *

The counting source emits two-field (key, value) rows; a map extends them to the table + * schema (id, part, sum_val) with a fixed partition value, so all records of a phase land in + * one pre-created partition. + * + * @param tablePath target table + * @param producerId producer ID for undo recovery + * @param restorePath savepoint (or externalized checkpoint) to restore from, or null + * @param maxRecordsPerKey records to emit per key + * @param parallelism job parallelism + * @param partitionValue the partition value written to + */ + private JobClient startPartitionedUndoJob( + TablePath tablePath, + String producerId, + @Nullable String restorePath, + int maxRecordsPerKey, + int parallelism, + String partitionValue) + throws Exception { + + Configuration conf = new Configuration(); + if (restorePath != null) { + conf.setString("execution.savepoint.path", restorePath); + } + + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(conf); + env.setParallelism(parallelism); + env.enableCheckpointing(1000); + + DataStreamSource stream = + env.fromSource( + CountingSource.multiKey(VALUE_PER_RECORD, maxRecordsPerKey), + WatermarkStrategy.noWatermarks(), + "counting-source"); + + TypeInformation rowTypeInfo = + InternalTypeInfo.of( + RowType.of( + new LogicalType[] { + DataTypes.BIGINT().getLogicalType(), + DataTypes.STRING().getLogicalType(), + DataTypes.BIGINT().getLogicalType() + }, + new String[] {"id", "part", "sum_val"})); + DataStream partitionedStream = + stream.map( + (MapFunction) + row -> { + GenericRowData newRow = new GenericRowData(3); + newRow.setField(0, row.getLong(0)); + newRow.setField( + 1, StringData.fromString(partitionValue)); + newRow.setField(2, row.getLong(1)); + return newRow; + }) + .returns(rowTypeInfo); + + FlussSink sink = + FlussSink.builder() + .setBootstrapServers(bootstrapServers) + .setDatabase(tablePath.getDatabaseName()) + .setTable(tablePath.getTableName()) + .setProducerId(producerId) + .setSerializationSchema(new RowDataSerializationSchema(false, true)) + .build(); + + partitionedStream.sinkTo(sink).name("Fluss Sink"); + + return env.executeAsync( + "Partitioned undo test (p=" + parallelism + ", part=" + partitionValue + ")"); + } + /** * Starts a job with FailingCountingSource for checkpoint-based failover testing. * @@ -1009,6 +1244,34 @@ private void verifyProducerOffsetsCleanedUp(String producerId) { LOG.info("Verified producer offsets cleaned up for {}", producerId); } + private void createPartition(TablePath tablePath, String partitionName) throws Exception { + admin.createPartition( + tablePath, + ResolvedPartitionSpec.fromPartitionName( + Collections.singletonList("part"), partitionName) + .toPartitionSpec(), + false) + .get(); + } + + /** + * Batch scans the whole table and folds the latest value of each primary key per partition. + * + * @return map from "{partition}:{id}" to the aggregated sum + */ + private Map scanSumsByPartitionAndKey(TablePath tablePath) throws Exception { + Map sums = new HashMap<>(); + try (Table table = conn.getTable(tablePath); + BatchScanner scanner = table.newScan().createBatchScanner()) { + for (InternalRow row : BatchScanUtils.collectRows(scanner)) { + String partition = row.getString(1).toString(); + long id = row.getLong(0); + sums.merge(partition + ":" + id, row.getLong(2), Long::sum); + } + } + return sums; + } + @Nullable private Long lookupSum(TablePath tablePath, Long key) throws Exception { try (Table table = conn.getTable(tablePath)) { From 1e021c1767c78da4609833ddb6db8cd54676502c Mon Sep 17 00:00:00 2001 From: duankaixuan <1417048384@qq.com> Date: Sun, 13 Sep 2026 14:20:11 +0800 Subject: [PATCH 2/2] [test] Merge redundant resolver ITCase --- .../PartitionBucketCountResolverITCase.java | 48 +++++-------------- 1 file changed, 12 insertions(+), 36 deletions(-) diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/PartitionBucketCountResolverITCase.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/PartitionBucketCountResolverITCase.java index 64a311f690c..c060d77e39d 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/PartitionBucketCountResolverITCase.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/PartitionBucketCountResolverITCase.java @@ -97,51 +97,23 @@ void setup() throws Exception { } @Test - void testResolveBucketCountsAgainstRealCluster() throws Exception { + void testResolveBucketCountsAndCachingWithRealCluster() throws Exception { TablePath tablePath = TablePath.of(DEFAULT_DB, "resolver_it_" + System.currentTimeMillis()); createTable(tablePath, tableDescriptor); - // Old partition, created before the rescale: 4 buckets. + // Old partition, created before the rescale: 2 buckets. createPartition(tablePath, "2024-01"); - // Rescale: partitions created afterwards use 8 buckets. + // Rescale: partitions created afterwards use 4 buckets. alterBucketCount(tablePath, NEW_BUCKET_NUM); createPartition(tablePath, "2024-02"); - // Cold start with an empty cache: the resolver must hit real cluster metadata. - PartitionBucketCountResolver resolver = - new PartitionBucketCountResolver(tablePath, clientConf, NEW_BUCKET_NUM); + // Cold start with an empty cache: the resolver must hit real cluster metadata. A single + // listPartitionInfos call warms up every existing partition of the table. + CountingResolver resolver = new CountingResolver(tablePath, admin, NEW_BUCKET_NUM); // The pre-rescale partition keeps its own count — not the rescaled table-level value. assertThat(resolver.bucketCountOf("2024-01")).isEqualTo(OLD_BUCKET_NUM); // The post-rescale partition resolves to the new count. assertThat(resolver.bucketCountOf("2024-02")).isEqualTo(NEW_BUCKET_NUM); - - // A partition that doesn't exist yet falls back to the current table-level count. - assertThat(resolver.bucketCountOf("2024-03")).isEqualTo(NEW_BUCKET_NUM); - - // The fallback matches the count the partition actually gets once created. - createPartition(tablePath, "2024-03"); - assertThat(bucketCountByPartitionName(tablePath)).containsEntry("2024-03", NEW_BUCKET_NUM); - - // Cached entries stay stable across repeated lookups. - assertThat(resolver.bucketCountOf("2024-01")).isEqualTo(OLD_BUCKET_NUM); - assertThat(resolver.bucketCountOf("2024-02")).isEqualTo(NEW_BUCKET_NUM); - assertThat(resolver.bucketCountOf("2024-03")).isEqualTo(NEW_BUCKET_NUM); - } - - @Test - void testCacheAvoidsRepeatedMetadataFetches() throws Exception { - TablePath tablePath = - TablePath.of(DEFAULT_DB, "resolver_cache_it_" + System.currentTimeMillis()); - createTable(tablePath, tableDescriptor); - createPartition(tablePath, "2024-01"); - alterBucketCount(tablePath, NEW_BUCKET_NUM); - createPartition(tablePath, "2024-02"); - - CountingResolver resolver = new CountingResolver(tablePath, admin, NEW_BUCKET_NUM); - - // Miss: one listPartitionInfos call warms up every existing partition of the table. - assertThat(resolver.bucketCountOf("2024-01")).isEqualTo(OLD_BUCKET_NUM); - assertThat(resolver.bucketCountOf("2024-02")).isEqualTo(NEW_BUCKET_NUM); assertThat(resolver.listCalls.get()).isEqualTo(1); // Cache hits: repeated lookups never trigger further metadata fetches. @@ -149,12 +121,16 @@ void testCacheAvoidsRepeatedMetadataFetches() throws Exception { assertThat(resolver.bucketCountOf("2024-02")).isEqualTo(NEW_BUCKET_NUM); assertThat(resolver.listCalls.get()).isEqualTo(1); - // A partition that doesn't exist yet falls back to the table-level count: the miss costs - // one list call (to confirm the partition is absent) plus one table-level lookup. + // A partition that doesn't exist yet falls back to the current table-level count; the + // miss costs one list call (to confirm the partition is absent) plus one table lookup. assertThat(resolver.bucketCountOf("2024-03")).isEqualTo(NEW_BUCKET_NUM); assertThat(resolver.listCalls.get()).isEqualTo(2); assertThat(resolver.tableCalls.get()).isEqualTo(1); + // The fallback matches the count the partition actually gets once created. + createPartition(tablePath, "2024-03"); + assertThat(bucketCountByPartitionName(tablePath)).containsEntry("2024-03", NEW_BUCKET_NUM); + // The fallback value is cached: repeated lookups of the same partition cost nothing. assertThat(resolver.bucketCountOf("2024-03")).isEqualTo(NEW_BUCKET_NUM); assertThat(resolver.listCalls.get()).isEqualTo(2);