[WIP] Split partition-level bucket-count support into focused PRs

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,40 @@ 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. Spark rejects writes
+  to a table with `'bucket.per-partition-count-enabled' = 'true'`; use Flink to write such a table.
+:::
+- **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 +136,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 +172,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>;

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

@@ -59,6 +59,11 @@ split one bucket into independent key ranges, so bucket count is not a hard read
 To change an existing layout, use the offline [Rescale Bucket](../maintenance/rescale-bucket)
 workflow.
 
+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
 
 Dynamic buckets are the default for primary-key tables (`bucket = -1`). Paimon maintains an index

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

@@ -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")
@@ -3155,6 +3165,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,47 @@ 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. A creation timestamp does not provide an
+        // ordering when the files were created in the same millisecond. In that case, choosing
+        // either bucket count (for example, the larger one) can silently discard a rescale.
+        if (newer.lastFileCreationTime == older.lastFileCreationTime
+                && newer.totalBuckets != older.totalBuckets) {
+            if (newer.totalBuckets > 0 && older.totalBuckets > 0) {
+                throw new IllegalStateException(
+                        String.format(
+                                "Cannot determine the bucket layout for partition %s: files created at %s "
+                                        + "have conflicting bucket counts (%s and %s).",
+                                partition,
+                                newer.lastFileCreationTime,
+                                newer.totalBuckets,
+                                older.totalBuckets));
+            }
+
+            // A non-positive bucket count is a mode marker, not a fixed-bucket layout. For
+            // example, dynamic bucket tables can contain -1 entries while their real buckets
+            // are being created. Preserve the previous deterministic tie-break for these modes.
+            return new PartitionEntry(
+                    partition,
+                    recordCount + entry.recordCount,
+                    fileSizeInBytes + entry.fileSizeInBytes,
+                    fileCount + entry.fileCount,
+                    newer.lastFileCreationTime,
+                    Math.max(newer.totalBuckets, older.totalBuckets));
+        }
+        int newTotalBuckets = 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

@@ -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;
@@ -503,7 +513,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 +531,20 @@ public RecordWriter<T> createWriter(BinaryRow partition, int bucket) {
     }
 
     public WriterContainer<T> createWriterContainer(BinaryRow partition, int bucket) {
-        return createWriterContainer(partition, bucket, numBuckets, !ignoreNumBucketCheck);
+        return createWriterContainer(partition, bucket, numBuckets, !ignoreNumBucketCheck, false);
     }
 
     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 +569,11 @@ private WriterContainer<T> createWriterContainer(
         if (!actualIgnorePreviousFiles) {
             restored =
                     scanExistingFileMetas(
-                            partition, bucket, expectedTotalBuckets, validateNumBuckets);
+                            partition,
+                            bucket,
+                            expectedTotalBuckets,
+                            validateNumBuckets,
+                            strictBucketCount);
         }
 
         DynamicBucketIndexMaintainer indexMaintainer =
@@ -647,7 +666,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 +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;
     }

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. */