[core][flink] Support Per-Partition Bucket Counts

docs/docs/flink/procedures.md

@@ -969,7 +969,7 @@ All available procedures are listed below.
          CALL [catalog.]sys.rescale(`table` => 'identifier', `bucket_num` => bucket_num, `partition` => 'partition', `scan_parallelism` => scan_parallelism, `sink_parallelism` => sink_parallelism)
       </td>
       <td>
-         Rescale one partition of a table. Arguments:
+         Rescale one partition of a table. For partitioned tables, different partitions can have different bucket counts after rescaling, provided the table option 'bucket.per-partition-count-enabled' is set to 'true'. Arguments:
          <li>table: The target table identifier. Cannot be empty.</li>
          <li>bucket_num: Resulting bucket number after rescale. The default value of argument bucket_num is the current bucket number of the table. Cannot be empty for postpone bucket tables.</li>
          <li>partition: What partition to rescale. For partitioned table this argument cannot be empty.</li>

docs/docs/maintenance/rescale-bucket.md

@@ -45,15 +45,12 @@ Please note that
 - `ALTER TABLE` only modifies the table's metadata and will **NOT** reorganize or reformat existing data. 
   Reorganize existing data must be achieved by `INSERT OVERWRITE`.
 - Rescale bucket number does not influence the read and running write jobs.
-- Once the bucket number is changed, any newly scheduled `INSERT INTO` jobs which write to without-reorganized 
-  existing table/partition will throw a `TableException` with message like 
-  ```text
-  Try to write table/partition ... with a new bucket num ..., 
-  but the previous bucket num is ... Please switch to batch mode, 
-  and perform INSERT OVERWRITE to rescale current data layout first.
-  ```
-- For partitioned table, it is possible to have different bucket number for different partitions. *E.g.*
+- For **partitioned tables**, it is possible to have different bucket numbers for different partitions.
+  This requires setting `'bucket.per-partition-count-enabled' = 'true'` on the table; otherwise every
+  partition uses the single table-level bucket count. *E.g.*
   ```sql
+  ALTER TABLE my_table SET ('bucket.per-partition-count-enabled' = 'true');
+
   ALTER TABLE my_table SET ('bucket' = '4');
   INSERT OVERWRITE my_table PARTITION (dt = '2022-01-01')
   SELECT * FROM ...;
@@ -62,7 +59,45 @@ Please note that
   INSERT OVERWRITE my_table PARTITION (dt = '2022-01-02')
   SELECT * FROM ...;
   ```
+  After these operations, partition `dt=2022-01-01` uses 4 buckets, `dt=2022-01-02` uses 8 buckets, and any
+  new partitions will use the latest table-level default (8 buckets in this case).
+  Each partition retains its own bucket count from its data files,
+  and the new bucket count only applies to newly created partitions or partitions that
+  have been reorganized with `INSERT OVERWRITE`.
+
+:::info
+  Per-partition bucket counts are disabled by default. Set `'bucket.per-partition-count-enabled' = 'true'`
+  to let partitions keep their own bucket count and be rescaled independently. When it is disabled, all
+  partitions share the table-level `bucket` value.
+
+  Note that enabling this option adds an extra manifest scan on write (to resolve each partition's bucket
+  count), so only enable it when you actually need different bucket counts across partitions.
+:::
+
+:::warning
+  Per-partition bucket counts are currently supported by the **Flink** engine only. The **Spark** writer
+  still derives the bucket from the single table-level bucket count, so when partitions have different
+  bucket counts (for example, after changing the table-level `bucket` while existing partitions keep
+  their previous count), Spark may route rows to buckets that do not belong to the partition and corrupt
+  the per-partition layout. Until Spark support is added, use Flink to write to tables that have
+  per-partition bucket counts, or perform a full-table rescale so every partition shares the same bucket
+  count before writing with Spark.
+:::
+- **Unpartitioned tables** require a full rescale before writing. If you change the bucket number and attempt
+  to write without reorganizing the data first, a `RuntimeException` will be thrown:
+  ```text
+  Try to write table/partition ... with a new bucket num ...,
+  but the previous bucket num is ... Please switch to batch mode,
+  and perform INSERT OVERWRITE to rescale current data layout first.
+  ```
 - During overwrite period, make sure there are no other jobs writing the same table/partition.
+- **Streaming jobs must be restarted after rescaling a partition.** The per-partition bucket mapping
+  is loaded once when the streaming job starts (from the manifest files at that point in time). If a
+  partition is rescaled while the streaming job is running, the job will continue routing rows using
+  the old bucket count for that partition, which can cause rows to land in wrong buckets and lead to
+  data correctness issues. The recommended workflow is: suspend the streaming job with a savepoint →
+  perform the rescale overwrite → restart from the savepoint.
+
 
 ## Use Case
 
@@ -106,10 +141,10 @@ SELECT trade_order_id,
 FROM raw_orders
 WHERE order_status = 'verified';
 ```
-The pipeline has been running well for the past few weeks. However, the data volume has grown fast recently, 
-and the job's latency keeps increasing. To improve the data freshness, users can 
-- Suspend the streaming job with a savepoint ( see 
-  [Suspended State](https://nightlies.apache.org/flink/flink-docs-stable/docs/internals/job_scheduling/) and 
+The pipeline has been running well for the past few weeks. However, the data volume has grown fast recently,
+and the job's latency keeps increasing. To improve the data freshness, users can
+- Suspend the streaming job with a savepoint ( see
+  [Suspended State](https://nightlies.apache.org/flink/flink-docs-stable/docs/internals/job_scheduling/) and
   [Stopping a Job Gracefully Creating a Final Savepoint](https://nightlies.apache.org/flink/flink-docs-stable/docs/deployment/cli/#terminating-a-job) )
   ```bash
   $ ./bin/flink stop \
@@ -142,8 +177,8 @@ and the job's latency keeps increasing. To improve the data freshness, users can
   FROM verified_orders
   WHERE dt IN ('2022-06-20', '2022-06-21', '2022-06-22');
   ```
-- After overwrite job has finished, switch back to streaming mode. And now, the parallelism can be increased alongside with bucket number to restore the streaming job from the savepoint 
-( see [Start a SQL Job from a savepoint](https://nightlies.apache.org/flink/flink-docs-stable/docs/dev/table/sqlclient/#start-a-sql-job-from-a-savepoint) )
+- After overwrite job has finished, switch back to streaming mode. And now, the parallelism can be increased alongside with bucket number to restore the streaming job from the savepoint
+  ( see [Start a SQL Job from a savepoint](https://nightlies.apache.org/flink/flink-docs-stable/docs/dev/table/sqlclient/#start-a-sql-job-from-a-savepoint) )
   ```sql
   SET 'execution.runtime-mode' = 'streaming';
   SET 'execution.savepoint.path' = <savepointPath>;
@@ -155,4 +190,4 @@ and the job's latency keeps increasing. To improve the data freshness, users can
        DATE_FORMAT(gmt_create, 'yyyy-MM-dd') AS dt
   FROM raw_orders
   WHERE order_status = 'verified';
-  ```
+  ```

docs/docs/primary-key-table/data-distribution.md

@@ -34,6 +34,11 @@ the bucket of record.
 Rescaling buckets can only be done through offline processes, see [Rescale Bucket](../maintenance/rescale-bucket).
 A too large number of buckets leads to too many small files, and a too small number of buckets leads to poor write performance.
 
+For partitioned tables, each partition can have its own bucket count when
+`'bucket.per-partition-count-enabled' = 'true'` is set. In that case, after a rescale operation existing
+partitions retain their original bucket count while newly created partitions use the updated table-level
+default. When the option is disabled (the default), all partitions share the single table-level bucket count.
+
 ## Dynamic Bucket
 
 Default mode for primary key table, or configure `'bucket' = '-1'`.

docs/generated/core_configuration.html

@@ -152,6 +152,12 @@
             <td>String</td>
             <td>Specify the paimon distribution policy. Data is assigned to each bucket according to the hash value of bucket-key.<br />If you specify multiple fields, delimiter is ','.<br />If not specified, the primary key will be used; if there is no primary key, the full row will be used.</td>
         </tr>
+        <tr>
+            <td><h5>bucket.per-partition-count-enabled</h5></td>
+            <td style="word-wrap: break-word;">false</td>
+            <td>Boolean</td>
+            <td>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.</td>
+        </tr>
         <tr>
             <td><h5>cache-page-size</h5></td>
             <td style="word-wrap: break-word;">64 kb</td>

paimon-api/src/main/java/org/apache/paimon/CoreOptions.java

@@ -155,6 +155,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")
@@ -3003,6 +3013,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/manifest/PartitionEntry.java

@@ -83,13 +83,27 @@ public int totalBuckets() {
     }
 
     public PartitionEntry merge(PartitionEntry entry) {
+        PartitionEntry newer = entry.lastFileCreationTime >= lastFileCreationTime ? entry : this;
+        PartitionEntry older = newer == entry ? this : entry;
+
+        // Use the totalBuckets from the most recently created file. This correctly handles
+        // the case where a partition has been overwritten with a different bucket count: the
+        // newer files carry the new totalBuckets, and their creation time is always later.
+        // When timestamps are equal (e.g., two files written in the same millisecond with
+        // different bucket counts), we take the larger totalBuckets value. This makes merge
+        // commutative and associative — a.merge(b) == b.merge(a)
+        int newTotalBuckets =
+                newer.lastFileCreationTime == older.lastFileCreationTime
+                        ? Math.max(newer.totalBuckets, older.totalBuckets)
+                        : newer.totalBuckets;
+
         return new PartitionEntry(
                 partition,
                 recordCount + entry.recordCount,
                 fileSizeInBytes + entry.fileSizeInBytes,
                 fileCount + entry.fileCount,
-                Math.max(lastFileCreationTime, entry.lastFileCreationTime),
-                entry.totalBuckets);
+                newer.lastFileCreationTime,
+                newTotalBuckets);
     }
 
     public Partition toPartition(InternalRowPartitionComputer computer) {

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

@@ -674,8 +674,31 @@ 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()) {
+                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/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;
@@ -40,6 +40,7 @@ public class FileSystemWriteRestore implements WriteRestore {
     private final SnapshotManager snapshotManager;
     private final FileStoreScan scan;
     private final IndexFileHandler indexFileHandler;
+    private final PartitionBucketMapping partitionBucketMapping;
     private final @Nullable Long snapshotId;
 
     public FileSystemWriteRestore(
@@ -74,6 +75,10 @@ private FileSystemWriteRestore(
                 this.scan.dropStats();
             }
         }
+        this.partitionBucketMapping =
+                options.bucketPerPartitionCountEnabled()
+                        ? PartitionBucketMapping.loadFromScan(scan, options.bucket())
+                        : PartitionBucketMapping.defaultBuckets(options.bucket());
     }
 
     @Override
@@ -101,10 +106,12 @@ public RestoreFiles restoreFiles(
             return RestoreFiles.empty();
         }
 
-        List<DataFileMeta> restoreFiles = new ArrayList<>();
         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, partitionBucketMapping);
 
         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,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.
+     *
+     * <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 extractDataFiles(List<ManifestEntry> entries, List<DataFileMeta> dataFiles) {
+    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 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) {
         Integer totalBuckets = null;
+        List<DataFileMeta> dataFiles = new ArrayList<>();
         for (ManifestEntry entry : entries) {
             if (totalBuckets != null && totalBuckets != entry.totalBuckets()) {
                 throw new RuntimeException(
@@ -51,6 +88,6 @@ static Integer extractDataFiles(List<ManifestEntry> entries, List<DataFileMeta>
             totalBuckets = entry.totalBuckets();
             dataFiles.add(entry.file());
         }
-        return totalBuckets;
+        return dataFiles;
     }
 }

paimon-core/src/main/java/org/apache/paimon/table/AbstractFileStoreTable.java

@@ -40,6 +40,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;
@@ -228,7 +229,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();
@@ -256,7 +259,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());