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

[mikedias]

Problem

In partitioned Paimon tables, all partitions share the same bucket count defined at the table level. This becomes a bottleneck when data is highly skewed: a "hot" partition (e.g., a large tenant) may receive orders of magnitude more data than other partitions, yet it is forced to use the same number of buckets. The only workaround was to increse the number of buckets for the entire table, but that in turn end up creating too many buckets for smaller partitions, leading to a small file problem.

Solution

This PR introduces per-partition bucket counts, allowing individual partitions to be independently rescaled. Skewed partitions can be split into more buckets without affecting the rest of the table.

The core idea is a new PartitionBucketMapping that maintains an explicit partition → bucket count map alongside a table-level default. Every component that needs to assign a bucket to a row (write selectors, key extractors) now consults this mapping rather than blindly using schema().numBuckets(). Each partition's bucket count is derived from the totalBuckets field already stamped on its data files in the manifest, so no schema migration is required.

Changes

Core (paimon-core)

• PartitionBucketMapping (new) — Serializable mapping of BinaryRow partition → int bucketCount, with a loadFromTable factory that scans the manifest to reconstruct the current per-partition layout and falls back to the schema default gracefully.
• SchemaBucketFileStoreTable (new) — A lightweight DelegatedFileStoreTable wrapper used during rescale/overwrite operations. It forces all writes to use the new target bucket count (ignoring the per-partition map), ensuring the overwrite lands in the right buckets.
• FixedBucketRowKeyExtractor / FixedBucketWriteSelector — Updated to accept a PartitionBucketMapping and call resolveNumBuckets(partition) per row instead of using a fixed global count.
• WriteRestore / FileSystemWriteRestore — Extended with extractTotalBuckets logic that correctly handles three cases: non-empty buckets (use the value from existing data files), empty buckets on partitioned tables (look up the per-partition override), and empty buckets on unpartitioned tables (fall back to schema default so the committer-side mismatch check still fires).
• PartitionEntry — Minor fix for correct behaviour in non-partitioned table corner cases.

Flink (paimon-flink)

• FlinkSinkBuilder — Wires PartitionBucketMapping into the streaming sink pipeline so that per-partition bucket routing is applied at ingest time.
• RescaleAction / CompactAction — Use RescaleFileStoreTable when performing rescale/overwrite so the new bucket count is applied only to the target partitions.
• RowDataChannelComputer — Updated to route rows to the correct sub-task using the per-partition bucket count.
• TableWriteCoordinator / PostponeFixedBucketChannelComputer — Fixed to handle the "empty bucket" scenario that can arise in write-restore flows when a partition exists in the mapping but has no files yet.
• RowDataKeyAndBucketExtractor (deleted) — Test helper class replaced with using the superclass types directly.

Behaviour

• Partitioned tables: each partition retains its own bucket count from its data files. New partitions use the current table-level default. Existing partitions are unaffected until explicitly rescaled.
• Unpartitioned tables: behaviour is unchanged — a full rescale is still required before writing with a new bucket count, and a RuntimeException is thrown if this is violated.
• Rescaling a single partition: use the rescale procedure or a manual INSERT OVERWRITE in batch mode:

  CALL sys.rescale(`table` => 'mydb.orders', `bucket_num` => 32, `partition` => 'tenant_id=123');

After the job completes, the rescaled partition uses 32 buckets while all other partitions are untouched.

Testing

We haven been soaking this change in our test environments and we are seeing good results. Plus, we add a bunch of new tests to validate we are not breaking anything:

• PartitionBucketMappingTest — unit tests for mapping resolution and loadFromTable.
• FixedBucketRowKeyExtractorTest — verifies correct bucket assignment with heterogeneous per-partition counts.
• FileStoreCommitTest— integration tests covering rescale commits with mixed bucket counts.
• FileSystemWriteRestoreTest — covers the empty-bucket write-restore scenario end-to-end, including the non-partitioned corner case.
• RescaleBucketITCase — end-to-end Flink integration tests for INSERT OVERWRITE-based rescale and streaming restore after rescale.
• RescaleActionITCase — end-to-end tests for the rescale procedure action with per-partition targeting.
• TableWriteCoordinatorTest — unit tests for coordinator behaviour under the new mapping.

[JingsongLi]

Review: Per-Partition Bucket Counts

Nice feature — per-partition rescaling without touching unaffected partitions addresses a real operational pain point for skewed workloads. The overall design (partition-bucket mapping loaded from manifests, SchemaBucketFileStoreTable wrapper for overwrite paths) is sound. A few observations below:

---

1. Silent exception swallowing in PartitionBucketMapping.loadFromScan

} catch (Exception e) {
    return new PartitionBucketMapping(defaultBuckets, Collections.emptyMap());
}

If the manifest scan fails (e.g., corrupted manifest, transient I/O error), the code silently falls back to an empty mapping. This means ALL writes would route using the table default bucket count, potentially placing rows in the wrong buckets for already-rescaled partitions — causing silent data corruption (duplicate keys across buckets). At minimum this should log the exception at WARN level so operators have a chance to detect it. Consider whether failing fast would be safer here.

---

2. PartitionBucketMapping staleness in long-running streaming jobs

The mapping is loaded once at job start (in AbstractFileStoreTable.newWriteSelector() / createRowKeyExtractor()). If a partition is rescaled while a streaming job is running (e.g., via the rescale procedure), the running job will continue routing rows using the old mapping. For a partition that was rescaled from 4 to 8 buckets, the streaming job will keep writing to buckets [0,3], potentially missing the new buckets or conflicting with the new layout.

Is there a plan for streaming jobs to detect/reload the mapping (e.g., on checkpoint/restore)? The TableWriteCoordinator does refresh on snapshot changes, which covers the coordinator-based path, but the normal streaming sink path (FlinkSinkBuilder.buildForFixedBucket) loads the mapping once. A documentation note about requiring a restart after rescale would help.

---

3. PartitionEntry.merge() tie-breaking semantics

When two entries have equal lastFileCreationTime, the receiver's (this) totalBuckets wins:

int newTotalBuckets =
    lastFileCreationTime >= entry.lastFileCreationTime
        ? totalBuckets
        : entry.totalBuckets;

This means a.merge(b) and b.merge(a) can produce different totalBuckets values when timestamps match but bucket counts differ. If the reduce/aggregation pipeline in PartitionEntry processing doesn't guarantee stable merge order, this could lead to nondeterministic partition bucket counts. The test testMergeWithEqualCreationTimeTakesFirstTotalBuckets documents the behavior, but it would be good to confirm that the scan pipeline guarantees a deterministic accumulation order for entries with identical creation timestamps.

---

4. TableWriteCoordinator scan reuse concern

In loadPartitionBucketMapping():

this.partitionBucketMapping = PartitionBucketMapping.loadFromScan(scan, defaultNumBuckets);

The scan field is shared between loadPartitionBucketMapping (called in refresh()) and the scan(ScanCoordinationRequest) method (which calls scan.withPartitionBucket(...)). If readPartitionEntries() mutates internal scan state (e.g., filters), subsequent withPartitionBucket calls in scan() could produce incorrect results. This appears safe given that withSnapshot is called before both uses, but worth a defensive comment or using a separate scan instance for the mapping load.

---

5. Serialization size for large partition maps

PartitionBucketMapping is Serializable and is embedded in FixedBucketWriteSelector / FixedBucketRowKeyExtractor, which get serialized across Flink task managers. For tables with many rescaled partitions, the Map<BinaryRow, Integer> could become non-trivial. The optimization to only store partitions that differ from the default is good. Just noting that for extreme cases (tens of thousands of individually rescaled partitions), this could impact checkpoint/serialization size.

---

6. Minor: SchemaBucketFileStoreTable missing newWrite(commitUser, writeId, rowKeyExtractor) override

SchemaBucketFileStoreTable overrides newWrite(String, Integer) to inject its own extractor, but inherits newWrite(String, Integer, RowKeyExtractor) from DelegatedFileStoreTable which just passes through to the wrapped table. If someone calls the 3-arg overload on a SchemaBucketFileStoreTable, the custom extractor logic is bypassed. This is unlikely in practice (callers use the 2-arg version), but could be a footgun for future changes. Consider adding:

@Override
public TableWriteImpl<?> newWrite(
        String commitUser, @Nullable Integer writeId, RowKeyExtractor rowKeyExtractor) {
    // Ignore the passed extractor; always use schema-bucket-based routing
    return wrapped().newWrite(commitUser, writeId, createRowKeyExtractor());
}

---

7. Test coverage

The test suite is thorough — particularly the RescaleBucketITCase.testWriteToEmptyBucketAfterRescaleKeepsPartitionBucketCount which carefully sets up the preconditions for the empty-bucket edge case. The PartitionEntryTest covering merge order independence is also valuable. One gap: there's no test for the scenario where loadFromScan hits an exception — a test verifying the fallback (or better, the logged warning) would strengthen confidence in that path.

---

Overall this is a well-structured contribution with clear separation of concerns. The main risk area is the silent fallback in loadFromScan (point 1) which could lead to hard-to-diagnose data issues in production. The rest are design/robustness suggestions.

[mikedias]

Thanks for the awesome review @JingsongLi, here is the follow-up:

1. Silent exception swallowing in PartitionBucketMapping.loadFromScan

Good point, I agree that failing fast is safer. Removed the try-catch block and let the exception propagate.

(That required changing how AppendOnlySimpleTableTest validates the BUCKET_APPEND_ORDERED flag because it was deleting the manifest rather than checking the order of sequence numbers. Please have a look at that to see if that is OK)

2. PartitionBucketMapping staleness in long-running streaming jobs

It is already expected that the rescale process can only be performed when no job is writing to the table, since it uses strict commit mode. Also, there are additional checks that will fail the commit if the bucket counts don't match. In any case, I've added a note on that in the rescale doc.

3. PartitionEntry.merge() tie-breaking semantics

Good point, no one wants non-deterministic behavior in their codebases 😄. Fixed that, making merge commutative.

4. TableWriteCoordinator scan reuse concern

Sounds good, added a small note in the code.

5. Serialization size for large partition maps

I suspect this scenario is rare, and the workaround would be to set the default bucket count to the most common distribution in the table. That would hit the optimization condition and reduce the memory pressure.

6. SchemaBucketFileStoreTable missing newWrite(commitUser, writeId, rowKeyExtractor) override

FIxed!

7. Test coverage

Covered the loadFromScan scenario.

[JingsongLi]

@mikedias If I understand correctly, this will apply to fixed buckets, which will affect a large number of historical tables and assignments, and the huge compatibility risk is almost unbearable.

Is it possible to continue from the perspective of Postpone, which supports the default number of buckets and allows writing to bypass Postpone and directly write to the buckets.

[mikedias]

@JingsongLi yes, it does apply to fixed buckets. However, for the historical tables and assignments that haven't had partitions rescaled independently, the logic remains the same. Assuming I've covered all scenarios here, historical tables will remain compatible. (FWIW, this patch is working well in our environments so far, and we haven't seen any issues)

Regarding Postpone buckets, we have explored using them, but we faced a few challenges:

• New data isn't visible before compaction, significantly reducing the freshness of these tables. Compacting at the cadence we need would be prohibitively expensive.
• When dealing with thousands of partitions, the Postpone Bucket compaction dag creates too many splits and crashes the Job Manager.

Hence, I'm unsure whether Postpone is the best way to handle heavily skewed, highly partitioned tables that can't wait for compaction to show data.

I understand the concern around changing the code for the fixed bucket assignments. Is there a way we can mitigate the risk here by breaking this PR into smaller ones or by adding more tests to areas of concern?

[JingsongLi]

This statement is broader than the implementation today. The Flink fixed-bucket path now routes through table.createRowKeyExtractor(), but the Spark extension fast path still computes BUCKET_COL with the table-level coreOptions.bucket() in PaimonSparkWriter (HASH_FIXED, paimonExtensionEnabled). After changing the table default from 4 to 8 while an existing partition still has 4 buckets, Spark can route keys for that old partition to buckets 4-7, and the write-side mismatch check is now relaxed for partitioned tables. That breaks the promised per-partition layout and can create files in buckets that do not belong to the partition. Please either update the Spark fixed-bucket path to use the same per-partition bucket mapping, or scope the docs/procedure support to engines that actually use it.

[mikedias]

Fair point. I had a quick look at updating the Spark fast path code, but I couldn't find a performant way to use the per-partition bucket mapping.

Given that this PR is already large, I'll update the docs to note that Spark isn't supported yet and raise an issue to implement it separately.

[JingsongLi]

This relaxation still needs to reject buckets outside the partition layout. When restoreFiles() returns no entries for an empty bucket, WriteRestore.extractTotalBuckets can now recover the partition bucket count from PartitionBucketMapping. For an old partition with 4 buckets after the table default was changed to 8, a caller that still computes buckets with the schema default can ask to write bucket 6. This branch only logs the 4-vs-8 mismatch and then continues, so the writer can create files in bucket 6 even though that partition should only have buckets 0-3. Please keep the per-partition mismatch allowance, but fail when bucket >= totalBuckets (unless ignoreNumBucketCheck is explicitly set for the controlled rescale/postpone path).

[mikedias]

Good point. Implemented the check + tests.

[JingsongLi]

Please remove the trailing spaces introduced in this doc. git diff --check origin/master...HEAD currently reports this line and lines 149, 170, and 171, so the final verification will fail.

[JingsongLi]

Please remove this diagnostic dump before merging. It prints multiple manifest-entry lines during normal test runs, and the assertions below already validate the p2 entries; if extra context is needed, prefer putting it into the assertion descriptions.

[JingsongLi]

@mikedias When considering the impact of changes, we must take into account the effects of super large tables, such as where this PR will affect PB level storage tables, hundreds of thousands of partitions, and thousands of buckets per partition.

[JingsongLi]

So I still -1 for this PR, I can imagine the huge impact of this PR on the existing tables in production.

[mikedias]

@JingsongLi are you against the whole idea of supporting per-partition bucket counts on fixed buckets, or are you just concerned about specific details on this PR?

If it's the latter, I'm happy to work with you to mitigate the concerns. Perhaps smaller incremental PRs or adding configurations would help to gain confidence here?

Regarding the comment about super large tables, I'm unable to see where the impact might be. The partition-bucket map only contains the overrides, so for existing large tables with no overwrites, the current logic and memory footprint are the same. Are there other parts of the code that you think could be problematic?

[JingsongLi]

git diff --check origin/master...HEAD is still failing on trailing whitespace in this document. It currently reports this line plus lines 63, 76, and 77. Please strip those spaces before merge; otherwise the final verification will fail.

[JingsongLi]

Hi @mikedias , I am opposing the default activation of this mode, as it will have a significant impact on the already displayed tables
For large tables, we will try to avoid scanning all manifests as much as possible.
2. If the old partition experiences batch overwrite and changes the number of buckets during the running of the stream job, the statistics recorded in memory by the stream job are completely incorrect.

[mikedias]

Understood, let me think about how we can put this behind a configuration in an elegant way, and whether there are ways to reduce the number of scans of the manifest when loading the partition mapping.

[JingsongLi]

@mikedias ❤️

[JingsongLi]

bucket.per-partition-count-enabled is set to false by default, and the documentation explicitly states that when this setting is disabled, all partitions must share the table-level bucket count; however, when AbstractFileStoreWrite.scanExistingFileMetas detects a discrepancy between the old and new bucket counts in a partitioned table, it still unconditionally allows writes where bucket < restoredTotalBuckets without checking this setting.

• For example, suppose an existing partition uses 4 buckets, the table-level bucket count is changed to 2, and the switch remains disabled. New writes are routed to buckets 0 and 1 using modulo 2, while old data is distributed using modulo 4; the current check would allow this, potentially causing the same primary key to land in different buckets, resulting in duplicate keys and data errors.
• It is recommended to allow partition-level bucket mismatches only when options.bucketPerPartitionCountEnabled() is set to true; when disabled, maintain the original mismatch rejection behavior and add a regression test that defaults to the switch being off.
• The existing PartitionBucketMappingTest and FileSystemWriteRestoreTest—a total of 9 tests—pass, but the tests involving different numbers of buckets per partition explicitly enable the switch, thus missing this specific scenario.

[mikedias]

Thanks for the super quick review @JingsongLi , added the missing validation!

With this configuration, the additional scan is now opt-in, removing the impact on existing tables. However, I feel it is still sub-optimal because it introduces a flag in a critical part of the system, which increases overall complexity and reduces maintainability... Am I being too harsh?

In my mind, the optimal solution would be to reuse existing manifest scans to derive the bucket count, so then we don't need a config or an additional scan. But I couldn't think of how to achieve that in the streaming case to distribute records to the right channels/buckets... Do you think it's possible to achieve that?

[JingsongLi]

Hi @mikedias , I think it would be better to introduce an option to keep things as they are. Paimon is already widely used in production, so we should ensure compatibility.