[core] Route overwrite writes by partition bucket layout

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<Boolean> 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<BucketFunctionType> 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);
     }

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<T> withWriteRestore(WriteRestore writeRestore) {
         return this;
     }
 
+    @Override
+    public FileStoreWrite<T> withPartitionBucketMapping(
+            PartitionBucketMapping partitionBucketMapping) {
+        if (restore instanceof FileSystemWriteRestore) {
+            ((FileSystemWriteRestore) restore).withPartitionBucketMapping(partitionBucketMapping);
+        }
+        return this;
+    }
+
     @Override
     public FileStoreWrite<T> withIOManager(IOManager ioManager) {
         this.ioManager = ioManager;
@@ -491,6 +501,7 @@ public Map<BinaryRow, List<Integer>> getActiveBuckets() {
     }
 
     protected WriterContainer<T> getWriterWrapper(BinaryRow partition, int bucket) {
+        requirePartitionBucketCount(partition);
         Map<Integer, WriterContainer<T>> buckets = getWriterContainers(partition);
         return buckets.computeIfAbsent(
                 bucket, k -> createWriterContainer(partition.copy(), bucket));
@@ -503,7 +514,8 @@ private WriterContainer<T> 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<Integer, WriterContainer<T>> getWriterContainers(BinaryRow partition) {
@@ -520,16 +532,32 @@ public RecordWriter<T> createWriter(BinaryRow partition, int bucket) {
     }
 
     public WriterContainer<T> createWriterContainer(BinaryRow partition, int bucket) {
-        return createWriterContainer(partition, bucket, numBuckets, !ignoreNumBucketCheck);
+        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<T> createWriterContainer(
             BinaryRow partition, int bucket, int totalBuckets) {
-        return createWriterContainer(partition, bucket, totalBuckets, true);
+        return createWriterContainer(partition, bucket, totalBuckets, true, true);
     }
 
     private WriterContainer<T> 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 +582,11 @@ private WriterContainer<T> createWriterContainer(
         if (!actualIgnorePreviousFiles) {
             restored =
                     scanExistingFileMetas(
-                            partition, bucket, expectedTotalBuckets, validateNumBuckets);
+                            partition,
+                            bucket,
+                            expectedTotalBuckets,
+                            validateNumBuckets,
+                            strictBucketCount);
         }
 
         DynamicBucketIndexMaintainer indexMaintainer =
@@ -647,7 +679,11 @@ public FileStoreWrite<T> 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<String> partInfo =
                 () ->
                         partitionType.getFieldCount() > 0
@@ -674,8 +710,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;
     }

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<T> extends Restorable<List<FileStoreWrite.State<
 
     FileStoreWrite<T> withWriteRestore(WriteRestore writeRestore);
 
+    /** Provides the preloaded partition-to-bucket mapping for fixed-bucket writes. */
+    default FileStoreWrite<T> withPartitionBucketMapping(
+            PartitionBucketMapping partitionBucketMapping) {
+        return this;
+    }
+
     FileStoreWrite<T> withIOManager(IOManager ioManager);
 
     /** Specified the write rowType. */

paimon-core/src/main/java/org/apache/paimon/operation/FileSystemWriteRestore.java

@@ -25,21 +25,23 @@
 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;
 
 /** {@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<DataFileMeta> restoreFiles = new ArrayList<>();
+        // load the mapping before narrowing the mutable scan to a single bucket
+        PartitionBucketMapping bucketMapping = partitionBucketMapping();
         List<ManifestEntry> entries =
                 scan.withSnapshot(snapshot).withPartitionBucket(partition, bucket).plan().files();
-        Integer totalBuckets = WriteRestore.extractDataFiles(entries, restoreFiles);
+        List<DataFileMeta> restoreFiles = WriteRestore.extractDataFiles(entries);
+
+        Integer totalBuckets = WriteRestore.extractTotalBuckets(entries, partition, bucketMapping);
 
         IndexFileMeta dynamicBucketIndex = null;
         if (scanDynamicBucketIndex) {

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,6 +40,40 @@ 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.
+     *
+     * <ul>
+     *   <li>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.
+     *   <li>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.
+     *   <li>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.
+     * </ul>
+     */
+    @Nullable
+    static Integer extractTotalBuckets(
+            List<ManifestEntry> entries, BinaryRow partition, PartitionBucketMapping mapping) {
+        if (!entries.isEmpty()) {
+            return entries.get(0).totalBuckets();
+        }
+        if (partition.getFieldCount() > 0) {
+            return mapping.getNumBucketsOverride(partition);
+        }
+        return null;
+    }
+
+    /**
+     * Extracts data files into the supplied list and returns their common bucket count.
+     *
+     * @param entries manifest entries for a single (partition, bucket) pair
+     * @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
+     */
     @Nullable
     static Integer extractDataFiles(List<ManifestEntry> entries, List<DataFileMeta> dataFiles) {
         Integer totalBuckets = null;
@@ -53,4 +89,19 @@ static Integer extractDataFiles(List<ManifestEntry> entries, List<DataFileMeta>
         }
         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<DataFileMeta> extractDataFiles(List<ManifestEntry> entries) {
+        List<DataFileMeta> dataFiles = new ArrayList<>();
+        extractDataFiles(entries, dataFiles);
+        return dataFiles;
+    }
 }

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> statistics() {
     public Optional<WriteSelector> 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());

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<InternalRow> newWrite(String commitUser) {
 
     @Override
     public TableWriteImpl<InternalRow> newWrite(String commitUser, @Nullable Integer writeId) {
+        return newWrite(commitUser, writeId, createRowKeyExtractor());
+    }
+
+    @Override
+    public TableWriteImpl<InternalRow> 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(),