[python] Write postpone batches to fixed buckets

[XiaoHongbo-Hope]

Purpose

Postpone-bucket files are not visible to normal readers until compaction. Java batch integrations solve this through a dedicated postpone fixed-bucket write builder.

This adds the same explicit builder to PyPaimon and uses it from the Ray batch integration. The regular new_batch_write_builder() keeps its existing bucket-postpone behavior.

Changes

• Add new_postpone_fixed_bucket_write_builder() for bucket=-2 tables.
• Reuse the current bucket count for partitions with real buckets.
• For partitions without real buckets, infer the bucket count from input row count or uncompressed Arrow size, capped by postpone.batch-write-fixed-bucket.max-parallelism.
• Include active postpone rows in append estimates; overwrite ignores them.
• Carry per-partition bucket counts through commit messages and reject conflicting concurrent commits.
• Let Ray plan bucket counts once on the driver, precluster by (partition, bucket), and pass the plan to explicit worker writers.
• Keep regular batch, streaming, and Daft writes on the existing postpone path.

Tests

• table_write_test.py
• ray_sink_test.py
• ray_repartition_test.py
• daft_pk_distributed_write_test.py

[leaves12138]

Thanks for working on this. The overall direction is aligned with the Java Spark fixed-bucket batch path (reuse existing bucket counts, include active postpone rows for append estimation, carry total_buckets, and detect bucket-count conflicts), but I found two correctness/parity issues that should be addressed before merging.

1. Bucket planning is not global for chunked or distributed batch writes

TableWrite.write_arrow_batch calls plan_input for each incoming batch, while PostponeFixedBucketRowKeyExtractor.plan_num_buckets skips a partition as soon as it is present in _known_num_buckets. Therefore, the first batch containing a new partition permanently determines its bucket count; later batches are not included in the estimate.

I reproduced this with the same logical input and postpone.target-size-per-bucket=100 b, max-parallelism=3:

write_arrow_batch(small_batch), then write_arrow_batch(large_batch): {1}
write_arrow(concatenated_input):                              {3}

This differs from Java Spark's preparePostponeBucketAssignment, which persists the complete DataFrame, collects all partition statistics, and resolves bucket counts before routing any rows.

This also affects Daft. PaimonDataSink.write creates an independent BatchTableWrite in each worker and feeds it record batches. Different workers can choose different bucket counts for the same new partition. A two-writer reproduction produced plans {1} and {3}, and the combined commit failed with:

CommitConflictError: Total buckets of partition ('p',) changed from 1 to 3 without overwrite.

Ray avoids this because this PR adds a driver-side plan, but the generic batch API and Daft do not. Bucket planning needs to happen once over the complete input before writing. Daft likely needs a coordinated driver plan similar to Ray, or should retain the legacy postpone-bucket path until such planning is available. Please also add coverage for multiple write_arrow_batch calls and multiple worker message sets for the same partition.

2. postpone.target-row-num-per-bucket is ignored

Java defines postpone.target-row-num-per-bucket and Spark gives it precedence over postpone.target-size-per-bucket. PyPaimon only adds the size option and always plans by Arrow size.

I reproduced this with target-row-num-per-bucket=1, target-size-per-bucket=1 gb, max-parallelism=8, and 8 rows in one new partition:

PyPaimon:   {1}
Java Spark: {8}

This is especially problematic for tables created/configured through Java: PyPaimon silently ignores the persisted table option when it creates the first real buckets for a new partition. Please add the option/accessor and choose the row-count calculation when configured, including active postpone rows for append and excluding them for overwrite, matching the Java Spark behavior.

For context, Java Flink's batch path uses capped sink parallelism for unknown partitions, while this implementation is primarily adapting Java Spark's data-statistics behavior, so it is not expected to be numerically identical to every Java engine. However, the global planning and row-count precedence above are correctness/configuration semantics that should remain consistent.

[XiaoHongbo-Hope]

Thanks for working on this. The overall direction is aligned with the Java Spark fixed-bucket batch path (reuse existing bucket counts, include active postpone rows for append estimation, carry total_buckets, and detect bucket-count conflicts), but I found two correctness/parity issues that should be addressed before merging.

1. Bucket planning is not global for chunked or distributed batch writes

TableWrite.write_arrow_batch calls plan_input for each incoming batch, while PostponeFixedBucketRowKeyExtractor.plan_num_buckets skips a partition as soon as it is present in _known_num_buckets. Therefore, the first batch containing a new partition permanently determines its bucket count; later batches are not included in the estimate.

I reproduced this with the same logical input and postpone.target-size-per-bucket=100 b, max-parallelism=3:

write_arrow_batch(small_batch), then write_arrow_batch(large_batch): {1}
write_arrow(concatenated_input):                              {3}

This differs from Java Spark's preparePostponeBucketAssignment, which persists the complete DataFrame, collects all partition statistics, and resolves bucket counts before routing any rows.

This also affects Daft. PaimonDataSink.write creates an independent BatchTableWrite in each worker and feeds it record batches. Different workers can choose different bucket counts for the same new partition. A two-writer reproduction produced plans {1} and {3}, and the combined commit failed with:

CommitConflictError: Total buckets of partition ('p',) changed from 1 to 3 without overwrite.

Ray avoids this because this PR adds a driver-side plan, but the generic batch API and Daft do not. Bucket planning needs to happen once over the complete input before writing. Daft likely needs a coordinated driver plan similar to Ray, or should retain the legacy postpone-bucket path until such planning is available. Please also add coverage for multiple write_arrow_batch calls and multiple worker message sets for the same partition.

2. postpone.target-row-num-per-bucket is ignored

Java defines postpone.target-row-num-per-bucket and Spark gives it precedence over postpone.target-size-per-bucket. PyPaimon only adds the size option and always plans by Arrow size.

I reproduced this with target-row-num-per-bucket=1, target-size-per-bucket=1 gb, max-parallelism=8, and 8 rows in one new partition:

PyPaimon:   {1}
Java Spark: {8}

This is especially problematic for tables created/configured through Java: PyPaimon silently ignores the persisted table option when it creates the first real buckets for a new partition. Please add the option/accessor and choose the row-count calculation when configured, including active postpone rows for append and excluding them for overwrite, matching the Java Spark behavior.

For context, Java Flink's batch path uses capped sink parallelism for unknown partitions, while this implementation is primarily adapting Java Spark's data-statistics behavior, so it is not expected to be numerically identical to every Java engine. However, the global planning and row-count precedence above are correctness/configuration semantics that should remain consistent.

Thanks for the review.

1. Multiple write_arrow_batch() inputs are now planned together before writing. Ray keeps its driver-side plan; Daft temporarily retains the legacy postpone path. Daft will be processed in following PR.
2. Added postpone.target-row-num-per-bucket with precedence over size, including existing postpone rows for append but not overwrite.
3. Added coverage for multi-batch, multi-worker conflicts, row-count planning, and Daft fallback.
   cc @JingsongLi

[XiaoHongbo-Hope]

Thanks for working on this. The overall direction is aligned with the Java Spark fixed-bucket batch path (reuse existing bucket counts, include active postpone rows for append estimation, carry total_buckets, and detect bucket-count conflicts), but I found two correctness/parity issues that should be addressed before merging.

1. Bucket planning is not global for chunked or distributed batch writes

TableWrite.write_arrow_batch calls plan_input for each incoming batch, while PostponeFixedBucketRowKeyExtractor.plan_num_buckets skips a partition as soon as it is present in _known_num_buckets. Therefore, the first batch containing a new partition permanently determines its bucket count; later batches are not included in the estimate.
I reproduced this with the same logical input and postpone.target-size-per-bucket=100 b, max-parallelism=3:

write_arrow_batch(small_batch), then write_arrow_batch(large_batch): {1}
write_arrow(concatenated_input):                              {3}

This differs from Java Spark's preparePostponeBucketAssignment, which persists the complete DataFrame, collects all partition statistics, and resolves bucket counts before routing any rows.
This also affects Daft. PaimonDataSink.write creates an independent BatchTableWrite in each worker and feeds it record batches. Different workers can choose different bucket counts for the same new partition. A two-writer reproduction produced plans {1} and {3}, and the combined commit failed with:

CommitConflictError: Total buckets of partition ('p',) changed from 1 to 3 without overwrite.

Ray avoids this because this PR adds a driver-side plan, but the generic batch API and Daft do not. Bucket planning needs to happen once over the complete input before writing. Daft likely needs a coordinated driver plan similar to Ray, or should retain the legacy postpone-bucket path until such planning is available. Please also add coverage for multiple write_arrow_batch calls and multiple worker message sets for the same partition.

2. postpone.target-row-num-per-bucket is ignored

Java defines postpone.target-row-num-per-bucket and Spark gives it precedence over postpone.target-size-per-bucket. PyPaimon only adds the size option and always plans by Arrow size.
I reproduced this with target-row-num-per-bucket=1, target-size-per-bucket=1 gb, max-parallelism=8, and 8 rows in one new partition:

PyPaimon:   {1}
Java Spark: {8}

This is especially problematic for tables created/configured through Java: PyPaimon silently ignores the persisted table option when it creates the first real buckets for a new partition. Please add the option/accessor and choose the row-count calculation when configured, including active postpone rows for append and excluding them for overwrite, matching the Java Spark behavior.
For context, Java Flink's batch path uses capped sink parallelism for unknown partitions, while this implementation is primarily adapting Java Spark's data-statistics behavior, so it is not expected to be numerically identical to every Java engine. However, the global planning and row-count precedence above are correctness/configuration semantics that should remain consistent.

Thanks for the review.

1. Multiple write_arrow_batch() inputs are now planned together before writing. Ray keeps its driver-side plan; Daft temporarily retains the legacy postpone path. Daft will be processed in following PR.
2. Added postpone.target-row-num-per-bucket with precedence over size, including existing postpone rows for append but not overwrite.
3. Added coverage for multi-batch, multi-worker conflicts, row-count planning, and Daft fallback.
   cc @JingsongLi

[leaves12138]

Thanks for the update. I re-tested the two original issues: multiple write_arrow_batch calls now produce the same bucket plan as the concatenated input, and postpone.target-row-num-per-bucket now correctly takes precedence. The Daft legacy fallback also looks reasonable.

I found one remaining API consistency issue: BatchTableWrite.write_row still bypasses the new postpone planning path. It immediately calls extract_partition_bucket_row, and PostponeFixedBucketRowKeyExtractor.num_buckets defaults an unknown partition to one bucket.

With postpone.target-row-num-per-bucket=1, postpone.batch-write-fixed-bucket.max-parallelism=8, and 8 rows in one new partition, I get:

write_arrow: {8}
write_row:   {1}

This means the bucket count depends on which public write API is used. Calling write_row first also records the new partition as one bucket, so subsequent Arrow batches cannot re-plan it. Please either buffer and plan row writes together with Arrow inputs, or explicitly reject write_row for postpone fixed-bucket batch writes, and add a regression test.

One additional resource concern: _postpone_batches retains all unknown-partition RecordBatch objects until prepare_commit, so the chunked API now has memory usage proportional to the complete input and cannot benefit from normal incremental file rolling. Spark can persist/spill and Ray can use object-store spilling; the local Python path has no equivalent. Please at least document this memory bound, or consider temporary spilling or an explicit preplanned bucket map.

[XiaoHongbo-Hope]

Thanks for the update. I re-tested the two original issues: multiple write_arrow_batch calls now produce the same bucket plan as the concatenated input, and postpone.target-row-num-per-bucket now correctly takes precedence. The Daft legacy fallback also looks reasonable.

I found one remaining API consistency issue: BatchTableWrite.write_row still bypasses the new postpone planning path. It immediately calls extract_partition_bucket_row, and PostponeFixedBucketRowKeyExtractor.num_buckets defaults an unknown partition to one bucket.

With postpone.target-row-num-per-bucket=1, postpone.batch-write-fixed-bucket.max-parallelism=8, and 8 rows in one new partition, I get:

write_arrow: {8}
write_row:   {1}

This means the bucket count depends on which public write API is used. Calling write_row first also records the new partition as one bucket, so subsequent Arrow batches cannot re-plan it. Please either buffer and plan row writes together with Arrow inputs, or explicitly reject write_row for postpone fixed-bucket batch writes, and add a regression test.

One additional resource concern: _postpone_batches retains all unknown-partition RecordBatch objects until prepare_commit, so the chunked API now has memory usage proportional to the complete input and cannot benefit from normal incremental file rolling. Spark can persist/spill and Ray can use object-store spilling; the local Python path has no equivalent. Please at least document this memory bound, or consider temporary spilling or an explicit preplanned bucket map.

Thanks for the follow-up. write_row() now joins the deferred bucket planning with Arrow inputs. Added pure-row and mixed-input tests, and documented the local buffering memory bound. Existing-bucket partitions still write incrementally. cc @JingsongLi

[JingsongLi]

You might need to introduce some abstraction—right now, a bunch of Postpone implementations are tightly coupled within the common classes TableWrite and BatchTableWrite.

[XiaoHongbo-Hope]

You might need to introduce some abstraction—right now, a bunch of Postpone implementations are tightly coupled within the common classes TableWrite and BatchTableWrite.

Got it, updated

[JingsongLi]

Is it possible to provide a higher-level API to accomplish this?

[XiaoHongbo-Hope]

Is it possible to provide a higher-level API to accomplish this?

Sure, done

[JingsongLi]

You can take a review to #8992

[XiaoHongbo-Hope]

You can take a review to #8992

Thanks. I aligned the implementation with #8992. Could you take another look?

[JingsongLi]

Please not do this. By default, we should not use this mode.

[XiaoHongbo-Hope]

Please not do this. By default, we should not use this mode.

Updated

[JingsongLi]

Thanks for making the fixed-bucket write path explicit. I revalidated the latest head (acf7e4e) against the Java implementation. The API-default concern is addressed, but the four issues below remain. The first can cause cross-engine correctness failures; the others can produce different bucket counts, stale partition metadata, or leaked files under concurrent commits.

[JingsongLi]

[P1] Respect the configured bucket function. This extractor always applies the default BinaryRow hash, while Java's PostponeFixBucketProcessor constructs BucketFunction from CoreOptions and therefore supports default, mod, and hive. For a valid bucket-function.type=mod or hive table, Python and Java can route the same primary key to different buckets; Java bucket pruning then follows the configured function and may miss rows written by Python. Please dispatch according to bucket-function.type (with Java-compatible semantics), or reject non-default functions until they are supported, and add cross-engine fixtures for all supported functions.

[XiaoHongbo-Hope]

[P1] Respect the configured bucket function. This extractor always applies the default BinaryRow hash, while Java's PostponeFixBucketProcessor constructs BucketFunction from CoreOptions and therefore supports default, mod, and hive. For a valid bucket-function.type=mod or hive table, Python and Java can route the same primary key to different buckets; Java bucket pruning then follows the configured function and may miss rows written by Python. Please dispatch according to bucket-function.type (with Java-compatible semantics), or reject non-default functions until they are supported, and add cross-engine fixtures for all supported functions.

Fixed

[JingsongLi]

[P2] Use the same size metric as Java when deriving the bucket count. RecordBatch.nbytes measures Arrow buffer memory, whereas Java sums BinaryRow.getSizeInBytes(). These values differ materially; for example, 1,000 (INT, STRING, STRING) rows occupy 14,000 bytes by this metric but 32,000 bytes as Java BinaryRows, so a 20 KiB target plans 1 bucket in Python and 2 in Java. Because the first writer persists the partition's bucket count, engine order changes the table layout. Please use a BinaryRow-compatible serialized-size calculation (or an exactly equivalent vectorized calculation) and verify it with shared fixtures.

[JingsongLi]

[P2] Preserve the replacement bucket count in partition statistics. This list emits DELETE entries first with the old total_buckets, then ADD entries with the new value. _generate_partition_statistics records only the first value it sees for a partition, so an overwrite from postpone mode (-2) to N leaves catalog partition statistics at -2 even though the new manifest entries use N. Java's PartitionEntry.merge takes the later entry's totalBuckets. Please make the statistics aggregation last-wins (and add a catalog-statistics assertion for -2 -> N, ideally also N -> M).

[XiaoHongbo-Hope]

[P2] Preserve the replacement bucket count in partition statistics. This list emits DELETE entries first with the old total_buckets, then ADD entries with the new value. _generate_partition_statistics records only the first value it sees for a partition, so an overwrite from postpone mode (-2) to N leaves catalog partition statistics at -2 even though the new manifest entries use N. Java's PartitionEntry.merge takes the later entry's totalBuckets. Please make the statistics aggregation last-wins (and add a catalog-statistics assertion for -2 -> N, ideally also N -> M).

Fixed

[JingsongLi]

[P2] Keep deterministic CAS failures on the abortable conflict path. Enabling conflict detection here exposes a cleanup gap: when atomic_commit returns False, _try_commit_once creates RetryResult(exception=None) because the snapshot definitely was not committed. If the retry then detects a bucket-count conflict, _try_commit converts it to a generic exception solely because retry_result is non-null. TableCommit aborts files only for CommitConflictError, so the losing writer's files can be orphaned. Please distinguish retry_result.exception is None from an uncertain commit exception and retain CommitConflictError for the former; add a concurrent different-plan test that also asserts cleanup.

[XiaoHongbo-Hope]

[P2] Keep deterministic CAS failures on the abortable conflict path. Enabling conflict detection here exposes a cleanup gap: when atomic_commit returns False, _try_commit_once creates RetryResult(exception=None) because the snapshot definitely was not committed. If the retry then detects a bucket-count conflict, _try_commit converts it to a generic exception solely because retry_result is non-null. TableCommit aborts files only for CommitConflictError, so the losing writer's files can be orphaned. Please distinguish retry_result.exception is None from an uncertain commit exception and retain CommitConflictError for the former; add a concurrent different-plan test that also asserts cleanup.

Fixed

[JingsongLi]

I re-reviewed the latest head (a80974c) against the Java implementation and revalidated the previous findings. The non-default bucket-function case now fails fast, the overwrite partition statistics fix matches Java, and the single deterministic CAS-failure cleanup case is covered. Two P1 issues remain: one can delete files from a commit whose result was previously uncertain, and the other makes size-based planning reject valid schemas supported by Java.

[JingsongLi]

[P1] Preserve commit uncertainty across all retries before aborting files. _try_commit replaces retry_result after every attempt, so an earlier commit exception (which may have occurred after the snapshot was successfully committed) can be overwritten by a later CAS False result with exception=None. If duplicate detection cannot read the original snapshot because it has expired or is temporarily unavailable, this branch raises CommitConflictError, and TableCommit aborts the data files even though a later snapshot can still reference them through its base manifests. I reproduced: attempt 1 commits snapshot 1 then loses the response; another writer creates snapshot 2; snapshot 1 is unavailable to _is_duplicate_commit; attempt 2 returns False; the final CommitConflictError deletes snapshot 1's data file while snapshot 2 remains latest. Please keep a monotonic commit_result_may_be_uncertain state across the whole retry loop (cleared only after duplicate-success confirmation) and add this multi-attempt regression test.

[JingsongLi]

[P1] Do not use the atomic-only GenericRowSerializer for full-row size planning. This now serializes every input field, but that serializer rejects non-AtomicType values and timezone-aware timestamps; Java uses InternalSerializers.create(rowType), which supports ARRAY, MAP, ROW, VECTOR, VARIANT, BLOB, and TIMESTAMP_LTZ. I reproduced an ARRAY<INT> value failing with ValueError: BinaryRow only support AtomicType, and an id primary-key table with a non-key TIMESTAMP_LTZ value succeeds through the legacy postpone writer but fails here with datetime tzinfo not supported yet. Ray enters this size-based planner by default, so valid existing postpone writes regress before routing any rows. Please use a Java-compatible full type-surface size calculator/serializer and add nested-value and non-key TIMESTAMP_LTZ fixtures in addition to the current all-atomic 32,000-byte fixture.