From 0ac88b2403f1711c5f105367685bcda927e985ac Mon Sep 17 00:00:00 2001 From: David Wang Date: Fri, 25 Sep 2026 09:45:34 +1000 Subject: [PATCH 1/5] [core] Route overwrite writes by partition bucket layout --- .../java/org/apache/paimon/CoreOptions.java | 14 + .../operation/AbstractFileStoreWrite.java | 64 +++- .../paimon/operation/FileStoreWrite.java | 7 + .../operation/FileSystemWriteRestore.java | 27 +- .../apache/paimon/operation/WriteRestore.java | 41 ++- .../paimon/table/AbstractFileStoreTable.java | 8 +- .../table/AppendOnlyFileStoreTable.java | 9 +- .../paimon/table/DelegatedFileStoreTable.java | 6 + .../apache/paimon/table/FileStoreTable.java | 7 + .../table/PrimaryKeyFileStoreTable.java | 16 +- .../table/SchemaBucketFileStoreTable.java | 100 ++++++ .../sink/FixedBucketRowKeyExtractor.java | 39 ++- .../table/sink/FixedBucketWriteSelector.java | 9 +- .../table/sink/PartitionBucketMapping.java | 166 ++++++++++ .../paimon/table/sink/RowKeyExtractor.java | 22 +- .../paimon/table/sink/TableWriteImpl.java | 6 + .../paimon/operation/FileStoreCommitTest.java | 42 +++ .../operation/FileSystemWriteRestoreTest.java | 303 +++++++++++++++++- .../table/AppendOnlySimpleTableTest.java | 61 ++-- .../sink/FixedBucketRowKeyExtractorTest.java | 53 ++- .../sink/FixedBucketWriteSelectorTest.java | 86 +++++ .../sink/PartitionBucketMappingTest.java | 87 +++++ 22 files changed, 1110 insertions(+), 63 deletions(-) create mode 100644 paimon-core/src/main/java/org/apache/paimon/table/SchemaBucketFileStoreTable.java create mode 100644 paimon-core/src/main/java/org/apache/paimon/table/sink/PartitionBucketMapping.java create mode 100644 paimon-core/src/test/java/org/apache/paimon/table/sink/FixedBucketWriteSelectorTest.java create mode 100644 paimon-core/src/test/java/org/apache/paimon/table/sink/PartitionBucketMappingTest.java diff --git a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java index 9818e68d9ca9..254a53455672 100644 --- a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java +++ b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java @@ -156,6 +156,16 @@ public class CoreOptions implements Serializable { .withDescription( "Whether to ignore the order of the buckets when reading data from an append-only table."); + public static final ConfigOption BUCKET_PER_PARTITION_COUNT_ENABLED = + key("bucket.per-partition-count-enabled") + .booleanType() + .defaultValue(false) + .withDescription( + "Whether to allow individual partitions of a fixed-bucket table to keep " + + "their own bucket count, so that a single partition can be rescaled " + + "independently. Enabling this scans the manifest to resolve each " + + "partition's bucket count on write."); + @Immutable public static final ConfigOption BUCKET_FUNCTION_TYPE = key("bucket-function.type") @@ -3201,6 +3211,10 @@ public int bucket() { return options.get(BUCKET); } + public boolean bucketPerPartitionCountEnabled() { + return options.get(BUCKET_PER_PARTITION_COUNT_ENABLED); + } + public BucketFunctionType bucketFunctionType() { return options.get(BUCKET_FUNCTION_TYPE); } diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java b/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java index 0c0738c6173e..d17490181b5e 100644 --- a/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java +++ b/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java @@ -40,6 +40,7 @@ import org.apache.paimon.partition.PartitionTimeExtractor; import org.apache.paimon.table.sink.CommitMessage; import org.apache.paimon.table.sink.CommitMessageImpl; +import org.apache.paimon.table.sink.PartitionBucketMapping; import org.apache.paimon.types.RowType; import org.apache.paimon.utils.CommitIncrement; import org.apache.paimon.utils.ExecutorThreadFactory; @@ -152,6 +153,15 @@ public FileStoreWrite withWriteRestore(WriteRestore writeRestore) { return this; } + @Override + public FileStoreWrite withPartitionBucketMapping( + PartitionBucketMapping partitionBucketMapping) { + if (restore instanceof FileSystemWriteRestore) { + ((FileSystemWriteRestore) restore).withPartitionBucketMapping(partitionBucketMapping); + } + return this; + } + @Override public FileStoreWrite withIOManager(IOManager ioManager) { this.ioManager = ioManager; @@ -503,7 +513,8 @@ private WriterContainer getWriterWrapper(BinaryRow partition, int bucket, int partition, totalBuckets, buckets.values().iterator().next().totalBuckets); } return buckets.computeIfAbsent( - bucket, k -> createWriterContainer(partition.copy(), bucket, totalBuckets)); + bucket, + k -> createWriterContainer(partition.copy(), bucket, totalBuckets, true, true)); } private Map> getWriterContainers(BinaryRow partition) { @@ -520,16 +531,20 @@ public RecordWriter createWriter(BinaryRow partition, int bucket) { } public WriterContainer createWriterContainer(BinaryRow partition, int bucket) { - return createWriterContainer(partition, bucket, numBuckets, !ignoreNumBucketCheck); + return createWriterContainer(partition, bucket, numBuckets, !ignoreNumBucketCheck, false); } private WriterContainer createWriterContainer( BinaryRow partition, int bucket, int totalBuckets) { - return createWriterContainer(partition, bucket, totalBuckets, true); + return createWriterContainer(partition, bucket, totalBuckets, true, true); } private WriterContainer createWriterContainer( - BinaryRow partition, int bucket, int expectedTotalBuckets, boolean validateNumBuckets) { + BinaryRow partition, + int bucket, + int expectedTotalBuckets, + boolean validateNumBuckets, + boolean strictBucketCount) { if (LOG.isDebugEnabled()) { LOG.debug("Creating writer for partition {}, bucket {}", partition, bucket); } @@ -554,7 +569,11 @@ private WriterContainer createWriterContainer( if (!actualIgnorePreviousFiles) { restored = scanExistingFileMetas( - partition, bucket, expectedTotalBuckets, validateNumBuckets); + partition, + bucket, + expectedTotalBuckets, + validateNumBuckets, + strictBucketCount); } DynamicBucketIndexMaintainer indexMaintainer = @@ -647,7 +666,11 @@ public FileStoreWrite withMetricRegistry(MetricRegistry metricRegistry) { } private RestoreFiles scanExistingFileMetas( - BinaryRow partition, int bucket, int expectedTotalBuckets, boolean validateNumBuckets) { + BinaryRow partition, + int bucket, + int expectedTotalBuckets, + boolean validateNumBuckets, + boolean strictBucketCount) { Supplier partInfo = () -> partitionType.getFieldCount() > 0 @@ -674,8 +697,33 @@ private RestoreFiles scanExistingFileMetas( partInfo.get(), bucket), e); } - if (restored.totalBuckets() != null && validateNumBuckets) { - checkNumBuckets(partInfo.get(), expectedTotalBuckets, restored.totalBuckets()); + Integer restoredTotalBuckets = restored.totalBuckets(); + if (restoredTotalBuckets != null + && validateNumBuckets + && expectedTotalBuckets != restoredTotalBuckets) { + if (partitionType.getFieldCount() > 0 + && options.bucketPerPartitionCountEnabled() + && !strictBucketCount) { + if (bucket >= restoredTotalBuckets) { + throw new RuntimeException( + String.format( + "Trying to write bucket %d to %s, but the partition only has %d " + + "buckets (table default: %d). Recompute the bucket using the " + + "partition's bucket count, or rescale the partition via " + + "INSERT OVERWRITE.", + bucket, + partInfo.get(), + restoredTotalBuckets, + expectedTotalBuckets)); + } + LOG.info( + "{} uses {} buckets (expected: {}). Accepting per-partition bucket count.", + partInfo.get(), + restoredTotalBuckets, + expectedTotalBuckets); + } else { + checkNumBuckets(partInfo.get(), expectedTotalBuckets, restoredTotalBuckets); + } } return restored; } diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/FileStoreWrite.java b/paimon-core/src/main/java/org/apache/paimon/operation/FileStoreWrite.java index 6945120d8ae5..3f4a7b9d2d75 100644 --- a/paimon-core/src/main/java/org/apache/paimon/operation/FileStoreWrite.java +++ b/paimon-core/src/main/java/org/apache/paimon/operation/FileStoreWrite.java @@ -30,6 +30,7 @@ import org.apache.paimon.mergetree.compact.CompactRewriterFactory; import org.apache.paimon.metrics.MetricRegistry; import org.apache.paimon.table.sink.CommitMessage; +import org.apache.paimon.table.sink.PartitionBucketMapping; import org.apache.paimon.table.sink.SinkRecord; import org.apache.paimon.types.RowType; import org.apache.paimon.utils.CommitIncrement; @@ -53,6 +54,12 @@ public interface FileStoreWrite extends Restorable withWriteRestore(WriteRestore writeRestore); + /** Provides the preloaded partition-to-bucket mapping for fixed-bucket writes. */ + default FileStoreWrite withPartitionBucketMapping( + PartitionBucketMapping partitionBucketMapping) { + return this; + } + FileStoreWrite withIOManager(IOManager ioManager); /** Specified the write rowType. */ diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/FileSystemWriteRestore.java b/paimon-core/src/main/java/org/apache/paimon/operation/FileSystemWriteRestore.java index 740e7399f087..522a61eda07b 100644 --- a/paimon-core/src/main/java/org/apache/paimon/operation/FileSystemWriteRestore.java +++ b/paimon-core/src/main/java/org/apache/paimon/operation/FileSystemWriteRestore.java @@ -25,11 +25,11 @@ import org.apache.paimon.index.IndexFileMeta; import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.manifest.ManifestEntry; +import org.apache.paimon.table.sink.PartitionBucketMapping; import org.apache.paimon.utils.SnapshotManager; import javax.annotation.Nullable; -import java.util.ArrayList; import java.util.List; import static org.apache.paimon.deletionvectors.DeletionVectorsIndexFile.DELETION_VECTORS_INDEX; @@ -37,9 +37,11 @@ /** {@link WriteRestore} to restore files directly from file system. */ public class FileSystemWriteRestore implements WriteRestore { + private final CoreOptions options; private final SnapshotManager snapshotManager; private final FileStoreScan scan; private final IndexFileHandler indexFileHandler; + @Nullable private PartitionBucketMapping partitionBucketMapping; private final @Nullable Long snapshotId; public FileSystemWriteRestore( @@ -65,6 +67,7 @@ private FileSystemWriteRestore( FileStoreScan scan, IndexFileHandler indexFileHandler, @Nullable Long snapshotId) { + this.options = options; this.snapshotManager = snapshotManager; this.scan = scan; this.indexFileHandler = indexFileHandler; @@ -74,6 +77,21 @@ private FileSystemWriteRestore( this.scan.dropStats(); } } + this.partitionBucketMapping = + options.bucketPerPartitionCountEnabled() + ? null + : PartitionBucketMapping.defaultBuckets(options.bucket()); + } + + public void withPartitionBucketMapping(PartitionBucketMapping partitionBucketMapping) { + this.partitionBucketMapping = partitionBucketMapping; + } + + private PartitionBucketMapping partitionBucketMapping() { + if (partitionBucketMapping == null) { + partitionBucketMapping = PartitionBucketMapping.loadFromScan(scan, options.bucket()); + } + return partitionBucketMapping; } @Override @@ -101,10 +119,13 @@ public RestoreFiles restoreFiles( return RestoreFiles.empty(); } - List restoreFiles = new ArrayList<>(); + // load the mapping before narrowing the mutable scan to a single bucket + PartitionBucketMapping bucketMapping = partitionBucketMapping(); List entries = scan.withSnapshot(snapshot).withPartitionBucket(partition, bucket).plan().files(); - Integer totalBuckets = WriteRestore.extractDataFiles(entries, restoreFiles); + List restoreFiles = WriteRestore.extractDataFiles(entries); + + Integer totalBuckets = WriteRestore.extractTotalBuckets(entries, partition, bucketMapping); IndexFileMeta dynamicBucketIndex = null; if (scanDynamicBucketIndex) { diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/WriteRestore.java b/paimon-core/src/main/java/org/apache/paimon/operation/WriteRestore.java index f57d9ab05515..f2ed55137827 100644 --- a/paimon-core/src/main/java/org/apache/paimon/operation/WriteRestore.java +++ b/paimon-core/src/main/java/org/apache/paimon/operation/WriteRestore.java @@ -21,9 +21,11 @@ import org.apache.paimon.data.BinaryRow; import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.manifest.ManifestEntry; +import org.apache.paimon.table.sink.PartitionBucketMapping; import javax.annotation.Nullable; +import java.util.ArrayList; import java.util.List; /** Restore for write to restore data files by partition and bucket from file system. */ @@ -38,9 +40,44 @@ RestoreFiles restoreFiles( boolean scanDeleteVectorsIndex, boolean scanSourceIndexPayloads); + /** + * Resolves the {@code totalBuckets} for a (partition, bucket) pair given the manifest entries + * for that bucket and the table's partition-bucket mapping. + * + *
    + *
  • Non-empty bucket: use the value stamped on the existing data files so that + * committer-side bucket-count mismatch detection (e.g. rescale-without-overwrite) still + * fires. + *
  • Empty bucket on a partitioned table: look up the per-partition override in {@code + * mapping}; returns {@code null} if the partition uses the table default. + *
  • Empty bucket on an unpartitioned table: returns {@code null} so the write path falls + * back to {@code numBuckets} and the committer-side check still fires. + *
+ */ @Nullable - static Integer extractDataFiles(List entries, List dataFiles) { + static Integer extractTotalBuckets( + List entries, BinaryRow partition, PartitionBucketMapping mapping) { + if (!entries.isEmpty()) { + return entries.get(0).totalBuckets(); + } + if (partition.getFieldCount() > 0) { + return mapping.getNumBucketsOverride(partition); + } + return null; + } + + /** + * Extracts the {@link DataFileMeta} list from the given manifest entries, validating that all + * entries agree on {@code totalBuckets}. + * + * @param entries manifest entries for a single (partition, bucket) pair + * @return the list of data files; empty if {@code entries} is empty + * @throws RuntimeException if entries carry inconsistent {@code totalBuckets} values, which + * indicates a corrupted manifest + */ + static List extractDataFiles(List entries) { Integer totalBuckets = null; + List dataFiles = new ArrayList<>(); for (ManifestEntry entry : entries) { if (totalBuckets != null && totalBuckets != entry.totalBuckets()) { throw new RuntimeException( @@ -51,6 +88,6 @@ static Integer extractDataFiles(List entries, List totalBuckets = entry.totalBuckets(); dataFiles.add(entry.file()); } - return totalBuckets; + return dataFiles; } } diff --git a/paimon-core/src/main/java/org/apache/paimon/table/AbstractFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/AbstractFileStoreTable.java index 93ecc173d5d2..c83486186c27 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/AbstractFileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/AbstractFileStoreTable.java @@ -44,6 +44,7 @@ import org.apache.paimon.table.sink.DynamicBucketRowKeyExtractor; import org.apache.paimon.table.sink.FixedBucketRowKeyExtractor; import org.apache.paimon.table.sink.FixedBucketWriteSelector; +import org.apache.paimon.table.sink.PartitionBucketMapping; import org.apache.paimon.table.sink.PostponeBucketRowKeyExtractor; import org.apache.paimon.table.sink.RowKeyExtractor; import org.apache.paimon.table.sink.RowKindGenerator; @@ -252,7 +253,9 @@ public Optional statistics() { public Optional newWriteSelector() { switch (bucketMode()) { case HASH_FIXED: - return Optional.of(new FixedBucketWriteSelector(schema())); + return Optional.of( + new FixedBucketWriteSelector( + schema(), PartitionBucketMapping.loadFromTable(this))); case BUCKET_UNAWARE: case POSTPONE_MODE: return Optional.empty(); @@ -280,7 +283,8 @@ protected CatalogEnvironment newCatalogEnvironment(String branch) { public RowKeyExtractor createRowKeyExtractor() { switch (bucketMode()) { case HASH_FIXED: - return new FixedBucketRowKeyExtractor(schema()); + return new FixedBucketRowKeyExtractor( + schema(), PartitionBucketMapping.loadFromTable(this)); case HASH_DYNAMIC: case KEY_DYNAMIC: return new DynamicBucketRowKeyExtractor(schema()); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/AppendOnlyFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/AppendOnlyFileStoreTable.java index 63710698ca3b..31ff9bd4f458 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/AppendOnlyFileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/AppendOnlyFileStoreTable.java @@ -30,6 +30,7 @@ import org.apache.paimon.predicate.Predicate; import org.apache.paimon.schema.TableSchema; import org.apache.paimon.table.query.LocalTableQuery; +import org.apache.paimon.table.sink.RowKeyExtractor; import org.apache.paimon.table.sink.TableWriteImpl; import org.apache.paimon.table.source.AppendBatchTableScan; import org.apache.paimon.table.source.AppendOnlySplitGenerator; @@ -162,11 +163,17 @@ public TableWriteImpl newWrite(String commitUser) { @Override public TableWriteImpl newWrite(String commitUser, @Nullable Integer writeId) { + return newWrite(commitUser, writeId, createRowKeyExtractor()); + } + + @Override + public TableWriteImpl newWrite( + String commitUser, @Nullable Integer writeId, RowKeyExtractor rowKeyExtractor) { BaseAppendFileStoreWrite writer = store().newWrite(commitUser, writeId); return new TableWriteImpl<>( rowType(), writer, - createRowKeyExtractor(), + rowKeyExtractor, (record, rowKind) -> { Preconditions.checkState( rowKind.isAdd(), diff --git a/paimon-core/src/main/java/org/apache/paimon/table/DelegatedFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/DelegatedFileStoreTable.java index ed541721fca0..f8c0856e7c48 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/DelegatedFileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/DelegatedFileStoreTable.java @@ -356,6 +356,12 @@ public TableWriteImpl newWrite(String commitUser, @Nullable Integer writeId) return wrapped.newWrite(commitUser, writeId); } + @Override + public TableWriteImpl newWrite( + String commitUser, @Nullable Integer writeId, RowKeyExtractor rowKeyExtractor) { + return wrapped.newWrite(commitUser, writeId, rowKeyExtractor); + } + @Override public TableWriteImpl newPostponeFixedBucketWrite( String commitUser, @Nullable Integer writeId) { diff --git a/paimon-core/src/main/java/org/apache/paimon/table/FileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/FileStoreTable.java index ddaa17407624..c19e0cd6a2dd 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/FileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/FileStoreTable.java @@ -144,6 +144,13 @@ default PostponeFixedBucketWriteBuilder newPostponeFixedBucketWriteBuilder() { TableWriteImpl newWrite(String commitUser, @Nullable Integer writeId); + /** + * Creates a new write with a custom {@link RowKeyExtractor}. This is useful for scenarios such + * as rescaling where the bucket assignment logic needs to be overridden. + */ + TableWriteImpl newWrite( + String commitUser, @Nullable Integer writeId, RowKeyExtractor rowKeyExtractor); + /** Creates a fixed-bucket merge-tree write for a postpone-bucket batch write. */ default TableWriteImpl newPostponeFixedBucketWrite( String commitUser, @Nullable Integer writeId) { diff --git a/paimon-core/src/main/java/org/apache/paimon/table/PrimaryKeyFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/PrimaryKeyFileStoreTable.java index 97c541175bd3..894308fb9dc8 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/PrimaryKeyFileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/PrimaryKeyFileStoreTable.java @@ -33,6 +33,7 @@ import org.apache.paimon.schema.KeyValueFieldsExtractor; import org.apache.paimon.schema.TableSchema; import org.apache.paimon.table.query.LocalTableQuery; +import org.apache.paimon.table.sink.RowKeyExtractor; import org.apache.paimon.table.sink.TableWriteImpl; import org.apache.paimon.table.source.DataTableScan; import org.apache.paimon.table.source.InnerTableRead; @@ -194,22 +195,29 @@ public TableWriteImpl newWrite(String commitUser) { @Override public TableWriteImpl newWrite(String commitUser, @Nullable Integer writeId) { - return newWrite(store().newWrite(commitUser, writeId)); + return newWrite(store().newWrite(commitUser, writeId), createRowKeyExtractor()); + } + + @Override + public TableWriteImpl newWrite( + String commitUser, @Nullable Integer writeId, RowKeyExtractor rowKeyExtractor) { + return newWrite(store().newWrite(commitUser, writeId), rowKeyExtractor); } @Override public TableWriteImpl newPostponeFixedBucketWrite( String commitUser, @Nullable Integer writeId) { - return newWrite(store().newPostponeFixedBucketWrite(commitUser)); + return newWrite(store().newPostponeFixedBucketWrite(commitUser), createRowKeyExtractor()); } - private TableWriteImpl newWrite(AbstractFileStoreWrite storeWrite) { + private TableWriteImpl newWrite( + AbstractFileStoreWrite storeWrite, RowKeyExtractor rowKeyExtractor) { KeyValue kv = new KeyValue(); List deleteNotNullFields = deleteNotNullFields(coreOptions()); return new TableWriteImpl<>( rowType(), storeWrite, - createRowKeyExtractor(), + rowKeyExtractor, (record, rowKind) -> kv.replace( record.primaryKey(), diff --git a/paimon-core/src/main/java/org/apache/paimon/table/SchemaBucketFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/SchemaBucketFileStoreTable.java new file mode 100644 index 000000000000..69ab8c4c809c --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/SchemaBucketFileStoreTable.java @@ -0,0 +1,100 @@ +/* + * 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.paimon.table; + +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.table.sink.FixedBucketRowKeyExtractor; +import org.apache.paimon.table.sink.FixedBucketWriteSelector; +import org.apache.paimon.table.sink.PartitionBucketMapping; +import org.apache.paimon.table.sink.RowKeyExtractor; +import org.apache.paimon.table.sink.TableWriteImpl; +import org.apache.paimon.table.sink.WriteSelector; + +import javax.annotation.Nullable; + +import java.util.Map; +import java.util.Optional; + +/** + * A {@link FileStoreTable} wrapper that uses the schema number of buckets assign writes instead of + * using the number of buckets defined in each partition. Useful for postpone buckets, overrides and + * rescales. + */ +public class SchemaBucketFileStoreTable extends DelegatedFileStoreTable { + + public SchemaBucketFileStoreTable(FileStoreTable wrapped) { + super(wrapped); + } + + @Override + public Optional newWriteSelector() { + return Optional.of( + new FixedBucketWriteSelector( + schema(), PartitionBucketMapping.defaultBuckets(schema().numBuckets()))); + } + + @Override + public RowKeyExtractor createRowKeyExtractor() { + return new FixedBucketRowKeyExtractor( + schema(), PartitionBucketMapping.defaultBuckets(schema().numBuckets())); + } + + @Override + public TableWriteImpl newWrite(String commitUser) { + return newWrite(commitUser, null); + } + + @Override + public TableWriteImpl newWrite(String commitUser, @Nullable Integer writeId) { + return wrapped().newWrite(commitUser, writeId, createRowKeyExtractor()); + } + + @Override + public TableWriteImpl newWrite( + String commitUser, @Nullable Integer writeId, RowKeyExtractor rowKeyExtractor) { + // Always use the schema-bucket-based extractor; ignore the caller-supplied extractor + // to ensure consistent per-partition bucket routing even when called via the 3-arg form. + return wrapped().newWrite(commitUser, writeId, createRowKeyExtractor()); + } + + @Override + public FileStoreTable copy(Map dynamicOptions) { + return new SchemaBucketFileStoreTable(wrapped().copy(dynamicOptions)); + } + + @Override + public FileStoreTable copy(TableSchema newTableSchema) { + return new SchemaBucketFileStoreTable(wrapped().copy(newTableSchema)); + } + + @Override + public FileStoreTable copyWithoutTimeTravel(Map dynamicOptions) { + return new SchemaBucketFileStoreTable(wrapped().copyWithoutTimeTravel(dynamicOptions)); + } + + @Override + public FileStoreTable copyWithLatestSchema() { + return new SchemaBucketFileStoreTable(wrapped().copyWithLatestSchema()); + } + + @Override + public FileStoreTable switchToBranch(String branchName) { + return new SchemaBucketFileStoreTable(wrapped().switchToBranch(branchName)); + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/sink/FixedBucketRowKeyExtractor.java b/paimon-core/src/main/java/org/apache/paimon/table/sink/FixedBucketRowKeyExtractor.java index 146a45b43713..8f78e702b07c 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/sink/FixedBucketRowKeyExtractor.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/sink/FixedBucketRowKeyExtractor.java @@ -29,24 +29,31 @@ /** {@link KeyAndBucketExtractor} for {@link InternalRow}. */ public class FixedBucketRowKeyExtractor extends RowKeyExtractor { - private final int numBuckets; + private transient Projection bucketKeyProjection; + private final boolean sameBucketKeyAndTrimmedPrimaryKey; - private final Projection bucketKeyProjection; + private final PartitionBucketMapping partitionBucketMapping; private BinaryRow reuseBucketKey; private Integer reuseBucket; private final BucketFunction bucketFunction; public FixedBucketRowKeyExtractor(TableSchema schema) { + this(schema, PartitionBucketMapping.defaultBuckets(schema.numBuckets())); + } + + public FixedBucketRowKeyExtractor( + TableSchema schema, PartitionBucketMapping partitionBucketMapping) { super(schema); - numBuckets = new CoreOptions(schema.options()).bucket(); - bucketFunction = - BucketFunction.create( - new CoreOptions(schema.options()), schema.logicalBucketKeyType()); - sameBucketKeyAndTrimmedPrimaryKey = schema.bucketKeys().equals(schema.trimmedPrimaryKeys()); - bucketKeyProjection = - CodeGenUtils.newProjection( - schema.logicalRowType(), schema.projection(schema.bucketKeys())); + this.bucketFunction = createBucketFunction(schema); + this.sameBucketKeyAndTrimmedPrimaryKey = + schema.bucketKeys().equals(schema.trimmedPrimaryKeys()); + this.partitionBucketMapping = partitionBucketMapping; + } + + private static BucketFunction createBucketFunction(TableSchema schema) { + return BucketFunction.create( + new CoreOptions(schema.options()), schema.logicalBucketKeyType()); } @Override @@ -62,7 +69,7 @@ private BinaryRow bucketKey() { } if (reuseBucketKey == null) { - reuseBucketKey = bucketKeyProjection.apply(record); + reuseBucketKey = bucketKeyProjection().apply(record); } return reuseBucketKey; } @@ -70,6 +77,7 @@ private BinaryRow bucketKey() { @Override public int bucket() { if (reuseBucket == null) { + int numBuckets = partitionBucketMapping.resolveNumBuckets(partition()); reuseBucket = bucket(numBuckets); } return reuseBucket; @@ -78,4 +86,13 @@ public int bucket() { public int bucket(int numBuckets) { return bucketFunction.bucket(bucketKey(), numBuckets); } + + private Projection bucketKeyProjection() { + if (bucketKeyProjection == null) { + bucketKeyProjection = + CodeGenUtils.newProjection( + schema.logicalRowType(), schema.projection(schema.bucketKeys())); + } + return bucketKeyProjection; + } } diff --git a/paimon-core/src/main/java/org/apache/paimon/table/sink/FixedBucketWriteSelector.java b/paimon-core/src/main/java/org/apache/paimon/table/sink/FixedBucketWriteSelector.java index e08841dd8cd3..893a6bf53f77 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/sink/FixedBucketWriteSelector.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/sink/FixedBucketWriteSelector.java @@ -28,17 +28,24 @@ public class FixedBucketWriteSelector implements WriteSelector { private static final long serialVersionUID = 1L; private final TableSchema schema; + private final PartitionBucketMapping partitionBucketMapping; private transient KeyAndBucketExtractor extractor; public FixedBucketWriteSelector(TableSchema schema) { + this(schema, PartitionBucketMapping.defaultBuckets(schema.numBuckets())); + } + + public FixedBucketWriteSelector( + TableSchema schema, PartitionBucketMapping partitionBucketMapping) { this.schema = schema; + this.partitionBucketMapping = partitionBucketMapping; } @Override public int select(InternalRow row, int numWriters) { if (extractor == null) { - extractor = new FixedBucketRowKeyExtractor(schema); + extractor = new FixedBucketRowKeyExtractor(schema, partitionBucketMapping); } extractor.setRecord(row); return ChannelComputer.select(extractor.partition(), extractor.bucket(), numWriters); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/sink/PartitionBucketMapping.java b/paimon-core/src/main/java/org/apache/paimon/table/sink/PartitionBucketMapping.java new file mode 100644 index 000000000000..666ddf127b2d --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/sink/PartitionBucketMapping.java @@ -0,0 +1,166 @@ +/* + * 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.paimon.table.sink; + +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.manifest.PartitionEntry; +import org.apache.paimon.operation.FileStoreScan; +import org.apache.paimon.table.FileStoreTable; + +import javax.annotation.Nullable; + +import java.io.Serializable; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * A serializable mapping that resolves the number of buckets for each partition in a table. + * + *

Different partitions may have different bucket counts (e.g., after a rescale operation). This + * class maintains a per-partition bucket count mapping and falls back to a default bucket count for + * partitions that are not explicitly mapped. + * + *

This is used by components such as {@link FixedBucketRowKeyExtractor} and {@link + * FixedBucketWriteSelector} to correctly determine the bucket assignment for rows in tables where + * partitions may have been rescaled independently. + * + * @see #loadFromTable(FileStoreTable) + * @see #resolveNumBuckets(BinaryRow) + */ +public class PartitionBucketMapping implements Serializable { + + private static final long serialVersionUID = 1L; + + /** The default number of buckets, used when a partition has no explicit mapping. */ + private final int defaultBucketCount; + + /** A map from partition to its specific bucket count. May be empty but never {@code null}. */ + private final Map partitionBucketMap; + + /** + * Creates a mapping with a default bucket count and an explicit per-partition bucket map. + * + * @param defaultBucketCount the default number of buckets, used as a fallback + * @param partitionBucketMap a map from partition (as {@link BinaryRow}) to its bucket count + */ + public PartitionBucketMapping( + int defaultBucketCount, Map partitionBucketMap) { + this.defaultBucketCount = defaultBucketCount; + this.partitionBucketMap = partitionBucketMap; + } + + /** + * Creates a mapping with only a default bucket count and no per-partition overrides. + * + *

Use this when per-partition bucket counts are not needed (i.e. the additional manifest + * scan is skipped), so that every partition resolves to {@code numBuckets}. + * + * @param numBuckets the default number of buckets for all partitions + * @return a mapping that resolves every partition to {@code numBuckets} + */ + public static PartitionBucketMapping defaultBuckets(int numBuckets) { + return new PartitionBucketMapping(numBuckets, Collections.emptyMap()); + } + + /** + * Loads a {@link PartitionBucketMapping} by scanning the manifest entries of the given table. + * + *

For non-partitioned tables, this returns a mapping with only the schema-defined default + * bucket count and an empty partition map. + * + *

For partitioned tables, the method reads {@link + * org.apache.paimon.manifest.PartitionEntry}s, which aggregate manifest entries per partition + * during the scan and therefore have a much smaller memory footprint than loading all data file + * entries. Any scan failure is propagated to the caller. + * + * @param table the {@link FileStoreTable} to load the mapping from + * @return a {@link PartitionBucketMapping} reflecting the current bucket layout of the table + */ + public static PartitionBucketMapping loadFromTable(FileStoreTable table) { + int defaultBuckets = table.schema().numBuckets(); + if (!table.coreOptions().bucketPerPartitionCountEnabled() + || table.partitionKeys().isEmpty()) { + return defaultBuckets(defaultBuckets); + } + return loadFromScan(table.store().newScan(), defaultBuckets); + } + + /** + * Loads a {@link PartitionBucketMapping} from the given scan by reading the manifest partition + * entries to resolve the per-partition bucket counts. + * + *

Callers should evaluate whether per-partition bucket counts are enabled before + * invoking this method, since it always triggers an additional scan. Use {@link + * #defaultBuckets(int)} instead when the scan should be skipped. + */ + public static PartitionBucketMapping loadFromScan(FileStoreScan scan, int defaultBuckets) { + if (scan == null) { + return defaultBuckets(defaultBuckets); + } + List partitionEntries = scan.readPartitionEntries(); + Map partitionBucketMap = new HashMap<>(); + for (PartitionEntry entry : partitionEntries) { + int totalBuckets = entry.totalBuckets(); + // Only store partitions whose bucket count differs from the default. + // This keeps the map empty for partitions that have never been rescaled, + // avoiding per-partition BinaryRow copies and Integer allocations entirely. + if (totalBuckets > 0 && totalBuckets != defaultBuckets) { + partitionBucketMap.put(entry.partition().copy(), totalBuckets); + } + } + return new PartitionBucketMapping(defaultBuckets, partitionBucketMap); + } + + /** + * Returns the explicit bucket-count override for the given partition, if one exists. + * + *

This method does not fall back to the table-level default. It is used by restore paths to + * distinguish an explicitly rescaled partition from an unseen partition which should continue + * using the writer's expected bucket count. + * + * @param partition the partition key as a {@link BinaryRow} + * @return the explicit bucket-count override, or {@code null} when the partition uses the + * table-level default + */ + @Nullable + public Integer getNumBucketsOverride(BinaryRow partition) { + return partitionBucketMap.get(partition); + } + + /** + * Resolves the number of buckets for the given partition. + * + *

If the partition has an explicit entry in the partition-to-bucket map, that value is + * returned. Otherwise, the default bucket count is returned. + * + * @param partition the partition key as a {@link BinaryRow} + * @return the number of buckets for the given partition + */ + public int resolveNumBuckets(BinaryRow partition) { + if (partitionBucketMap != null) { + Integer partitionBucketCount = partitionBucketMap.get(partition); + if (partitionBucketCount != null) { + return partitionBucketCount; + } + } + return defaultBucketCount; + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/sink/RowKeyExtractor.java b/paimon-core/src/main/java/org/apache/paimon/table/sink/RowKeyExtractor.java index 455aaa4aa5e9..697734ca10b1 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/sink/RowKeyExtractor.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/sink/RowKeyExtractor.java @@ -22,18 +22,23 @@ import org.apache.paimon.data.InternalRow; import org.apache.paimon.schema.TableSchema; +import java.io.Serializable; + /** {@link KeyAndBucketExtractor} for {@link InternalRow}. */ -public abstract class RowKeyExtractor implements KeyAndBucketExtractor { +public abstract class RowKeyExtractor implements KeyAndBucketExtractor, Serializable { + + private static final long serialVersionUID = 1L; - private final RowPartitionKeyExtractor partitionKeyExtractor; + private transient RowPartitionKeyExtractor partitionKeyExtractor; + protected final TableSchema schema; protected InternalRow record; private BinaryRow partition; private BinaryRow trimmedPrimaryKey; public RowKeyExtractor(TableSchema schema) { - this.partitionKeyExtractor = new RowPartitionKeyExtractor(schema); + this.schema = schema; } @Override @@ -46,7 +51,7 @@ public void setRecord(InternalRow record) { @Override public BinaryRow partition() { if (partition == null) { - partition = partitionKeyExtractor.partition(record); + partition = partitionKeyExtractor().partition(record); } return partition; } @@ -54,8 +59,15 @@ public BinaryRow partition() { @Override public BinaryRow trimmedPrimaryKey() { if (trimmedPrimaryKey == null) { - trimmedPrimaryKey = partitionKeyExtractor.trimmedPrimaryKey(record); + trimmedPrimaryKey = partitionKeyExtractor().trimmedPrimaryKey(record); } return trimmedPrimaryKey; } + + private RowPartitionKeyExtractor partitionKeyExtractor() { + if (partitionKeyExtractor == null) { + partitionKeyExtractor = new RowPartitionKeyExtractor(schema); + } + return partitionKeyExtractor; + } } diff --git a/paimon-core/src/main/java/org/apache/paimon/table/sink/TableWriteImpl.java b/paimon-core/src/main/java/org/apache/paimon/table/sink/TableWriteImpl.java index 3d728aeed79a..daf99340e721 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/sink/TableWriteImpl.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/sink/TableWriteImpl.java @@ -123,6 +123,12 @@ public TableWriteImpl withIgnorePreviousFiles(boolean ignorePreviousFiles) { return this; } + public TableWriteImpl withPartitionBucketMapping( + PartitionBucketMapping partitionBucketMapping) { + write.withPartitionBucketMapping(partitionBucketMapping); + return this; + } + @Override public TableWriteImpl withIOManager(IOManager ioManager) { write.withIOManager(ioManager); diff --git a/paimon-core/src/test/java/org/apache/paimon/operation/FileStoreCommitTest.java b/paimon-core/src/test/java/org/apache/paimon/operation/FileStoreCommitTest.java index 7a0d26b1607e..4695b32f3efe 100644 --- a/paimon-core/src/test/java/org/apache/paimon/operation/FileStoreCommitTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/operation/FileStoreCommitTest.java @@ -1764,6 +1764,48 @@ public void testCommitRetryAfterFalseSuccessDoesNotCleanManifest() throws Except assertThat(store.readKvsFromSnapshot(latestSnapshot.id())).hasSize(1); } + @Test + public void testBucketCountConsistencyValidation() throws Exception { + TestFileStore store = createStore(false); + + // Commit initial data + List data = generateDataList(10); + store.commitData(data, gen::getPartition, kv -> 0); + + // Re-commit the same data but with a different totalBuckets value. + // This simulates a stale writer that loaded an old bucket mapping. + assertThatThrownBy( + () -> + store.commitDataImpl( + data, + gen::getPartition, + kv -> 0, + false, + null, + null, + Collections.emptyList(), + (commit, committable) -> { + ManifestCommittable tampered = + new ManifestCommittable( + committable.identifier(), + committable.watermark()); + for (CommitMessage msg : + committable.fileCommittables()) { + CommitMessageImpl impl = (CommitMessageImpl) msg; + tampered.addFileCommittable( + new CommitMessageImpl( + impl.partition(), + impl.bucket(), + 99, + impl.newFilesIncrement(), + impl.compactIncrement())); + } + commit.commit(tampered, true); + })) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("without overwrite"); + } + @Test public void testCommitRetryReusePreviousManifestMergeResultWhenBeforeStillExists() throws Exception { diff --git a/paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java b/paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java index 69d41094c4b6..767a6f114b8c 100644 --- a/paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java @@ -20,24 +20,75 @@ import org.apache.paimon.CoreOptions; import org.apache.paimon.Snapshot; +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.data.BinaryRowWriter; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.local.LocalFileIO; import org.apache.paimon.index.IndexFileHandler; import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.options.Options; +import org.apache.paimon.schema.FileSystemSchemaManager; +import org.apache.paimon.schema.Schema; +import org.apache.paimon.schema.SchemaUtils; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.table.CatalogEnvironment; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.FileStoreTableFactory; +import org.apache.paimon.table.sink.InnerTableWrite; +import org.apache.paimon.table.sink.PartitionBucketMapping; +import org.apache.paimon.table.sink.StreamTableCommit; +import org.apache.paimon.table.sink.StreamTableWrite; +import org.apache.paimon.table.sink.TableWriteImpl; +import org.apache.paimon.types.DataType; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; import org.apache.paimon.utils.SnapshotManager; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import java.util.Arrays; import java.util.Collections; import java.util.HashMap; +import java.util.UUID; import static org.apache.paimon.data.BinaryRow.EMPTY_ROW; +import static org.apache.paimon.table.BucketMode.POSTPONE_BUCKET; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -/** Tests for {@link FileSystemWriteRestore}. */ -class FileSystemWriteRestoreTest { +/** + * Tests for {@link FileSystemWriteRestore}, covering the {@code totalBuckets} resolution logic for + * both empty and non-empty buckets across partitioned and unpartitioned tables. + * + *

When restoring files for a {@code (partition, bucket)} that has no existing data files, there + * are no manifest entries to derive {@code totalBuckets} from. For partitioned tables, {@link + * WriteRestore#extractTotalBuckets} falls back to {@link + * org.apache.paimon.table.sink.PartitionBucketMapping} to correctly return the per-partition bucket + * count (e.g. after a rescale). For unpartitioned tables, {@code null} is returned so the write + * path falls back to {@code numBuckets} and the committer-side mismatch check still fires. + */ +public class FileSystemWriteRestoreTest { + + @TempDir java.nio.file.Path tempDir; + + private static final RowType ROW_TYPE = + RowType.of( + new DataType[] {DataTypes.INT(), DataTypes.INT(), DataTypes.BIGINT()}, + new String[] {"pt", "k", "v"}); + + @Test + void testEmptyBucketDoesNotReportPostponeTableDefault() { + PartitionBucketMapping mapping = PartitionBucketMapping.defaultBuckets(POSTPONE_BUCKET); + + assertThat(WriteRestore.extractTotalBuckets(Collections.emptyList(), binaryRow(1), mapping)) + .isNull(); + } @Test void testRestoreFromPinnedSnapshotForPostponeBucket() { @@ -102,4 +153,252 @@ void testRestoreSourceIndexPayloadsWithoutDirectory() { assertThat(restored.sourceIndexPayloads()).containsExactly(ann); } + + @Test + public void testEmptyBucketUsesPartitionBucketMapping() throws Exception { + // Build a table with default bucket=4 and write data into partition 1. + // Some buckets within partition 1 will end up with files (bucket 0 OR + // bucket 1, depending on hash); the OTHER bucket will be empty. Then + // "rescale" the table-level default to 32 (without rewriting partition 1) + // and ask the WriteRestore for an empty bucket. It must return + // totalBuckets=4 (the partition's actual bucket count), NOT 32 (the new + // table default). + FileStoreTable table = createPartitionedPkTable(4); + + // Write enough rows to populate at least one bucket within partition 1. + commitOneRow(table, /* pt */ 1, /* k */ 1); + commitOneRow(table, /* pt */ 1, /* k */ 2); + + // Find an empty bucket in partition 1 by inspecting the existing files. + int emptyBucket = findEmptyBucket(table, 1, /* totalBuckets */ 4); + + // Simulate a rescale by raising the table-level default bucket count + // (without rewriting existing files). Existing manifest entries still + // carry totalBuckets=4. + table = withBucket(table, 32); + + WriteRestore restore = newWriteRestore(table); + + RestoreFiles restored = + restore.restoreFiles(binaryRow(1), emptyBucket, false, false, false); + + assertThat(restored.totalBuckets()) + .as( + "Empty (partition 1, bucket %d): totalBuckets must be inferred from " + + "PartitionBucketMapping (4), not the new table default (32).", + emptyBucket) + .isEqualTo(4); + assertThat(restored.dataFiles()).isNullOrEmpty(); + } + + @Test + public void testEmptyBucketInUnseenPartitionDoesNotReportTableDefault() throws Exception { + // For an entirely unseen partition (no files anywhere), no per-partition override exists. + // Return null so the writer falls back to its expected table-level bucket count. In + // particular, this preserves postpone-bucket writes whose table default is -2 while the + // staged writer assigns real fixed buckets. + FileStoreTable table = createPartitionedPkTable(8); + commitOneRow(table, 1, 100); // ensures the snapshot exists + + WriteRestore restore = newWriteRestore(table); + RestoreFiles restored = + restore.restoreFiles(binaryRow(/* unseen */ 999), 0, false, false, false); + + assertThat(restored.totalBuckets()).isNull(); + assertThat(restored.dataFiles()).isNullOrEmpty(); + } + + @Test + public void testWriteRejectsBucketOutsidePartitionLayout() throws Exception { + // Partition 1 is created with 2 buckets. + FileStoreTable table = createPartitionedPkTable(2); + commitOneRow(table, /* pt */ 1, /* k */ 1); + + // Simulate a rescale: the table default is raised to 8 buckets, but partition 1 + // still only has 2 buckets. writeOnly=false so the writer scans previous files and + // runs the per-partition bucket-layout check in AbstractFileStoreWrite. + FileStoreTable rescaledTable = withBucket(table, 8); + + // Writing an out-of-range bucket (>= the partition's 2 buckets) into an empty bucket of + // partition 1 must be rejected, even though the bucket id is valid for the 8-bucket + // default. + // This is the bucket that PartitionBucketMapping recovery would otherwise silently accept. + // write(row, bucket) routes the row (partition pt=1) to the explicitly given bucket. + try (InnerTableWrite write = rescaledTable.newWrite(UUID.randomUUID().toString())) { + assertThatThrownBy(() -> write.write(GenericRow.of(1, 1, 1L), /* bucket */ 6)) + .hasMessageContaining("only has 2 buckets") + .hasMessageContaining("table default: 8"); + } + + // Writing an in-range bucket (< the partition's 2 buckets) into an empty bucket of the same + // partition is accepted: per-partition bucket counts are still honored. + int emptyBucket = findEmptyBucket(rescaledTable, 1, /* totalBuckets */ 2); + try (TableWriteImpl write = rescaledTable.newWrite(UUID.randomUUID().toString())) { + assertThatThrownBy(() -> write.writeAndReturn(GenericRow.of(1, 2, 2L), emptyBucket, 8)) + .hasMessageContaining("new bucket num 8") + .hasMessageContaining("previous bucket num is 2"); + } + + String user = UUID.randomUUID().toString(); + long id = rescaledTable.snapshotManager().latestSnapshotId(); + try (InnerTableWrite write = rescaledTable.newWrite(user); + StreamTableCommit commit = rescaledTable.newCommit(user)) { + write.write(GenericRow.of(1, 2, 2L), emptyBucket); + commit.commit(id, write.prepareCommit(true, id)); + } + } + + @Test + public void testNonEmptyBucketReportsManifestTotalBuckets() throws Exception { + // Sanity test: when a bucket has files, totalBuckets must come from the + // manifest entries (not from the fallback path). This guards against + // accidentally always overriding totalBuckets via PartitionBucketMapping. + FileStoreTable table = createPartitionedPkTable(2); + commitOneRow(table, 1, 1); + commitOneRow(table, 1, 2); + + // Locate a non-empty bucket within partition 1. + int nonEmptyBucket = findNonEmptyBucket(table, 1, 2); + + // Change the table default to ensure the returned totalBuckets is from the + // manifest entry, not the schema. + table = withBucket(table, 32); + + WriteRestore restore = newWriteRestore(table); + RestoreFiles restored = + restore.restoreFiles(binaryRow(1), nonEmptyBucket, false, false, false); + + assertThat(restored.totalBuckets()).isEqualTo(2); + assertThat(restored.dataFiles()).isNotEmpty(); + } + + @Test + public void testWriteRejectsBucketMismatchWhenPerPartitionCountDisabledByDefault() + throws Exception { + FileStoreTable table = createPartitionedPkTableWithDefaultOptions(4); + commitOneRow(table, 1, 1); + commitOneRow(table, 1, 2); + + int nonEmptyBucket = findNonEmptyBucket(table, 1, 4); + + FileStoreTable rescaledTable = withBucket(table, 2); + + try (InnerTableWrite write = rescaledTable.newWrite(UUID.randomUUID().toString())) { + assertThatThrownBy(() -> write.write(GenericRow.of(1, 1, 1L), nonEmptyBucket)) + .hasMessageContaining("a new bucket num 2") + .hasMessageContaining("the previous bucket num is 4"); + } + } + + // ------------------------------------------------------------------------ + // helpers + // ------------------------------------------------------------------------ + + private FileStoreTable createPartitionedPkTable(int bucket) throws Exception { + return createPartitionedPkTable(bucket, true); + } + + private FileStoreTable createPartitionedPkTable(int bucket, boolean perPartitionCountEnabled) + throws Exception { + return createPartitionedPkTable(bucket, Boolean.valueOf(perPartitionCountEnabled)); + } + + private FileStoreTable createPartitionedPkTableWithDefaultOptions(int bucket) throws Exception { + return createPartitionedPkTable(bucket, null); + } + + private FileStoreTable createPartitionedPkTable(int bucket, Boolean perPartitionCountEnabled) + throws Exception { + Path path = new Path(tempDir.toString()); + Options options = new Options(); + options.set(CoreOptions.PATH, path.toString()); + options.set(CoreOptions.BUCKET, bucket); + if (perPartitionCountEnabled != null) { + options.set(CoreOptions.BUCKET_PER_PARTITION_COUNT_ENABLED, perPartitionCountEnabled); + } + + TableSchema tableSchema = + SchemaUtils.forceCommit( + new FileSystemSchemaManager(LocalFileIO.create(), path), + new Schema( + ROW_TYPE.getFields(), + Collections.singletonList("pt"), + Arrays.asList("pt", "k"), + options.toMap(), + "")); + + return FileStoreTableFactory.create( + LocalFileIO.create(), path, tableSchema, CatalogEnvironment.empty()); + } + + private FileStoreTable withBucket(FileStoreTable table, int newBucket) { + Options options = new Options(table.options()); + options.set(CoreOptions.BUCKET, newBucket); + return table.copy(table.schema().copy(options.toMap())); + } + + private WriteRestore newWriteRestore(FileStoreTable table) { + return new FileSystemWriteRestore( + table.store().options(), + table.snapshotManager(), + table.store().newScan(), + table.store().newIndexFileHandler()); + } + + private void commitOneRow(FileStoreTable table, int pt, int k) throws Exception { + String user = UUID.randomUUID().toString(); + Long latest = table.snapshotManager().latestSnapshotId(); + long id = latest == null ? 0L : latest; + try (StreamTableWrite write = table.newWrite(user); + StreamTableCommit commit = table.newCommit(user)) { + write.write(GenericRow.of(pt, k, (long) k)); + commit.commit(id, write.prepareCommit(true, id)); + } + } + + /** Returns a bucket id (0..totalBuckets-1) that has no data files within the partition. */ + private int findEmptyBucket(FileStoreTable table, int pt, int totalBuckets) throws Exception { + BinaryRow partition = binaryRow(pt); + for (int b = 0; b < totalBuckets; b++) { + int bucket = b; + boolean nonEmpty = + table.newSnapshotReader() + .withPartitionFilter(Collections.singletonList(partition)) + .withBucket(bucket).read().dataSplits().stream() + .anyMatch(s -> !s.dataFiles().isEmpty()); + if (!nonEmpty) { + return bucket; + } + } + throw new IllegalStateException( + "Could not find an empty bucket in partition " + + pt + + " (every bucket has files); test scenario could not be set up."); + } + + /** Returns a bucket id (0..totalBuckets-1) that has at least one data file. */ + private int findNonEmptyBucket(FileStoreTable table, int pt, int totalBuckets) + throws Exception { + BinaryRow partition = binaryRow(pt); + for (int b = 0; b < totalBuckets; b++) { + int bucket = b; + boolean nonEmpty = + table.newSnapshotReader() + .withPartitionFilter(Collections.singletonList(partition)) + .withBucket(bucket).read().dataSplits().stream() + .anyMatch(s -> !s.dataFiles().isEmpty()); + if (nonEmpty) { + return bucket; + } + } + throw new IllegalStateException("Could not find a non-empty bucket in partition " + pt); + } + + private static BinaryRow binaryRow(int pt) { + BinaryRow row = new BinaryRow(1); + BinaryRowWriter writer = new BinaryRowWriter(row); + writer.writeInt(0, pt); + writer.complete(); + return row; + } } diff --git a/paimon-core/src/test/java/org/apache/paimon/table/AppendOnlySimpleTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/AppendOnlySimpleTableTest.java index e6352b8d91fb..43e1e6783ada 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/AppendOnlySimpleTableTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/AppendOnlySimpleTableTest.java @@ -158,16 +158,16 @@ public void testOverwriteSameFiles() throws Exception { } @Test - public void testBucketedAppendTableWriteWithInit() throws Exception { - innerTestBucketedAppendTableWriteInit(true); + public void testBucketedAppendOrderedSequenceNumbers() throws Exception { + innerTestBucketedAppendSequenceNumbers(true); } @Test - public void testBucketedAppendTableWriteNoInit() throws Exception { - innerTestBucketedAppendTableWriteInit(false); + public void testBucketedAppendUnorderedSequenceNumbers() throws Exception { + innerTestBucketedAppendSequenceNumbers(false); } - public void innerTestBucketedAppendTableWriteInit(boolean ordered) throws Exception { + public void innerTestBucketedAppendSequenceNumbers(boolean ordered) throws Exception { FileStoreTable table = createFileStoreTable( options -> { @@ -179,32 +179,47 @@ public void innerTestBucketedAppendTableWriteInit(boolean ordered) throws Except BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder(); - // 1. first write + // 1. first write - use a=1 so both batches land in the same bucket try (BatchTableWrite write = writeBuilder.newWrite(); BatchTableCommit commit = writeBuilder.newCommit()) { write.write(rowData(1, 10, 100L)); commit.commit(write.prepareCommit()); } - // 2. delete all manifests - ManifestList manifestList = table.store().manifestListFactory().create(); - ManifestFile manifestFile = table.store().manifestFileFactory().create(); - List manifests = - manifestList.readAllManifests(table.latestSnapshot().get()); - for (ManifestFileMeta manifest : manifests) { - manifestFile.delete(manifest.fileName()); + // collect sequence numbers from batch 1 + List batch1Files = + table.newReadBuilder().newScan().plan().splits().stream() + .flatMap(s -> ((DataSplit) s).dataFiles().stream()) + .collect(Collectors.toList()); + long batch1MaxSeq = + batch1Files.stream().mapToLong(DataFileMeta::maxSequenceNumber).max().getAsLong(); + Set batch1FileNames = + batch1Files.stream().map(DataFileMeta::fileName).collect(Collectors.toSet()); + + // 2. second write - same a=1 value ensures same bucket as batch 1 + try (BatchTableWrite write = writeBuilder.newWrite(); + BatchTableCommit commit = writeBuilder.newCommit()) { + write.write(rowData(1, 20, 200L)); + commit.commit(write.prepareCommit()); } - // 3. check new write - try (BatchTableWrite write = writeBuilder.newWrite()) { - if (ordered) { - assertThatThrownBy(() -> write.write(rowData(1, 10, 100L))) - .hasMessageContaining("Failed to restore existing files") - .hasRootCauseInstanceOf(java.io.FileNotFoundException.class); - } else { - // no exception - write.write(rowData(1, 10, 100L)); - } + // collect sequence numbers from batch 2 only (exclude batch 1 files by name) + List batch2Files = + table.newReadBuilder().newScan().plan().splits().stream() + .flatMap(s -> ((DataSplit) s).dataFiles().stream()) + .filter(f -> !batch1FileNames.contains(f.fileName())) + .collect(Collectors.toList()); + long batch2MinSeq = + batch2Files.stream().mapToLong(DataFileMeta::minSequenceNumber).min().getAsLong(); + + if (ordered) { + // ordered mode always restores previous files and continues sequence numbers, + // so batch 2 sequence numbers are strictly greater than batch 1's + assertThat(batch2MinSeq).isGreaterThan(batch1MaxSeq); + } else { + // unordered+writeOnly mode skips restoring previous files (ignorePreviousFiles=true), + // so sequence numbers reset to 0 each session + assertThat(batch2MinSeq).isEqualTo(0L); } } diff --git a/paimon-core/src/test/java/org/apache/paimon/table/sink/FixedBucketRowKeyExtractorTest.java b/paimon-core/src/test/java/org/apache/paimon/table/sink/FixedBucketRowKeyExtractorTest.java index 5c551dac7004..48d29d20ebb6 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/sink/FixedBucketRowKeyExtractorTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/sink/FixedBucketRowKeyExtractorTest.java @@ -101,6 +101,30 @@ public void testUnCompactDecimalAndTimestampNullValueBucketNumber() { } } + @Test + public void testPerPartitionBucketCount() { + int defaultBuckets = 100; + int partition1Buckets = 4; + + // Build a BinaryRow for partition value = 1 + BinaryRow partitionRow = BinaryRow.singleColumn(1); + + Map partitionMap = new HashMap<>(); + partitionMap.put(partitionRow, partition1Buckets); + PartitionBucketMapping mapping = new PartitionBucketMapping(defaultBuckets, partitionMap); + + // Schema: partition key "a", bucket key "b", primary key "a,b" + FixedBucketRowKeyExtractor extractor = extractor("a", "b", "a,b", defaultBuckets, mapping); + + // Same bucket key (b=456) in both partitions, different bucket counts produce + // different bucket assignments: hash(456) % 4 = 3, hash(456) % 100 = 47 + GenericRow rowInMappedPartition = GenericRow.of(1, 456, 7); + assertThat(bucket(extractor, rowInMappedPartition)).isEqualTo(3); + + GenericRow rowInDefaultPartition = GenericRow.of(99, 456, 7); + assertThat(bucket(extractor, rowInDefaultPartition)).isEqualTo(47); + } + private int bucket(FixedBucketRowKeyExtractor extractor, InternalRow row) { extractor.setRecord(row); return extractor.bucket(); @@ -125,8 +149,35 @@ private FixedBucketRowKeyExtractor extractor( return extractor(rowType, partK, bk, pk, numBucket); } + private FixedBucketRowKeyExtractor extractor( + String partK, String bk, String pk, int numBucket, PartitionBucketMapping mapping) { + RowType rowType = + new RowType( + Arrays.asList( + new DataField(0, "a", new IntType()), + new DataField(1, "b", new IntType()), + new DataField(2, "c", new IntType()))); + return extractor(rowType, partK, bk, pk, numBucket, mapping); + } + private FixedBucketRowKeyExtractor extractor( RowType rowType, String partK, String bk, String pk, int numBucket) { + return extractor( + rowType, + partK, + bk, + pk, + numBucket, + PartitionBucketMapping.defaultBuckets(numBucket)); + } + + private FixedBucketRowKeyExtractor extractor( + RowType rowType, + String partK, + String bk, + String pk, + int numBucket, + PartitionBucketMapping mapping) { List fields = TableSchema.newFields(rowType); Map options = new HashMap<>(); options.put(BUCKET_KEY.key(), bk); @@ -142,6 +193,6 @@ private FixedBucketRowKeyExtractor extractor( "".equals(pk) ? Collections.emptyList() : Arrays.asList(pk.split(",")), options, ""); - return new FixedBucketRowKeyExtractor(schema); + return new FixedBucketRowKeyExtractor(schema, mapping); } } diff --git a/paimon-core/src/test/java/org/apache/paimon/table/sink/FixedBucketWriteSelectorTest.java b/paimon-core/src/test/java/org/apache/paimon/table/sink/FixedBucketWriteSelectorTest.java new file mode 100644 index 000000000000..da524a5f59fc --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/sink/FixedBucketWriteSelectorTest.java @@ -0,0 +1,86 @@ +/* + * 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.paimon.table.sink; + +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.types.DataField; +import org.apache.paimon.types.IntType; +import org.apache.paimon.types.RowType; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.HashMap; +import java.util.Map; + +import static org.apache.paimon.CoreOptions.BUCKET; +import static org.apache.paimon.CoreOptions.BUCKET_KEY; +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link FixedBucketWriteSelector}. */ +public class FixedBucketWriteSelectorTest { + + @Test + public void testLegacyConstructorUsesTableBucketCount() { + TableSchema schema = schema(100); + GenericRow row = GenericRow.of(1, 456, 7); + + int actual = new FixedBucketWriteSelector(schema).select(row, 128); + int expected = ChannelComputer.select(BinaryRow.singleColumn(1), 47, 128); + + assertThat(actual).isEqualTo(expected); + } + + @Test + public void testPerPartitionBucketCount() { + TableSchema schema = schema(100); + BinaryRow partition = BinaryRow.singleColumn(1); + Map partitionBuckets = new HashMap<>(); + partitionBuckets.put(partition, 4); + PartitionBucketMapping mapping = new PartitionBucketMapping(100, partitionBuckets); + GenericRow row = GenericRow.of(1, 456, 7); + + int actual = new FixedBucketWriteSelector(schema, mapping).select(row, 128); + int expected = ChannelComputer.select(partition, 3, 128); + + assertThat(actual).isEqualTo(expected); + } + + private TableSchema schema(int numBuckets) { + RowType rowType = + new RowType( + Arrays.asList( + new DataField(0, "a", new IntType()), + new DataField(1, "b", new IntType()), + new DataField(2, "c", new IntType()))); + Map options = new HashMap<>(); + options.put(BUCKET_KEY.key(), "b"); + options.put(BUCKET.key(), String.valueOf(numBuckets)); + return new TableSchema( + 0, + TableSchema.newFields(rowType), + RowType.currentHighestFieldId(TableSchema.newFields(rowType)), + Arrays.asList("a"), + Arrays.asList("a", "b"), + options, + ""); + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/table/sink/PartitionBucketMappingTest.java b/paimon-core/src/test/java/org/apache/paimon/table/sink/PartitionBucketMappingTest.java new file mode 100644 index 000000000000..e272544e335e --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/sink/PartitionBucketMappingTest.java @@ -0,0 +1,87 @@ +/* + * 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.paimon.table.sink; + +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.operation.FileStoreScan; + +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.util.HashMap; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for {@link PartitionBucketMapping}. */ +public class PartitionBucketMappingTest { + + @Test + public void testDefaultBucketCount() { + PartitionBucketMapping mapping = PartitionBucketMapping.defaultBuckets(16); + + // Any partition should resolve to the default, but no partition has an explicit override. + assertThat(mapping.resolveNumBuckets(BinaryRow.EMPTY_ROW)).isEqualTo(16); + assertThat(mapping.resolveNumBuckets(partition(1))).isEqualTo(16); + assertThat(mapping.resolveNumBuckets(partition(42))).isEqualTo(16); + assertThat(mapping.getNumBucketsOverride(BinaryRow.EMPTY_ROW)).isNull(); + assertThat(mapping.getNumBucketsOverride(partition(1))).isNull(); + } + + @Test + public void testExplicitPartitionMapping() { + BinaryRow partA = partition(1); + BinaryRow partB = partition(2); + BinaryRow partC = partition(3); + + Map partitionMap = new HashMap<>(); + partitionMap.put(partA, 32); + partitionMap.put(partB, 64); + + PartitionBucketMapping mapping = new PartitionBucketMapping(16, partitionMap); + + // Mapped partitions return their specific bucket counts and expose explicit overrides. + assertThat(mapping.resolveNumBuckets(partA)).isEqualTo(32); + assertThat(mapping.resolveNumBuckets(partB)).isEqualTo(64); + assertThat(mapping.getNumBucketsOverride(partA)).isEqualTo(32); + assertThat(mapping.getNumBucketsOverride(partB)).isEqualTo(64); + + // Unmapped partition falls back to the default without reporting an explicit override. + assertThat(mapping.resolveNumBuckets(partC)).isEqualTo(16); + assertThat(mapping.getNumBucketsOverride(partC)).isNull(); + } + + @Test + public void testLoadFromScanPropagatesException() { + // Simulate a scan that throws (e.g. corrupted manifest, transient I/O error). + // loadFromScan must fail fast so the job does not silently write to wrong buckets. + FileStoreScan failingScan = Mockito.mock(FileStoreScan.class); + Mockito.when(failingScan.readPartitionEntries()) + .thenThrow(new RuntimeException("simulated manifest scan failure")); + + assertThatThrownBy(() -> PartitionBucketMapping.loadFromScan(failingScan, 8)) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("simulated manifest scan failure"); + } + + private static BinaryRow partition(int value) { + return BinaryRow.singleColumn(value); + } +} From 838f2dc4ad7906b390aa1a270c23168ceed5f613 Mon Sep 17 00:00:00 2001 From: David Wang Date: Fri, 25 Sep 2026 11:37:41 +1000 Subject: [PATCH 2/5] [core] Reject legacy writes for partition bucket layouts --- .../operation/AbstractFileStoreWrite.java | 13 +++++++ .../operation/FileSystemWriteRestoreTest.java | 36 ++++++++++--------- 2 files changed, 33 insertions(+), 16 deletions(-) diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java b/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java index d17490181b5e..6b7ec86840ad 100644 --- a/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java +++ b/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java @@ -501,6 +501,7 @@ public Map> getActiveBuckets() { } protected WriterContainer getWriterWrapper(BinaryRow partition, int bucket) { + requirePartitionBucketCount(partition); Map> buckets = getWriterContainers(partition); return buckets.computeIfAbsent( bucket, k -> createWriterContainer(partition.copy(), bucket)); @@ -531,9 +532,21 @@ public RecordWriter createWriter(BinaryRow partition, int bucket) { } public WriterContainer createWriterContainer(BinaryRow partition, int bucket) { + requirePartitionBucketCount(partition); return createWriterContainer(partition, bucket, numBuckets, !ignoreNumBucketCheck, false); } + private void requirePartitionBucketCount(BinaryRow partition) { + if (partitionType.getFieldCount() > 0 && options.bucketPerPartitionCountEnabled()) { + throw new UnsupportedOperationException( + "Writing partition " + + partition + + " with per-partition bucket counts requires the partition-level " + + "total bucket count. Use write(partition, bucket, totalBuckets, data) " + + "instead."); + } + } + private WriterContainer createWriterContainer( BinaryRow partition, int bucket, int totalBuckets) { return createWriterContainer(partition, bucket, totalBuckets, true, true); diff --git a/paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java b/paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java index 767a6f114b8c..4750dd71fe23 100644 --- a/paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java @@ -211,27 +211,23 @@ public void testEmptyBucketInUnseenPartitionDoesNotReportTableDefault() throws E @Test public void testWriteRejectsBucketOutsidePartitionLayout() throws Exception { // Partition 1 is created with 2 buckets. - FileStoreTable table = createPartitionedPkTable(2); + FileStoreTable table = createPartitionedPkTable(2, false); commitOneRow(table, /* pt */ 1, /* k */ 1); // Simulate a rescale: the table default is raised to 8 buckets, but partition 1 - // still only has 2 buckets. writeOnly=false so the writer scans previous files and - // runs the per-partition bucket-layout check in AbstractFileStoreWrite. - FileStoreTable rescaledTable = withBucket(table, 8); - - // Writing an out-of-range bucket (>= the partition's 2 buckets) into an empty bucket of - // partition 1 must be rejected, even though the bucket id is valid for the 8-bucket - // default. - // This is the bucket that PartitionBucketMapping recovery would otherwise silently accept. - // write(row, bucket) routes the row (partition pt=1) to the explicitly given bucket. + // still only has 2 buckets. Enable per-partition bucket counts after the existing layout + // is in place. + FileStoreTable rescaledTable = withBucket(table, 8, true); + + // The legacy two-argument write does not carry the bucket count used to route the row and + // must be rejected for a per-partition bucket table. try (InnerTableWrite write = rescaledTable.newWrite(UUID.randomUUID().toString())) { assertThatThrownBy(() -> write.write(GenericRow.of(1, 1, 1L), /* bucket */ 6)) - .hasMessageContaining("only has 2 buckets") - .hasMessageContaining("table default: 8"); + .hasMessageContaining("requires the partition-level total bucket count"); } - // Writing an in-range bucket (< the partition's 2 buckets) into an empty bucket of the same - // partition is accepted: per-partition bucket counts are still honored. + // A caller which supplies a stale layout is still rejected, even when the bucket happens + // to be in range for the old partition layout. int emptyBucket = findEmptyBucket(rescaledTable, 1, /* totalBuckets */ 2); try (TableWriteImpl write = rescaledTable.newWrite(UUID.randomUUID().toString())) { assertThatThrownBy(() -> write.writeAndReturn(GenericRow.of(1, 2, 2L), emptyBucket, 8)) @@ -241,9 +237,9 @@ public void testWriteRejectsBucketOutsidePartitionLayout() throws Exception { String user = UUID.randomUUID().toString(); long id = rescaledTable.snapshotManager().latestSnapshotId(); - try (InnerTableWrite write = rescaledTable.newWrite(user); + try (TableWriteImpl write = rescaledTable.newWrite(user); StreamTableCommit commit = rescaledTable.newCommit(user)) { - write.write(GenericRow.of(1, 2, 2L), emptyBucket); + write.writeAndReturn(GenericRow.of(1, 2, 2L), emptyBucket, 2); commit.commit(id, write.prepareCommit(true, id)); } } @@ -332,8 +328,16 @@ private FileStoreTable createPartitionedPkTable(int bucket, Boolean perPartition } private FileStoreTable withBucket(FileStoreTable table, int newBucket) { + return withBucket(table, newBucket, null); + } + + private FileStoreTable withBucket( + FileStoreTable table, int newBucket, Boolean perPartitionCountEnabled) { Options options = new Options(table.options()); options.set(CoreOptions.BUCKET, newBucket); + if (perPartitionCountEnabled != null) { + options.set(CoreOptions.BUCKET_PER_PARTITION_COUNT_ENABLED, perPartitionCountEnabled); + } return table.copy(table.schema().copy(options.toMap())); } From 655cf5e29da25e2a0cb959450774143a7b567152 Mon Sep 17 00:00:00 2001 From: David Wang Date: Fri, 25 Sep 2026 11:53:24 +1000 Subject: [PATCH 3/5] [core] Remove duplicate append test coverage --- .../operation/FileSystemWriteRestoreTest.java | 12 ++-- .../table/AppendOnlySimpleTableTest.java | 66 ------------------- 2 files changed, 7 insertions(+), 71 deletions(-) diff --git a/paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java b/paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java index 4750dd71fe23..9cbe94e5f1a4 100644 --- a/paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java @@ -163,7 +163,7 @@ public void testEmptyBucketUsesPartitionBucketMapping() throws Exception { // and ask the WriteRestore for an empty bucket. It must return // totalBuckets=4 (the partition's actual bucket count), NOT 32 (the new // table default). - FileStoreTable table = createPartitionedPkTable(4); + FileStoreTable table = createPartitionedPkTable(4, false); // Write enough rows to populate at least one bucket within partition 1. commitOneRow(table, /* pt */ 1, /* k */ 1); @@ -175,7 +175,7 @@ public void testEmptyBucketUsesPartitionBucketMapping() throws Exception { // Simulate a rescale by raising the table-level default bucket count // (without rewriting existing files). Existing manifest entries still // carry totalBuckets=4. - table = withBucket(table, 32); + table = withBucket(table, 32, true); WriteRestore restore = newWriteRestore(table); @@ -197,9 +197,11 @@ public void testEmptyBucketInUnseenPartitionDoesNotReportTableDefault() throws E // Return null so the writer falls back to its expected table-level bucket count. In // particular, this preserves postpone-bucket writes whose table default is -2 while the // staged writer assigns real fixed buckets. - FileStoreTable table = createPartitionedPkTable(8); + FileStoreTable table = createPartitionedPkTable(8, false); commitOneRow(table, 1, 100); // ensures the snapshot exists + table = withBucket(table, 8, true); + WriteRestore restore = newWriteRestore(table); RestoreFiles restored = restore.restoreFiles(binaryRow(/* unseen */ 999), 0, false, false, false); @@ -249,7 +251,7 @@ public void testNonEmptyBucketReportsManifestTotalBuckets() throws Exception { // Sanity test: when a bucket has files, totalBuckets must come from the // manifest entries (not from the fallback path). This guards against // accidentally always overriding totalBuckets via PartitionBucketMapping. - FileStoreTable table = createPartitionedPkTable(2); + FileStoreTable table = createPartitionedPkTable(2, false); commitOneRow(table, 1, 1); commitOneRow(table, 1, 2); @@ -258,7 +260,7 @@ public void testNonEmptyBucketReportsManifestTotalBuckets() throws Exception { // Change the table default to ensure the returned totalBuckets is from the // manifest entry, not the schema. - table = withBucket(table, 32); + table = withBucket(table, 32, true); WriteRestore restore = newWriteRestore(table); RestoreFiles restored = diff --git a/paimon-core/src/test/java/org/apache/paimon/table/AppendOnlySimpleTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/AppendOnlySimpleTableTest.java index 43e1e6783ada..56d249ef4bf2 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/AppendOnlySimpleTableTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/AppendOnlySimpleTableTest.java @@ -167,72 +167,6 @@ public void testBucketedAppendUnorderedSequenceNumbers() throws Exception { innerTestBucketedAppendSequenceNumbers(false); } - public void innerTestBucketedAppendSequenceNumbers(boolean ordered) throws Exception { - FileStoreTable table = - createFileStoreTable( - options -> { - options.set(BUCKET, 2); - options.set(BUCKET_KEY, "a"); - options.set(WRITE_ONLY, true); - options.set(BUCKET_APPEND_ORDERED, ordered); - }); - - BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder(); - - // 1. first write - use a=1 so both batches land in the same bucket - try (BatchTableWrite write = writeBuilder.newWrite(); - BatchTableCommit commit = writeBuilder.newCommit()) { - write.write(rowData(1, 10, 100L)); - commit.commit(write.prepareCommit()); - } - - // collect sequence numbers from batch 1 - List batch1Files = - table.newReadBuilder().newScan().plan().splits().stream() - .flatMap(s -> ((DataSplit) s).dataFiles().stream()) - .collect(Collectors.toList()); - long batch1MaxSeq = - batch1Files.stream().mapToLong(DataFileMeta::maxSequenceNumber).max().getAsLong(); - Set batch1FileNames = - batch1Files.stream().map(DataFileMeta::fileName).collect(Collectors.toSet()); - - // 2. second write - same a=1 value ensures same bucket as batch 1 - try (BatchTableWrite write = writeBuilder.newWrite(); - BatchTableCommit commit = writeBuilder.newCommit()) { - write.write(rowData(1, 20, 200L)); - commit.commit(write.prepareCommit()); - } - - // collect sequence numbers from batch 2 only (exclude batch 1 files by name) - List batch2Files = - table.newReadBuilder().newScan().plan().splits().stream() - .flatMap(s -> ((DataSplit) s).dataFiles().stream()) - .filter(f -> !batch1FileNames.contains(f.fileName())) - .collect(Collectors.toList()); - long batch2MinSeq = - batch2Files.stream().mapToLong(DataFileMeta::minSequenceNumber).min().getAsLong(); - - if (ordered) { - // ordered mode always restores previous files and continues sequence numbers, - // so batch 2 sequence numbers are strictly greater than batch 1's - assertThat(batch2MinSeq).isGreaterThan(batch1MaxSeq); - } else { - // unordered+writeOnly mode skips restoring previous files (ignorePreviousFiles=true), - // so sequence numbers reset to 0 each session - assertThat(batch2MinSeq).isEqualTo(0L); - } - } - - @Test - public void testBucketedAppendOrderedSequenceNumbers() throws Exception { - innerTestBucketedAppendSequenceNumbers(true); - } - - @Test - public void testBucketedAppendUnorderedSequenceNumbers() throws Exception { - innerTestBucketedAppendSequenceNumbers(false); - } - private void innerTestBucketedAppendSequenceNumbers(boolean ordered) throws Exception { FileStoreTable table = createFileStoreTable( From dcb0695817b6042a445218c4ed6d6092c228ce70 Mon Sep 17 00:00:00 2001 From: David Wang Date: Fri, 25 Sep 2026 12:21:23 +1000 Subject: [PATCH 4/5] [core] Preserve WriteRestore extraction API --- .../apache/paimon/operation/WriteRestore.java | 28 ++++++++++++++----- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/WriteRestore.java b/paimon-core/src/main/java/org/apache/paimon/operation/WriteRestore.java index f2ed55137827..00a1f1f205ff 100644 --- a/paimon-core/src/main/java/org/apache/paimon/operation/WriteRestore.java +++ b/paimon-core/src/main/java/org/apache/paimon/operation/WriteRestore.java @@ -67,17 +67,16 @@ static Integer extractTotalBuckets( } /** - * Extracts the {@link DataFileMeta} list from the given manifest entries, validating that all - * entries agree on {@code totalBuckets}. + * Extracts data files into the supplied list and returns their common bucket count. * * @param entries manifest entries for a single (partition, bucket) pair - * @return the list of data files; empty if {@code entries} is empty - * @throws RuntimeException if entries carry inconsistent {@code totalBuckets} values, which - * indicates a corrupted manifest + * @param dataFiles destination for the extracted data files + * @return the common bucket count, or {@code null} when {@code entries} is empty + * @throws RuntimeException if entries carry inconsistent {@code totalBuckets} values */ - static List extractDataFiles(List entries) { + @Nullable + static Integer extractDataFiles(List entries, List dataFiles) { Integer totalBuckets = null; - List dataFiles = new ArrayList<>(); for (ManifestEntry entry : entries) { if (totalBuckets != null && totalBuckets != entry.totalBuckets()) { throw new RuntimeException( @@ -88,6 +87,21 @@ static List extractDataFiles(List entries) { totalBuckets = entry.totalBuckets(); dataFiles.add(entry.file()); } + return totalBuckets; + } + + /** + * Extracts the {@link DataFileMeta} list from the given manifest entries, validating that all + * entries agree on {@code totalBuckets}. + * + * @param entries manifest entries for a single (partition, bucket) pair + * @return the list of data files; empty if {@code entries} is empty + * @throws RuntimeException if entries carry inconsistent {@code totalBuckets} values, which + * indicates a corrupted manifest + */ + static List extractDataFiles(List entries) { + List dataFiles = new ArrayList<>(); + extractDataFiles(entries, dataFiles); return dataFiles; } } From 2a1a9b2336442da5e35d6a06c27618f2cf1df157 Mon Sep 17 00:00:00 2001 From: David Wang Date: Fri, 25 Sep 2026 14:47:56 +1000 Subject: [PATCH 5/5] [ci] Retry GitHub Actions