This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new b0af45681f [python] Expand native cross-partition and postpone write
support (#10354)
b0af45681f is described below
commit b0af45681fffbb1827a23b2f2435098c479dbefa
Author: Jingsong Lee <[email protected]>
AuthorDate: Sun Oct 4 09:03:46 2026 +0800
[python] Expand native cross-partition and postpone write support (#10354)
---
docs/docs/pypaimon/ray-data.md | 7 +-
paimon-python/README.md | 49 ++
.../pypaimon/common/options/core_options.py | 16 +
paimon-python/pypaimon/ray/ray_paimon.py | 4 +-
paimon-python/pypaimon/ray/shuffle.py | 57 +-
.../pypaimon/read/scanner/file_scanner.py | 5 +
.../tests/native_cross_partition_write_test.py | 193 +++++++
.../tests/native_postpone_fixed_write_test.py | 573 +++++++++++++++++++++
.../pypaimon/tests/native_postpone_write_test.py | 186 +++++++
.../pypaimon/tests/ray_repartition_test.py | 36 +-
.../pypaimon/tests/write/dynamic_bucket_test.py | 1 +
.../pypaimon/tests/write/table_write_test.py | 19 +-
.../pypaimon/write/commit/commit_scanner.py | 4 +-
.../pypaimon/write/commit/conflict_detection.py | 8 +
.../write/commit/fixed_bucket_commit_check.py | 72 +++
.../write/commit/overwrite_changes_provider.py | 1 +
paimon-python/pypaimon/write/file_store_commit.py | 14 +-
paimon-python/pypaimon/write/file_store_write.py | 24 +-
paimon-python/pypaimon/write/native_commit.py | 4 +-
paimon-python/pypaimon/write/native_write.py | 51 +-
.../pypaimon/write/postpone_batch_table_write.py | 11 +-
paimon-python/pypaimon/write/postpone_bucket.py | 76 ++-
paimon-python/pypaimon/write/ray_datasink.py | 80 +--
paimon-python/pypaimon/write/table_update.py | 2 +-
paimon-python/pypaimon/write/write_builder.py | 4 +-
25 files changed, 1300 insertions(+), 197 deletions(-)
diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md
index 241c8014cc..5566ee9e36 100644
--- a/docs/docs/pypaimon/ray-data.md
+++ b/docs/docs/pypaimon/ray-data.md
@@ -267,9 +267,10 @@ For non-HASH_FIXED append-only tables, the dataset is
written as-is.
Postpone-bucket tables (`bucket = -2`) follow
`postpone.batch-write-fixed-bucket` (default: `true`). Existing partitions
reuse their bucket count; new partitions infer one from the configured target
-row count or size. Ray materializes the input for this global plan, then sorts
-by partition, bucket, and primary-key hash. One primary key stays with one
-writer, while a large bucket can span multiple Ray blocks. Set
+row count or size. Ray materializes the input for this global plan, then groups
+by partition and bucket. Each complete group is written by one writer inside
+`map_groups()`, including when its input spans multiple Ray blocks. Each group
+must fit in memory on one node. Set
`hash_fixed_precluster="off"` to retain `bucket-postpone` writes. Fixed-bucket
postpone writes support `bucket-function.type=default` only. HASH_DYNAMIC and
CROSS_PARTITION primary-key Ray writes are not supported and fail fast,
diff --git a/paimon-python/README.md b/paimon-python/README.md
index b1b904d6c1..45151f5cab 100644
--- a/paimon-python/README.md
+++ b/paimon-python/README.md
@@ -177,6 +177,55 @@ owner: a HASH ADD replaces the complete previous index and
does not carry a
concurrent-writer baseline. Partitions with existing data but missing HASH
indexes
must be rewritten before incremental writes.
+Cross-partition primary keys (`bucket=-1`, with some partition columns outside
+of the primary key) also support native writes. Rust rebuilds the global index
+from live rows on writer startup and preserves the full primary key.
Deduplicate
+moves a key by deleting its old location; first-row keeps its first partition;
+partial-update and aggregation apply changes in the existing partition,
matching
+Java. Use one writer owner for the global key space. Index TTL is not
supported,
+and `sequence.field` and `bucket-key` are invalid for this mode. The Python
writer
+does not implement cross-partition routing, so this mode requires the native
runtime.
+
+Ordinary postpone writes (`bucket=-2`) support native Parquet deduplicate
writes
+without BLOB columns. Pending files retain input order, row kinds and
duplicates, and
+roll at a batch boundary after reaching `target-file-row-num`. Normal scans
expose
+only real buckets after deferred bucket assignment. Other postpone merge
engines,
+BLOB writes use the Python writer.
+
+The fixed-bucket postpone builder uses native Parquet writes for deduplicate,
+first-row, partial-update and aggregation when supplied with a shared bucket
plan
+or `postpone.default-bucket-num`. The default count is used exactly, without
+rounding or applying the inference limit. Append reuses each existing
partition's
+bucket count; overwrite uses the new count for overwritten partitions. Pending
+files may coexist with real buckets and remain available for later assignment.
+
+```python
+from pypaimon.write.postpone_bucket import PostponeBucketPlan
+
+# A driver can serialize this small plan and give it to every worker.
+plan = PostponeBucketPlan({("2026-10-01",): 3, ("2026-10-02",): 5})
+builder =
table.new_postpone_fixed_bucket_write_builder().with_bucket_plan(plan)
+writer = builder.new_write()
+try:
+ writer.write_arrow(arrow_table)
+ messages = writer.prepare_commit()
+finally:
+ writer.close()
+builder.new_commit().commit(messages)
+```
+
+All rows for one partition/bucket must have one writer owner. Python and native
+committers reject overlapping owners or writes that became stale after their
+baseline snapshot. Different real buckets and pending files can be appended
+concurrently. Overwrite detects concurrent changes to any overwritten
partition.
+A shared plan must cover every input partition; a configured default does not
+fill missing entries in an explicit plan. Ray workers use native routing and
+writing with the driver's shared plan.
+
+Local automatic size estimation still uses the Python planner. Fixed-bucket
+native writes currently exclude BLOB and deletion-vector tables. They do not
+perform staging, compaction or automatic rescaling of existing partitions.
+
Native writes honor `data-file.path-directory` and the configured
`data-file.external-paths` strategy. Existing files keep their recorded
locations
when the write destinations change. Python and native readers and committers
diff --git a/paimon-python/pypaimon/common/options/core_options.py
b/paimon-python/pypaimon/common/options/core_options.py
index ef8c553c28..e419a70415 100644
--- a/paimon-python/pypaimon/common/options/core_options.py
+++ b/paimon-python/pypaimon/common/options/core_options.py
@@ -264,6 +264,16 @@ class CoreOptions:
)
)
+ POSTPONE_DEFAULT_BUCKET_NUM: ConfigOption[int] = (
+ ConfigOptions.key("postpone.default-bucket-num")
+ .int_type()
+ .no_default_value()
+ .with_description(
+ "Exact bucket count for new postpone partitions and overwrite
writes. "
+ "Takes precedence over automatic bucket estimation."
+ )
+ )
+
POSTPONE_TARGET_ROW_NUM_PER_BUCKET: ConfigOption[int] = (
ConfigOptions.key("postpone.target-row-num-per-bucket")
.long_type()
@@ -1308,6 +1318,12 @@ class CoreOptions:
default,
)
+ def postpone_default_bucket_num(self):
+ value = self.options.get(CoreOptions.POSTPONE_DEFAULT_BUCKET_NUM)
+ if value is not None and not 0 < value <= 2147483647:
+ raise ValueError('postpone.default-bucket-num must be a positive
32-bit integer')
+ return value
+
def postpone_target_row_num_per_bucket(self, default=None):
return self.options.get(
CoreOptions.POSTPONE_TARGET_ROW_NUM_PER_BUCKET, default
diff --git a/paimon-python/pypaimon/ray/ray_paimon.py
b/paimon-python/pypaimon/ray/ray_paimon.py
index f14bc5a0ae..65380919d6 100644
--- a/paimon-python/pypaimon/ray/ray_paimon.py
+++ b/paimon-python/pypaimon/ray/ray_paimon.py
@@ -547,8 +547,8 @@ def write_paimon(
writer. For primary-key tables, ``map_groups`` writes each complete
``(partition_keys..., bucket)`` group in one Ray task. Postpone-bucket
writes follow ``postpone.batch-write-fixed-bucket`` by default. Their
- bucket plan is resolved once on the driver and workers write sorted
- blocks to real buckets.
+ bucket plan is resolved once on the driver, then each complete
+ partition/bucket group is written by one task to a real bucket.
HASH_DYNAMIC and CROSS_PARTITION primary-key Ray writes are rejected
because Ray write tasks create independent Paimon writers.
diff --git a/paimon-python/pypaimon/ray/shuffle.py
b/paimon-python/pypaimon/ray/shuffle.py
index f1a7fd6898..c256cc3ad7 100644
--- a/paimon-python/pypaimon/ray/shuffle.py
+++ b/paimon-python/pypaimon/ray/shuffle.py
@@ -22,8 +22,9 @@ The legacy ``map_groups`` strategy groups rows by
``(partition_keys..., bucket)`` so every distinct group lands in a
single Ray task. Primary-key writes consume the complete group in that
task; append-only writes use the regrouped rows as a file-count
-optimization. Ray requires each ``map_groups`` group to fit in memory
-on one node, so keep that strategy behind an explicit opt-in.
+optimization. Postpone fixed-bucket writes use the same grouping to keep
+one writer per partition/bucket. Ray requires each ``map_groups`` group
+to fit in memory on one node.
For append-only tables in any other bucket mode the dataset is returned
unchanged.
@@ -45,7 +46,6 @@ if TYPE_CHECKING:
# runtime by ``_pick_bucket_col_name`` so user tables that happen to
# contain a column with this name still work correctly.
BUCKET_KEY_COL = "__paimon_bucket__"
-WRITER_KEY_COL = "__paimon_writer_key__"
HASH_FIXED_PRECLUSTER_AUTO = "auto"
HASH_FIXED_PRECLUSTER_OFF = "off"
HASH_FIXED_PRECLUSTER_MAP_GROUPS = "map_groups"
@@ -146,57 +146,6 @@ def _group_by_partition_bucket(
return ds_with_bucket.groupby(group_keys), bucket_col
-def _sort_by_partition_bucket_primary_key(
- dataset: "ray.data.Dataset",
- table: "Table",
- extractor,
-):
- """Sort rows so one primary key is owned by one Ray block."""
- partition_keys = list(table.table_schema.partition_keys or [])
- existing_names = set(f.name for f in table.table_schema.fields)
- bucket_col = _pick_bucket_col_name(existing_names)
- existing_names.add(bucket_col)
- writer_key_col = _pick_internal_col_name(
- existing_names, WRITER_KEY_COL
- )
- key_columns = list(table.trimmed_primary_keys)
- key_fields = table.trimmed_primary_keys_fields
-
- def _routing_keys(batch: pa.Table) -> pa.Table:
- if batch.num_rows == 0:
- buckets = []
- writer_keys = []
- else:
- record_batch = batch.combine_chunks().to_batches()[0]
- _, buckets = extractor.extract_partition_bucket_batch(
- record_batch
- )
- columns = [batch.column(name) for name in key_columns]
- writer_keys = [
- extractor._binary_row_hash_code(
- tuple(
- column[row_index].as_py()
- for column in columns
- ),
- key_fields,
- )
- for row_index in range(batch.num_rows)
- ]
- return batch.append_column(
- bucket_col, pa.array(buckets, type=pa.int32())
- ).append_column(
- writer_key_col, pa.array(writer_keys, type=pa.uint32())
- )
-
- with_keys = dataset.map_batches(
- _routing_keys, batch_format="pyarrow", zero_copy_batch=True,
- )
- # Ray keeps equal sort keys in one block. Hash collisions only
- # co-locate additional primary keys.
- sort_keys: List[str] = partition_keys + [bucket_col, writer_key_col]
- return with_keys.sort(sort_keys), [bucket_col, writer_key_col]
-
-
def _identity_batch(batch: pa.Table) -> pa.Table:
return batch
diff --git a/paimon-python/pypaimon/read/scanner/file_scanner.py
b/paimon-python/pypaimon/read/scanner/file_scanner.py
index 17599d970d..503cc78d0a 100755
--- a/paimon-python/pypaimon/read/scanner/file_scanner.py
+++ b/paimon-python/pypaimon/read/scanner/file_scanner.py
@@ -653,6 +653,11 @@ class FileScanner:
return _filter
+ def with_all_buckets(self) -> 'FileScanner':
+ """Include pending buckets for commit metadata scans, as in Java."""
+ self.only_read_real_buckets = False
+ return self
+
def with_shard(self, idx_of_this_subtask: int, number_of_para_subtasks:
int) -> 'FileScanner':
validate_shard(idx_of_this_subtask, number_of_para_subtasks)
if self.start_pos_of_this_subtask is not None:
diff --git a/paimon-python/pypaimon/tests/native_cross_partition_write_test.py
b/paimon-python/pypaimon/tests/native_cross_partition_write_test.py
new file mode 100644
index 0000000000..04b6ae1059
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_cross_partition_write_test.py
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Native cross-partition writes use the Java global-key routing contract."""
+
+from unittest.mock import patch
+
+import pyarrow as pa
+import pytest
+
+from pypaimon import CatalogFactory, Schema
+from pypaimon.table.row.generic_row import GenericRow
+from pypaimon.write.native_write import NativeTableWrite
+
+pytestmark = pytest.mark.native_plan
+_SCHEMA = pa.schema([('tenant', pa.string()), ('p', pa.string()),
+ ('id', pa.int32()), ('v', pa.int32())])
+
+
+def _table(tmp_path, options=None, catalog=None, rowkind=False):
+ catalog = catalog or CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_database('default', True)
+ opts = {'bucket': '-1', 'file.format': 'parquet',
+ 'dynamic-bucket.target-row-num': '2', 'write.native.enabled':
'true'}
+ opts.update(options or {})
+ schema = _SCHEMA.append(pa.field('op', pa.string())) if rowkind else
_SCHEMA
+ catalog.create_table('default.t', Schema.from_pyarrow_schema(
+ schema, primary_keys=['tenant', 'id'], partition_keys=['tenant', 'p'],
options=opts), False)
+ return catalog.get_table('default.t')
+
+
+def _batch(rows, rowkind=False):
+ schema = _SCHEMA.append(pa.field('op', pa.string())) if rowkind else
_SCHEMA
+ return pa.RecordBatch.from_pylist([dict(zip(schema.names, row)) for row in
rows], schema=schema)
+
+
+def _rows(table, native):
+ builder = table.copy({'scan.native-plan.enabled': str(native).lower(),
+ 'read.native.enabled':
str(native).lower()}).new_read_builder()
+ # Java first-row scans hide un-compacted L0 files; a write scan includes
them.
+ splits = builder.new_scan().plan_for_write().splits()
+ reader = builder.new_read()
+ if native:
+ with patch.object(reader, '_create_split_read',
side_effect=AssertionError('fallback')):
+ rows = reader.to_arrow(splits).to_pylist()
+ else:
+ rows = reader.to_arrow(splits).to_pylist()
+ return sorted(tuple(row[name] for name in _SCHEMA.names) for row in rows)
+
+
+def _commit(commit, messages, native_commit=False, checkpoint=None):
+ def publish():
+ if checkpoint is None:
+ commit.commit(messages)
+ else:
+ commit.commit(messages, checkpoint)
+ if native_commit:
+ with patch.object(commit.file_store_commit, 'commit',
+ side_effect=AssertionError('Python commit
fallback')):
+ publish()
+ else:
+ publish()
+
+
[email protected]('native_commit', [False, True])
[email protected]('engine,expected', [
+ ('deduplicate', [('x', 'b', 1, 20), ('y', 'b', 1, 200)]),
+ ('first-row', [('x', 'a', 1, 10), ('y', 'a', 1, 100)]),
+ ('partial-update', [('x', 'a', 1, 20), ('y', 'a', 1, 200)]),
+ ('aggregation', [('x', 'a', 1, 30), ('y', 'a', 1, 300)]),
+])
+def test_cross_partition_engines_rebuild_full_primary_keys(
+ tmp_path, native_rest_catalog, native_commit, engine, expected):
+ opts = {'merge-engine': engine, 'commit.native.enabled':
str(native_commit).lower()}
+ if engine == 'aggregation':
+ opts['fields.v.aggregate-function'] = 'sum'
+ table = _table(tmp_path, opts, native_rest_catalog)
+ for data in ([('x', 'a', 1, 10), ('y', 'a', 1, 100)],
+ [('x', 'b', 1, 20), ('y', 'b', 1, 200)]):
+ builder = table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ assert isinstance(writer, NativeTableWrite)
+ try:
+ writer.write_arrow_batch(_batch(data))
+ assert writer._python_writer is None
+ messages = writer.prepare_commit()
+ # Global index is rebuilt from data, not the HASH index of local
dynamic buckets.
+ assert all(not m.index_adds and not m.index_deletes for m in
messages)
+ _commit(commit, messages, native_commit)
+ finally:
+ writer.close()
+ commit.close()
+ for native in (False, True):
+ assert _rows(table, native) == expected
+
+
[email protected]('chunk_size', [1, 2, 3, 4])
[email protected]('stream', [False, True])
+def test_cross_partition_migration_event_order(tmp_path, chunk_size, stream):
+ table = _table(tmp_path)
+ builder = table.new_stream_write_builder() if stream else
table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ assert isinstance(writer, NativeTableWrite)
+ data = [('x', 'a/b', 1, 10), ('x', 'c=d', 1, 20),
+ ('x', 'a/b', 1, 30), ('x', 'a/b', 1, 40)]
+ try:
+ for start in range(0, len(data), chunk_size):
+ writer.write_arrow_batch(_batch(data[start:start + chunk_size]))
+ _commit(commit, writer.prepare_commit(1) if stream else
writer.prepare_commit(),
+ checkpoint=1 if stream else None)
+ if stream:
+ # The same assigner must retain locations after preparing a
checkpoint.
+ writer.write_arrow_batch(_batch([('x', 'c=d', 1, 50)]))
+ _commit(commit, writer.prepare_commit(2), checkpoint=2)
+ assert writer.prepare_commit(3) == []
+ finally:
+ writer.close()
+ commit.close()
+ expected = [('x', 'c=d', 1, 50)] if stream else [data[-1]]
+ for native in (False, True):
+ assert _rows(table, native) == expected
+
+
+def test_cross_partition_rowkind_and_row_input(tmp_path):
+ table = _table(tmp_path, {'rowkind.field': 'op'}, rowkind=True)
+ builder = table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ assert isinstance(writer, NativeTableWrite)
+ try:
+ writer.write_row(GenericRow(['x', 'a', 1, 10, '+I'], table.fields))
+ writer.write_arrow_batch(_batch([
+ ('x', 'b', 1, 20, '+U'), ('x', 'a', 1, 30, '+U'),
+ ('x', 'a', 2, 40, '+I'), ('x', 'b', 2, 50, '-D')], True))
+ assert writer._python_writer is None
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+ for native in (False, True):
+ assert _rows(table, native) == [('x', 'a', 1, 30)]
+
+
[email protected]('committed', [False, True])
+def test_cross_partition_abort_preserves_committed_files(tmp_path, committed):
+ table = _table(tmp_path)
+ builder = table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ assert isinstance(writer, NativeTableWrite)
+ try:
+ writer.write_arrow_batch(_batch([('x', 'a', 1, 10), ('x', 'b', 1,
20)]))
+ messages = writer.prepare_commit()
+ paths = [f.file_path for m in messages for f in m.new_files]
+ assert paths and all(table.file_io.exists(p) for p in paths)
+ if committed:
+ commit.commit(messages)
+ writer.abort()
+ assert all(table.file_io.exists(p) == committed for p in paths)
+ finally:
+ writer.close()
+ commit.close()
+ assert _rows(table, True) == ([('x', 'b', 1, 20)] if committed else [])
+
+
+def test_cross_partition_dynamic_overwrite_keeps_migration_deletes(tmp_path):
+ table = _table(tmp_path, {'dynamic-partition-overwrite': 'true'})
+ for data, overwrite in [([('x', 'a', 1, 10), ('y', 'c', 2, 20)], False),
+ ([('x', 'b', 1, 30)], True)]:
+ builder = table.new_batch_write_builder()
+ if overwrite:
+ builder.overwrite()
+ writer, commit = builder.new_write(), builder.new_commit()
+ assert isinstance(writer, NativeTableWrite)
+ try:
+ writer.write_arrow_batch(_batch(data))
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+ for native in (False, True):
+ assert _rows(table, native) == [('x', 'b', 1, 30), ('y', 'c', 2, 20)]
diff --git a/paimon-python/pypaimon/tests/native_postpone_fixed_write_test.py
b/paimon-python/pypaimon/tests/native_postpone_fixed_write_test.py
new file mode 100644
index 0000000000..d9dafe6148
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_postpone_fixed_write_test.py
@@ -0,0 +1,573 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Fixed-bucket native writes use shared routing and Java commit semantics."""
+
+import pickle
+from unittest.mock import Mock, patch
+
+import pyarrow as pa
+import pytest
+
+from pypaimon import CatalogFactory, Schema
+from pypaimon.table.row.generic_row import GenericRow
+from pypaimon.write.native_write import NativePostponeFixedBucketTableWrite
+from pypaimon.write.postpone_bucket import PostponeBucketPlan,
PostponeBucketPlanner
+
+pytestmark = pytest.mark.native_plan
+_SCHEMA = pa.schema([('id', pa.int32()), ('p', pa.string()), ('v',
pa.int32())])
+
+
+def _table(tmp_path, options=None, catalog=None, schema=_SCHEMA,
partition_keys=None):
+ catalog = catalog or CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_database('default', True)
+ opts = {'bucket': '-2', 'file.format': 'parquet', 'write.native.enabled':
'true',
+ 'postpone.default-bucket-num': '3'}
+ opts.update(options or {})
+ keys = ['p'] if partition_keys is None else partition_keys
+ catalog.create_table('default.t', Schema.from_pyarrow_schema(
+ schema, primary_keys=['id'] + keys, partition_keys=keys,
options=opts), False)
+ return catalog.get_table('default.t')
+
+
+def _batch(rows):
+ return pa.RecordBatch.from_pylist([dict(zip(_SCHEMA.names, row)) for row
in rows], schema=_SCHEMA)
+
+
+def _rows(table):
+ read = table.new_read_builder()
+ return
sorted(read.new_read().to_arrow(read.new_scan().plan().splits()).to_pylist(),
+ key=lambda row: (row['p'], row['id']))
+
+
+def _write(table, rows, plan=None, overwrite=False):
+ builder = table.new_postpone_fixed_bucket_write_builder()
+ if plan is not None:
+ builder.with_bucket_plan(plan)
+ if overwrite:
+ builder.overwrite()
+ writer, commit = builder.new_write(), builder.new_commit()
+ try:
+ writer.write_arrow_batch(_batch(rows))
+ messages = writer.prepare_commit()
+ commit.commit(messages)
+ return messages
+ finally:
+ writer.close()
+ commit.close()
+
+
[email protected]('native_commit', [False, True])
[email protected]('explicit_plan', [False, True])
+def test_fixed_buckets_are_visible_to_both_readers(
+ tmp_path, native_rest_catalog, native_commit, explicit_plan):
+ table = _table(tmp_path, {'commit.native.enabled':
str(native_commit).lower()}, native_rest_catalog)
+ plan = PostponeBucketPlan({('a/b',): 3, ('b',): 5}) if explicit_plan else
None
+ builder = table.new_postpone_fixed_bucket_write_builder()
+ if plan is not None:
+ builder.with_bucket_plan(plan)
+ writer, commit = builder.new_write(), builder.new_commit()
+ assert isinstance(writer, NativePostponeFixedBucketTableWrite)
+ expected = [dict(id=i, p=p, v=i * 10) for p in ['a/b', 'b'] for i in
range(12)]
+ try:
+ writer.write_arrow(pa.Table.from_pylist(expected, schema=_SCHEMA))
+ writer.write_row(GenericRow([11, 'b', 999], table.fields))
+ messages = writer.prepare_commit()
+ assert writer._python_writer is None
+ assert messages and all(m.bucket >= 0 for m in messages)
+ for message in messages:
+ count = 5 if explicit_plan and message.partition == ('b',) else 3
+ assert message.total_buckets == count
+ assert 0 <= message.bucket < count
+ assert message.check_from_snapshot == 0
+ writer.close()
+ if native_commit:
+ with patch.object(commit.file_store_commit, 'commit',
+ side_effect=AssertionError('Python commit
fallback')):
+ commit.commit(messages)
+ else:
+ commit.commit(messages)
+ finally:
+ writer.close()
+ commit.close()
+ expected[-1]['v'] = 999
+ for native in (False, True):
+ copied = table.copy({'read.native.enabled': str(native).lower(),
+ 'scan.native-plan.enabled': str(native).lower()})
+ assert _rows(copied) == expected
+
+
[email protected]('native', [False, True])
+def test_default_is_exact_and_overwrite_replaces_old_layout(tmp_path, native):
+ table = _table(tmp_path, {'write.native.enabled': str(native).lower(),
+
'postpone.batch-write-fixed-bucket.max-parallelism': '1',
+ 'postpone.target-row-num-per-bucket': '0',
+ 'postpone.target-size-per-bucket': 'invalid'})
+ first = _write(table, [(1, 'a', 10), (2, 'b', 20)])
+ assert {message.total_buckets for message in first} == {3}
+ changed = table.copy({'postpone.default-bucket-num': '5'})
+ appended = _write(changed, [(1, 'a', 11), (3, 'c', 30)])
+ assert {m.partition: m.total_buckets for m in appended} == {('a',): 3,
('c',): 5}
+ replaced = _write(changed, [(4, 'a', 40)], overwrite=True)
+ assert {message.total_buckets for message in replaced} == {5}
+ assert _rows(changed) == [dict(id=4, p='a', v=40), dict(id=2, p='b',
v=20), dict(id=3, p='c', v=30)]
+
+
[email protected]('engine,expected', [('deduplicate', 30),
('partial-update', 30),
+ ('aggregation', 60),
('first-row', 10)])
+def test_fixed_writer_uses_core_merge_engine(tmp_path, engine, expected):
+ options = {'merge-engine': engine}
+ if engine == 'aggregation':
+ options['fields.v.aggregate-function'] = 'sum'
+ table = _table(tmp_path, options)
+ builder = table.new_postpone_fixed_bucket_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ assert isinstance(writer, NativePostponeFixedBucketTableWrite)
+ try:
+ for value in [10, 20, 30]:
+ writer.write_arrow_batch(_batch([(1, 'a', value)]))
+ messages = writer.prepare_commit()
+ assert sum(f.row_count for m in messages for f in m.new_files) == 1
+ # Inspect output even for first-row, whose normal read hides level
zero.
+ import pyarrow.parquet as pq
+ values = []
+ for message in messages:
+ for file in message.new_files:
+ with table.file_io.new_input_stream(file.file_path) as source:
+
values.extend(pq.ParquetFile(source).read().column('v').to_pylist())
+ assert values == [expected]
+ commit.commit(messages)
+ finally:
+ writer.close()
+ commit.close()
+ if engine != 'first-row':
+ assert _rows(table) == [dict(id=1, p='a', v=expected)]
+
+
[email protected]('native', [False, True])
[email protected]('count', ['0', '-1', '2147483648', 'x'])
+def test_default_bucket_validation_is_not_silently_ignored(tmp_path, native,
count):
+ table = _table(tmp_path, {'write.native.enabled': str(native).lower(),
+ 'postpone.default-bucket-num': count})
+ with pytest.raises((ValueError, TypeError),
match='(postpone.default-bucket-num|invalid literal)'):
+ table.new_postpone_fixed_bucket_write_builder().new_write()
+
+
+def test_provided_plan_does_not_fill_missing_partitions_from_default(tmp_path):
+ table = _table(tmp_path)
+ builder = table.new_postpone_fixed_bucket_write_builder().with_bucket_plan(
+ PostponeBucketPlan({('a',): 2}))
+ writer = builder.new_write()
+ assert isinstance(writer, NativePostponeFixedBucketTableWrite)
+ try:
+ writer.write_arrow_batch(_batch([(1, 'a', 10)]))
+ with pytest.raises(Exception, match='does not contain an input
partition'):
+ writer.write_arrow_batch(_batch([(2, 'missing', 20)]))
+ with pytest.raises(Exception, match='closed or failed'):
+ writer.prepare_commit()
+ assert table.snapshot_manager().get_latest_snapshot() is None
+ assert writer._python_writer is None
+ finally:
+ writer.abort()
+
+
[email protected]('prepared', [False, True])
+def test_abort_and_close_keep_commit_message_ownership(tmp_path, prepared):
+ table = _table(tmp_path)
+ builder = table.new_postpone_fixed_bucket_write_builder()
+ writer = builder.new_write()
+ writer.write_arrow_batch(_batch([(1, 'a', 10)]))
+ paths = []
+ if prepared:
+ messages = writer.prepare_commit()
+ paths = [f.file_path for m in messages for f in m.new_files]
+ assert paths and all(table.file_io.exists(path) for path in paths)
+ writer.abort()
+ writer.close()
+ assert all(not table.file_io.exists(path) for path in paths)
+ assert table.snapshot_manager().get_latest_snapshot() is None
+
+
+def test_close_preserves_prepared_and_committed_files(tmp_path):
+ table = _table(tmp_path)
+ builder = table.new_postpone_fixed_bucket_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ try:
+ writer.write_arrow_batch(_batch([(1, 'a', 10)]))
+ messages = writer.prepare_commit()
+ writer.close()
+ paths = [f.file_path for m in messages for f in m.new_files]
+ assert paths and all(table.file_io.exists(path) for path in paths)
+ commit.commit(messages)
+ writer.abort()
+ assert all(table.file_io.exists(path) for path in paths)
+ assert _rows(table) == [dict(id=1, p='a', v=10)]
+ finally:
+ writer.close()
+ commit.close()
+
+
+def test_fixed_writer_is_one_shot(tmp_path):
+ table = _table(tmp_path)
+ writer = table.new_postpone_fixed_bucket_write_builder().new_write()
+ try:
+ assert writer.prepare_commit() == []
+ with pytest.raises(Exception, match='one-time|one prepare_commit'):
+ writer.prepare_commit()
+ with pytest.raises(Exception, match='one prepare_commit'):
+ writer.write_arrow_batch(_batch([(1, 'a', 10)]))
+ finally:
+ writer.close()
+
+
+def test_python_fallback_retains_specialized_writer_and_plan(tmp_path):
+ from pypaimon.write.postpone_batch_table_write import
PostponeFixedBucketBatchTableWrite
+
+ table = _table(tmp_path)
+ plan = PostponeBucketPlan({('a',): 7})
+ writer =
table.new_postpone_fixed_bucket_write_builder().with_bucket_plan(plan).new_write()
+ assert isinstance(writer, NativePostponeFixedBucketTableWrite)
+ try:
+ # Distributed coordinators and advanced methods may select Python
before data.
+ writer.file_store_write
+ assert isinstance(writer._python_writer,
PostponeFixedBucketBatchTableWrite)
+ writer.write_row(GenericRow([1, 'a', 10], table.fields))
+ messages = writer.prepare_commit()
+ assert {message.total_buckets for message in messages} == {7}
+ finally:
+ writer.abort()
+
+
+def test_cannot_fallback_after_native_data(tmp_path):
+ table = _table(tmp_path)
+ writer = table.new_postpone_fixed_bucket_write_builder().new_write()
+ try:
+ writer.write_arrow_batch(_batch([(1, 'a', 10)]))
+ with pytest.raises(RuntimeError, match='Cannot switch'):
+ writer.file_store_write
+ assert writer._python_writer is None
+ finally:
+ writer.abort()
+
+
+def test_distributed_worker_uses_pickled_plan_without_python_routing(tmp_path):
+ from pypaimon.write.ray_datasink import PaimonDatasink
+ from pypaimon.write.row_key_extractor import
PostponeFixedBucketRowKeyExtractor
+
+ table = _table(tmp_path)
+ plan = pickle.loads(pickle.dumps(PostponeBucketPlan({('a',): 3})))
+ sink = PaimonDatasink(table, postpone_bucket_plan=plan)
+ with patch.object(PostponeBucketPlanner, '_load_bucket_metadata',
+ side_effect=AssertionError('worker planner scan')), \
+ patch.object(PostponeFixedBucketRowKeyExtractor, 'extract',
+ side_effect=AssertionError('Python routing'),
create=True):
+ messages = sink.write([pa.Table.from_batches([_batch([(i, 'a', i) for
i in range(20)])])], Mock())
+ assert {m.total_buckets for m in messages} == {3}
+ assert {m.bucket for m in messages} == {0, 1, 2}
+ assert sum(f.row_count for m in messages for f in m.new_files) == 20
+
+
+def test_plan_arrow_encoding_preserves_colliding_partition_name(tmp_path):
+ schema = pa.schema([('id', pa.int32()), ('total_buckets', pa.string()),
('v', pa.int32())])
+ table = _table(tmp_path, schema=schema, partition_keys=['total_buckets'])
+ plan = PostponeBucketPlan({('a',): 3})
+ encoded = plan.to_arrow(table)
+ assert encoded.column(0).to_pylist() == ['a']
+ assert encoded.column(1).to_pylist() == [3]
+ builder =
table.new_postpone_fixed_bucket_write_builder().with_bucket_plan(plan)
+ writer, commit = builder.new_write(), builder.new_commit()
+ try:
+ assert isinstance(writer, NativePostponeFixedBucketTableWrite)
+ writer.write_arrow_batch(pa.RecordBatch.from_pylist(
+ [dict(id=1, total_buckets='a', v=10)], schema=schema))
+ messages = writer.prepare_commit()
+ assert [m.partition for m in messages] == [('a',)]
+ assert {m.total_buckets for m in messages} == {3}
+ commit.commit(messages)
+ finally:
+ writer.close()
+ commit.close()
+
+
[email protected]('partition, count, message', [
+ ((), 1, 'arity'), (('a',), 0, 'positive'), (('a',), -1, 'positive'),
+ (('a',), True, 'positive'), (('a',), 1.5, 'positive'),
+ (('a',), 2147483648, 'positive'),
+])
+def test_invalid_shared_plan_fails_before_writing(tmp_path, partition, count,
message):
+ table = _table(tmp_path)
+ with pytest.raises(ValueError, match=message):
+ PostponeBucketPlan({partition: count}).to_arrow(table)
+
+
+def test_empty_plan_has_typed_schema(tmp_path):
+ table = _table(tmp_path)
+ plan = PostponeBucketPlan({}).to_arrow(table)
+ assert plan.num_rows == 0
+ assert plan.schema.names == ['p', 'total_buckets']
+ assert plan.schema.field(0).type == pa.string()
+ assert plan.schema.field(1).type == pa.int32()
+
+
+def test_unpartitioned_default_writer(tmp_path):
+ table = _table(tmp_path, partition_keys=[])
+ messages = _write(table, [(i, 'a', i) for i in range(20)])
+ assert {m.total_buckets for m in messages} == {3}
+ assert {m.bucket for m in messages} == {0, 1, 2}
+ assert all(m.partition == () for m in messages)
+ assert len(_rows(table)) == 20
+
+
[email protected]('native_commit', [False, True])
[email protected]('combined', [False, True])
+def test_concurrent_bucket_owners_cannot_silently_lose_update(
+ tmp_path, native_rest_catalog, native_commit, combined):
+ table = _table(tmp_path, {'commit.native.enabled':
str(native_commit).lower(),
+ 'postpone.default-bucket-num': '1'},
native_rest_catalog)
+ builders = [table.new_postpone_fixed_bucket_write_builder() for _ in
range(2)]
+ writers = [builder.new_write() for builder in builders]
+ commits = [builder.new_commit() for builder in builders]
+ try:
+ for i in range(10):
+ writers[0].write_arrow_batch(_batch([(1, 'a', i)]))
+ writers[1].write_arrow_batch(_batch([(1, 'a', 999)]))
+ messages = [writer.prepare_commit() for writer in writers]
+ assert {m.check_from_snapshot for group in messages for m in group} ==
{0}
+ if combined:
+ with pytest.raises(Exception, match='ownership conflict'):
+ commits[0].commit(messages[0] + messages[1])
+ assert table.snapshot_manager().get_latest_snapshot() is None
+ else:
+ commits[0].commit(messages[0])
+ with pytest.raises(Exception, match='ownership conflict'):
+ commits[1].commit(messages[1])
+ assert _rows(table) == [dict(id=1, p='a', v=9)]
+ assert table.snapshot_manager().get_latest_snapshot().id == 1
+ finally:
+ for writer in writers:
+ writer.close()
+ for commit in commits:
+ commit.close()
+
+
[email protected]('native_commit', [False, True])
+def test_disjoint_bucket_owners_can_commit_concurrently(tmp_path,
native_rest_catalog, native_commit):
+ from pypaimon.write.row_key_extractor import
PostponeFixedBucketRowKeyExtractor
+
+ table = _table(tmp_path, {'commit.native.enabled':
str(native_commit).lower()}, native_rest_catalog)
+ extractor = PostponeFixedBucketRowKeyExtractor(table,
PostponeBucketPlan({('a',): 3}))
+ representatives = {}
+ for i in range(50):
+ _, bucket = extractor.extract_partition_bucket_row(dict(id=i, p='a',
v=i))
+ representatives.setdefault(bucket, i)
+ assert set(representatives) == {0, 1, 2}
+ builders = [table.new_postpone_fixed_bucket_write_builder() for _ in
representatives]
+ writers = [builder.new_write() for builder in builders]
+ commits = [builder.new_commit() for builder in builders]
+ try:
+ for writer, i in zip(writers, representatives.values()):
+ writer.write_arrow_batch(_batch([(i, 'a', i)]))
+ messages = [writer.prepare_commit() for writer in writers]
+ assert len({m.bucket for group in messages for m in group}) == 3
+ for commit, group in zip(commits, messages):
+ commit.commit(group)
+ assert _rows(table) == [dict(id=i, p='a', v=i) for i in
sorted(representatives.values())]
+ finally:
+ for writer in writers:
+ writer.close()
+ for commit in commits:
+ commit.close()
+
+
[email protected]('native_commit', [False, True])
[email protected]('order', ['pending-first', 'fixed-first',
'concurrent'])
+def test_pending_and_real_buckets_coexist_without_rewriting_history(
+ tmp_path, native_rest_catalog, native_commit, order):
+ table = _table(tmp_path, {'commit.native.enabled':
str(native_commit).lower()}, native_rest_catalog)
+ fixed = table.new_postpone_fixed_bucket_write_builder()
+ pending = table.new_batch_write_builder()
+ fixed_writer, pending_writer = fixed.new_write(), pending.new_write()
+ fixed_commit, pending_commit = fixed.new_commit(), pending.new_commit()
+ try:
+ if order == 'pending-first':
+ pending_writer.write_arrow_batch(_batch([(2, 'a', 20)]))
+ pending_messages = pending_writer.prepare_commit()
+ pending_commit.commit(pending_messages)
+ fixed_writer.write_arrow_batch(_batch([(1, 'a', 10)]))
+ fixed_messages = fixed_writer.prepare_commit()
+ if order == 'fixed-first':
+ fixed_commit.commit(fixed_messages)
+ if order != 'pending-first':
+ pending_writer.write_arrow_batch(_batch([(2, 'a', 20)]))
+ pending_messages = pending_writer.prepare_commit()
+ pending_commit.commit(pending_messages)
+ if order != 'fixed-first':
+ fixed_commit.commit(fixed_messages)
+ snapshot = table.snapshot_manager().get_latest_snapshot()
+ assert snapshot.total_record_count == 2
+ for message in pending_messages + fixed_messages:
+ for file in message.new_files:
+ assert table.file_io.exists(file.file_path)
+ # Ordinary reads expose real buckets; committed pending rows await
assignment.
+ assert _rows(table) == [dict(id=1, p='a', v=10)]
+ scanner = table.new_read_builder().new_scan().file_scanner
+ manifests, _ = scanner.manifest_scanner()
+ entries = scanner.with_all_buckets().read_manifest_entries(manifests)
+ assert {entry.bucket for entry in entries} == {-2,
fixed_messages[0].bucket}
+ assert sum(entry.file.row_count for entry in entries) == 2
+ finally:
+ fixed_writer.close()
+ pending_writer.close()
+ fixed_commit.close()
+ pending_commit.close()
+
+
[email protected]('native_commit', [False, True])
+def test_fixed_overwrite_rejects_concurrent_pending_rows(tmp_path,
native_rest_catalog, native_commit):
+ table = _table(tmp_path, {'commit.native.enabled':
str(native_commit).lower()}, native_rest_catalog)
+ _write(table, [(1, 'a', 10)])
+ builder = table.new_postpone_fixed_bucket_write_builder().overwrite()
+ writer, commit = builder.new_write(), builder.new_commit()
+ pending = table.new_batch_write_builder()
+ pending_writer, pending_commit = pending.new_write(), pending.new_commit()
+ try:
+ writer.write_arrow_batch(_batch([(1, 'a', 30)]))
+ messages = writer.prepare_commit()
+ pending_writer.write_arrow_batch(_batch([(2, 'a', 20)]))
+ pending_commit.commit(pending_writer.prepare_commit())
+ with pytest.raises(Exception, match='overwrite conflict'):
+ commit.commit(messages)
+ assert
table.snapshot_manager().get_latest_snapshot().total_record_count == 2
+ assert _rows(table) == [dict(id=1, p='a', v=10)]
+ finally:
+ writer.close()
+ commit.close()
+ pending_writer.close()
+ pending_commit.close()
+
+
[email protected]('overwrite,static_partition,expected', [
+ (False, None, 3), (True, None, 5), (False, {}, 5), (False, {'p': 'a'}, 5),
+])
+def test_ray_coordinator_uses_new_default_for_overwrite(
+ tmp_path, overwrite, static_partition, expected):
+ from pypaimon.write.ray_datasink import write_paimon_dataset
+
+ table = _table(tmp_path)
+ _write(table, [(1, 'a', 10)])
+ changed = table.copy({'postpone.default-bucket-num': '5'})
+ dataset = Mock()
+ with patch('pypaimon.write.ray_datasink._collect_partition_stats',
+ return_value=(dataset, {('a',): (1, 0)})), \
+ patch('pypaimon.write.ray_datasink._write_primary_key_groups') as
write:
+ write_paimon_dataset(dataset, changed, overwrite=overwrite,
static_partition=static_partition)
+ assert write.call_args.kwargs['postpone_bucket_plan'].as_dict() ==
{('a',): expected}
+
+
[email protected]('native', [False, True])
+def test_fixed_writer_baseline_survives_message_serialization(tmp_path,
native):
+ from pypaimon.write.commit_message_serializer import
serialize_commit_message, deserialize_commit_message
+
+ table = _table(tmp_path, {'write.native.enabled': str(native).lower()})
+ _write(table, [(1, 'a', 10)])
+ writer = table.new_postpone_fixed_bucket_write_builder().new_write()
+ try:
+ writer.write_arrow_batch(_batch([(1, 'a', 20)]))
+ messages = writer.prepare_commit()
+ for message in messages:
+ imported = deserialize_commit_message(
+ serialize_commit_message(message, table.partition_keys_fields),
+ table.partition_keys_fields, table.trimmed_primary_keys_fields)
+ assert imported.check_from_snapshot == 1
+ assert imported.total_buckets == 3
+ assert imported.bucket == message.bucket
+ finally:
+ writer.abort()
+
+
[email protected]('bucket,count', [(-1, 3), (3, 3), (0, 0)])
+def test_python_commit_rejects_invalid_fixed_bucket_metadata(tmp_path, bucket,
count):
+ table = _table(tmp_path, {'commit.native.enabled': 'false'})
+ builder = table.new_postpone_fixed_bucket_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ try:
+ writer.write_arrow_batch(_batch([(1, 'a', 10)]))
+ messages = writer.prepare_commit()
+ messages[0].bucket = bucket
+ messages[0].total_buckets = count
+ with pytest.raises((ValueError, RuntimeError), match='Invalid fixed
bucket'):
+ commit.commit(messages)
+ assert table.snapshot_manager().get_latest_snapshot() is None
+ finally:
+ writer.close()
+ commit.close()
+
+
[email protected]('kind', ['int', 'date', 'decimal', 'timestamp',
'unicode'])
+def test_shared_plan_partition_encoding_matches_python(tmp_path, kind):
+ from datetime import date, datetime
+ from decimal import Decimal
+
+ types = {'int': (pa.int32(), [-7, 9]),
+ 'date': (pa.date32(), [date(1969, 12, 31), date(2026, 10, 3)]),
+ 'decimal': (pa.decimal128(12, 2), [Decimal('-1.25'),
Decimal('9.50')]),
+ 'timestamp': (pa.timestamp('us'), [datetime(1969, 12, 31, 23, 59,
59),
+ datetime(2026, 10, 3, 12,
30)]),
+ 'unicode': (pa.string(), ['a/b=%', '中文'])}
+ type_, values = types[kind]
+ schema = pa.schema([('id', pa.int32()), ('p', type_), ('v', pa.int32())])
+ table = _table(tmp_path, schema=schema)
+ plan = PostponeBucketPlan({(value,): i + 2 for i, value in
enumerate(values)})
+ builder =
table.new_postpone_fixed_bucket_write_builder().with_bucket_plan(plan)
+ writer, commit = builder.new_write(), builder.new_commit()
+ expected = [dict(id=i, p=value, v=i * 10) for i, value in
enumerate(values)]
+ try:
+ assert isinstance(writer, NativePostponeFixedBucketTableWrite)
+ writer.write_arrow(pa.Table.from_pylist(expected, schema=schema))
+ messages = writer.prepare_commit()
+ assert {m.partition: m.total_buckets for m in messages} ==
plan.as_dict()
+ commit.commit(messages)
+ for native in (False, True):
+ copied = table.copy({'read.native.enabled': str(native).lower(),
+ 'scan.native-plan.enabled':
str(native).lower()})
+ assert _rows(copied) == expected
+ finally:
+ writer.close()
+ commit.close()
+
+
[email protected]('native', [False, True])
+def test_reused_python_committer_drops_prior_fixed_bucket_baseline(tmp_path,
native):
+ table = _table(tmp_path, {'write.native.enabled': str(native).lower()})
+ first = table.new_postpone_fixed_bucket_write_builder()
+ writer = first.new_write()
+ try:
+ writer.write_arrow_batch(_batch([(1, 'a', 10)]))
+ messages = writer.prepare_commit()
+ # Exercise FileStoreCommit reuse directly: the public batch facade is
one-shot.
+ commit = first.new_commit().file_store_commit
+ commit.commit(messages, 1)
+ pending = table.new_batch_write_builder().new_write()
+ try:
+ pending.write_arrow_batch(_batch([(2, 'a', 20)]))
+ commit.commit(pending.prepare_commit(), 2)
+ finally:
+ pending.close()
+ assert commit.conflict_detection.fixed_bucket_commit_check is None
+ assert
table.snapshot_manager().get_latest_snapshot().total_record_count == 2
+ assert _rows(table) == [dict(id=1, p='a', v=10)]
+ finally:
+ writer.close()
diff --git a/paimon-python/pypaimon/tests/native_postpone_write_test.py
b/paimon-python/pypaimon/tests/native_postpone_write_test.py
new file mode 100644
index 0000000000..2c77b224ac
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_postpone_write_test.py
@@ -0,0 +1,186 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Native pending-bucket writes preserve Java replay order and file
metadata."""
+
+from unittest.mock import patch
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pytest
+
+from pypaimon import CatalogFactory, Schema
+from pypaimon.write.native_write import NativeTableWrite
+
+pytestmark = pytest.mark.native_plan
+_SCHEMA = pa.schema([('id', pa.int32()), ('p', pa.string()),
+ ('v', pa.int32()), ('op', pa.string())])
+
+
+def _table(tmp_path, options=None, catalog=None):
+ catalog = catalog or CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_database('default', True)
+ opts = {'bucket': '-2', 'file.format': 'parquet', 'rowkind.field': 'op',
+ 'target-file-row-num': '3', 'write.native.enabled': 'true'}
+ opts.update(options or {})
+ catalog.create_table('default.t', Schema.from_pyarrow_schema(
+ _SCHEMA, primary_keys=['id', 'p'], partition_keys=['p'],
options=opts), False)
+ return catalog.get_table('default.t')
+
+
+def _batch(rows):
+ return pa.RecordBatch.from_pylist([dict(zip(_SCHEMA.names, row)) for row
in rows], schema=_SCHEMA)
+
+
+def _file_rows(table, files):
+ rows = []
+ for file in files:
+ # Inspect the unsorted physical file, before the deferred compaction.
+ with table.file_io.new_input_stream(file.file_path) as source:
+ values = pq.ParquetFile(source).read().to_pylist()
+ assert len(values) == file.row_count
+ assert sum(row['_VALUE_KIND'] in (1, 3) for row in values) ==
file.delete_row_count
+ assert [row['_SEQUENCE_NUMBER'] for row in values] == list(range(
+ file.min_sequence_number, file.max_sequence_number + 1))
+ rows.extend(tuple(row[name] for name in _SCHEMA.names) for row in
values)
+ return rows
+
+
[email protected]('native_commit', [False, True])
[email protected]('stream', [False, True])
+def test_postpone_native_files_roll_and_preserve_replay_order(
+ tmp_path, native_rest_catalog, native_commit, stream):
+ table = _table(tmp_path, {'commit.native.enabled':
str(native_commit).lower()}, native_rest_catalog)
+ builder = table.new_stream_write_builder() if stream else
table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ assert isinstance(writer, NativeTableWrite)
+ try:
+ for checkpoint in range(1, 3 if stream else 2):
+ data = [(5, 'a/b', 10, '+I'), (1, 'a/b', 20, '-U'),
+ (3, 'a/b', 30, '+U'), (2, 'a/b', 40, '-D'), (1, 'a/b', 50,
'+I')]
+ for start in (0, 2, 4):
+ writer.write_arrow_batch(_batch(data[start:start + 2]))
+ messages = writer.prepare_commit(checkpoint) if stream else
writer.prepare_commit()
+ assert writer._python_writer is None
+ assert {m.bucket for m in messages} == {-2}
+ files = [f for m in messages for f in m.new_files]
+ assert [f.row_count for f in files] == [4, 1]
+ assert all('-u-' + builder.commit_user + '-s-' in f.file_name for
f in files)
+ assert [(f.min_key.values, f.max_key.values) for f in files] ==
[([5], [2]), ([1], [1])]
+ assert _file_rows(table, files) == data
+ if native_commit:
+ with patch.object(commit.file_store_commit, 'commit',
+ side_effect=AssertionError('Python commit
fallback')):
+ commit.commit(messages, checkpoint) if stream else
commit.commit(messages)
+ else:
+ commit.commit(messages, checkpoint) if stream else
commit.commit(messages)
+ if stream:
+ assert writer.prepare_commit(3) == []
+ finally:
+ writer.close()
+ commit.close()
+ # Postpone records become visible only after assignment to real buckets.
+ for native in (False, True):
+ read = table.copy({'scan.native-plan.enabled': str(native).lower(),
+ 'read.native.enabled':
str(native).lower()}).new_read_builder()
+ assert read.new_scan().plan().splits() == []
+
+
[email protected]('native_commit', [False, True])
+def test_postpone_native_overwrite_and_abort(tmp_path, native_rest_catalog,
native_commit):
+ table = _table(tmp_path, {'commit.native.enabled':
str(native_commit).lower()}, native_rest_catalog)
+ all_files = []
+ for overwrite in (False, True):
+ builder = table.new_batch_write_builder()
+ if overwrite:
+ builder.overwrite()
+ writer, commit = builder.new_write(), builder.new_commit()
+ assert isinstance(writer, NativeTableWrite)
+ try:
+ writer.write_arrow_batch(_batch([(1, 'a', 20 if overwrite else 10,
'+I')]))
+ messages = writer.prepare_commit()
+ files = [f for m in messages for f in m.new_files]
+ assert _file_rows(table, files) == [(1, 'a', 20 if overwrite else
10, '+I')]
+ all_files.extend(files)
+ if native_commit:
+ with patch.object(commit.file_store_commit, 'commit',
+ side_effect=AssertionError('Python commit
fallback')):
+ commit.commit(messages)
+ else:
+ commit.commit(messages)
+ writer.abort()
+ assert all(table.file_io.exists(f.file_path) for f in files)
+ finally:
+ writer.close()
+ commit.close()
+ assert table.snapshot_manager().get_latest_snapshot().total_record_count
== 1
+ assert len({f.file_name.split('-w-')[0] for f in all_files}) == 2
+ writer = table.new_batch_write_builder().new_write()
+ try:
+ writer.write_arrow_batch(_batch([(2, 'b', 30, '+I')]))
+ messages = writer.prepare_commit()
+ paths = [f.file_path for m in messages for f in m.new_files]
+ writer.abort()
+ assert paths and all(not table.file_io.exists(path) for path in paths)
+ finally:
+ writer.close()
+
+
[email protected]('engine', ['first-row', 'partial-update',
'aggregation'])
+def test_postpone_engine_retract_validation_stays_on_python(tmp_path, engine):
+ table = _table(tmp_path, {'merge-engine': engine})
+ writer = table.new_batch_write_builder().new_write()
+ try:
+ assert not isinstance(writer, NativeTableWrite)
+ finally:
+ writer.close()
+
+
[email protected]('operation', ['truncate_table', 'truncate_partition',
'delete_partition'])
+def test_pending_files_participate_in_partition_metadata_deletes(tmp_path,
operation):
+ table = _table(tmp_path, {'commit.native.enabled': 'false'})
+ builder = table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ try:
+ writer.write_arrow_batch(_batch([(1, 'a', 10, '+I'), (2, 'b', 20,
'+I')]))
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+ builder = table.new_batch_write_builder()
+ commit = builder.new_commit()
+ try:
+ if operation == 'truncate_table':
+ commit.truncate_table()
+ elif operation == 'truncate_partition':
+ commit.truncate_partitions([{'p': 'a'}])
+ else:
+ predicate =
table.new_read_builder().new_predicate_builder().equal('p', 'a')
+ messages = builder.new_update().delete_by_predicate(predicate)
+ assert len(messages) == 1
+ assert messages[0].bucket == -2
+ assert len(messages[0].deleted_files) == 1
+ commit.commit(messages)
+ finally:
+ commit.close()
+ expected = 0 if operation == 'truncate_table' else 1
+ assert table.snapshot_manager().get_latest_snapshot().total_record_count
== expected
+ from pypaimon.manifest.manifest_list_manager import ManifestListManager
+ from pypaimon.read.scanner.file_scanner import FileScanner
+ snapshot = table.snapshot_manager().get_latest_snapshot()
+ entries = FileScanner(table, lambda: ([],
None)).with_all_buckets().read_manifest_entries(
+ ManifestListManager(table).read_all(snapshot))
+ assert [(tuple(e.partition.values), e.bucket) for e in entries] == ([] if
expected == 0 else [(('b',), -2)])
diff --git a/paimon-python/pypaimon/tests/ray_repartition_test.py
b/paimon-python/pypaimon/tests/ray_repartition_test.py
index 6cbbc1053d..3b17784216 100644
--- a/paimon-python/pypaimon/tests/ray_repartition_test.py
+++ b/paimon-python/pypaimon/tests/ray_repartition_test.py
@@ -84,6 +84,10 @@ class RayShuffleTest(unittest.TestCase):
def _make_table(self, table_name, pa_schema, *, primary_keys=None,
partition_keys=None, options=None):
identifier = 'default.{}'.format(table_name)
+ options = dict(options or {})
+ # Pytest's patched defaults do not propagate to Ray worker processes.
+ options.setdefault('write.native.enabled', str(
+ os.environ.get('PYPAIMON_TEST_NATIVE_WRITE') == '1').lower())
schema = Schema.from_pyarrow_schema(
pa_schema,
primary_keys=primary_keys,
@@ -384,7 +388,7 @@ class RayShuffleTest(unittest.TestCase):
)
self.assertEqual(len(self._read_table(identifier)), 100)
- def test_postpone_large_bucket_spans_writer_blocks(self):
+ def test_postpone_large_bucket_keeps_one_writer_across_input_blocks(self):
from ray.data import DataContext
from pypaimon.ray import write_paimon
@@ -411,8 +415,10 @@ class RayShuffleTest(unittest.TestCase):
previous_target = context.target_max_block_size
context.target_max_block_size = 64 * 1024
try:
+ dataset = ray.data.from_arrow(rows).repartition(8).materialize()
+ self.assertGreater(dataset.num_blocks(), 1)
write_paimon(
- ray.data.from_arrow(rows).repartition(8),
+ dataset,
identifier,
self.catalog_options,
concurrency=4,
@@ -421,12 +427,14 @@ class RayShuffleTest(unittest.TestCase):
context.target_max_block_size = previous_target
files = self._count_data_files(table_name)
- self.assertGreater(len(files), 1)
+ self.assertTrue(files)
self.assertEqual(
{'bucket-0'},
{os.path.basename(os.path.dirname(path)) for path in files},
)
- self.assertEqual(400, len(self._read_table(identifier)))
+ result = self._read_table(identifier)
+ self.assertEqual(list(range(400)), sorted(result['id'].tolist()))
+ self.assertEqual({'x' * 4096}, set(result['value']))
def test_explicit_postpone_writer_uses_real_buckets(self):
pa_schema = pa.schema([
@@ -598,16 +606,26 @@ class RayShuffleTest(unittest.TestCase):
for data_file in split.files
]
self.assertEqual(len(data_files), 1)
- self.assertEqual(
- (data_files[0].min_sequence_number,
- data_files[0].max_sequence_number),
- (1, 801),
- )
+ self.assertEqual(data_files[0].row_count, 1)
result = self._read_table(identifier)
self.assertEqual(len(result), 1)
self.assertTrue(result.iloc[0]['value'].startswith('0799-'))
+ # A subsequent writer must continue after the persisted sequence,
+ # including when native buffering removed the earlier duplicate rows.
+ write_paimon(
+ ray.data.from_arrow(pa.Table.from_pydict(
+ {'id': [1], 'value': ['updated']}, schema=pa_schema)),
+ identifier,
+ self.catalog_options,
+ hash_fixed_precluster='map_groups',
+ )
+ self.assertEqual(
+ self._read_table(identifier).to_dict('records'),
+ [{'id': 1, 'value': 'updated'}],
+ )
+
def test_primary_key_group_writer_overwrite(self):
from pypaimon.ray import write_paimon
diff --git a/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py
b/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py
index e32d5f5a72..1ce9541e8c 100644
--- a/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py
+++ b/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py
@@ -379,6 +379,7 @@ class DynamicBucketTest(unittest.TestCase):
):
new_writer.write_arrow(pa.table({'id': [1], 'value': ['new']}))
+ @pytest.mark.python_write
def test_cross_partition_write_requires_global_index(self):
with tempfile.TemporaryDirectory() as root:
catalog = CatalogFactory.create({'warehouse': root})
diff --git a/paimon-python/pypaimon/tests/write/table_write_test.py
b/paimon-python/pypaimon/tests/write/table_write_test.py
index ee0ee42fd9..8006269bea 100644
--- a/paimon-python/pypaimon/tests/write/table_write_test.py
+++ b/paimon-python/pypaimon/tests/write/table_write_test.py
@@ -862,16 +862,12 @@ class TableWriteTest(unittest.TestCase):
'dt': ['p1'],
}, schema=self.pk_pa_schema)
- self._commit_arrow(table, expected)
-
- self.assertEqual(
- 1,
- len(glob.glob(
- self.warehouse
- + "/default.db/test_postpone_default_builder/user_id=1/"
- + "bucket-postpone/*.avro"
- )),
- )
+ messages = self._commit_arrow(table, expected)
+ self.assertEqual({-2}, {message.bucket for message in messages})
+ files = [file for message in messages for file in message.new_files]
+ self.assertEqual(1, len(files))
+ self.assertIn('/bucket-postpone/', files[0].file_path)
+ self.assertTrue(table.file_io.exists(files[0].file_path))
splits = table.new_read_builder().new_scan().plan().splits()
self.assertTrue(not
table.new_read_builder().new_read().to_arrow(splits))
@@ -1752,7 +1748,8 @@ class TableWriteTest(unittest.TestCase):
# Verify file name format:
{table_prefix}-u-{commit_user}-s-{random_number}-w--{uuid}-0.{format}
# Expected pattern: data--u-{user}-s-{random}-w--{uuid}-0.{format}
- expected_pattern = r'^data--u-.+-s-\d+-w-.+-0\.avro$'
+ # Native postpone writes use Parquet; the Python writer may select
Avro.
+ expected_pattern = r'^data--u-.+-s-\d+-w-.+-0\.(avro|parquet)$'
for file_name in data_files:
self.assertRegex(file_name, expected_pattern,
diff --git a/paimon-python/pypaimon/write/commit/commit_scanner.py
b/paimon-python/pypaimon/write/commit/commit_scanner.py
index 835df4bf61..77a7c9c74b 100644
--- a/paimon-python/pypaimon/write/commit/commit_scanner.py
+++ b/paimon-python/pypaimon/write/commit/commit_scanner.py
@@ -72,7 +72,7 @@ class CommitScanner:
all_manifests = self.manifest_list_manager.read_all(latest_snapshot)
return FileScanner(
self.table, lambda: ([], None),
partition_predicate=partition_filter
- ).read_manifest_entries(all_manifests)
+ ).with_all_buckets().read_manifest_entries(all_manifests)
def read_incremental_entries_from_changed_partitions(self,
snapshot: Snapshot,
@@ -101,7 +101,7 @@ class CommitScanner:
return FileScanner(
self.table, lambda: ([], None),
partition_predicate=partition_filter
- ).read_manifest_entries(delta_manifests)
+ ).with_all_buckets().read_manifest_entries(delta_manifests)
def read_incremental_raw_entries_from_changed_partitions(self, snapshot:
Snapshot,
commit_entries:
List[ManifestEntry],
diff --git a/paimon-python/pypaimon/write/commit/conflict_detection.py
b/paimon-python/pypaimon/write/commit/conflict_detection.py
index f15917c3c2..19d831cf60 100644
--- a/paimon-python/pypaimon/write/commit/conflict_detection.py
+++ b/paimon-python/pypaimon/write/commit/conflict_detection.py
@@ -192,6 +192,7 @@ class ConflictDetection:
self.manifest_list_manager = manifest_list_manager
self.table = table
self._row_id_check_from_snapshot = None
+ self.fixed_bucket_commit_check = None
self.commit_scanner = commit_scanner
def should_be_overwrite_commit(self, append_file_entries=None,
append_index_files=None):
@@ -248,6 +249,13 @@ class ConflictDetection:
if conflict is not None:
return conflict
+ if self.fixed_bucket_commit_check is not None:
+ conflict = self.fixed_bucket_commit_check.check(
+ latest_snapshot, delta_entries, commit_kind,
+ self.snapshot_manager, self.commit_scanner)
+ if conflict is not None:
+ return conflict
+
conflict = self.check_overwrite_from_snapshot(
latest_snapshot, delta_entries, commit_kind)
if conflict is not None:
diff --git a/paimon-python/pypaimon/write/commit/fixed_bucket_commit_check.py
b/paimon-python/pypaimon/write/commit/fixed_bucket_commit_check.py
new file mode 100644
index 0000000000..322da1eb16
--- /dev/null
+++ b/paimon-python/pypaimon/write/commit/fixed_bucket_commit_check.py
@@ -0,0 +1,72 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Conflict checks for batches that assign real buckets to postpone tables."""
+
+
+class FixedBucketCommitCheck:
+ """Retain the native writer's single-owner requirement across
serialization.
+
+ A writer seeds automatic sequence numbers from its baseline snapshot. Two
+ owners of the same bucket can overlap those sequences and silently lose an
+ update, even when their total bucket counts agree. Pending files do not own
+ a real bucket and can coexist, as in Java's postpone writes.
+ """
+
+ def __init__(self, messages):
+ self.owners = set()
+ baselines = []
+ for message in messages:
+ count = message.total_buckets
+ if count is None:
+ continue
+ if (isinstance(count, bool) or not isinstance(count, int)
+ or not 0 < count <= 2147483647
+ or not 0 <= message.bucket < count):
+ raise ValueError('Invalid fixed bucket {} for
total_buckets={}'.format(message.bucket, count))
+ if not message.new_files:
+ continue
+ owner = (tuple(message.partition), message.bucket)
+ if owner in self.owners:
+ raise ValueError('Postpone fixed-bucket writer ownership
conflict: '
+ 'one partition and bucket must have a single
writer')
+ self.owners.add(owner)
+ if message.check_from_snapshot is not None:
+ if message.check_from_snapshot < 0:
+ raise ValueError('Invalid fixed-bucket check snapshot')
+ baselines.append(message.check_from_snapshot)
+ self.baseline = min(baselines) if baselines else None
+
+ def check(self, latest, entries, commit_kind, snapshot_manager, scanner):
+ if (not self.owners or self.baseline is None or latest is None
+ or latest.id <= self.baseline):
+ return None
+ partitions = {partition for partition, _ in self.owners}
+ for snapshot_id in range(self.baseline + 1, latest.id + 1):
+ snapshot = snapshot_manager.get_snapshot_by_id(snapshot_id)
+ if snapshot is None:
+ return RuntimeError('Fixed-bucket conflict check cannot find
snapshot {}'.format(snapshot_id))
+ changes =
scanner.read_incremental_raw_entries_from_changed_partitions(snapshot, entries)
+ for change in changes:
+ partition = tuple(change.partition.values)
+ if commit_kind == 'OVERWRITE' and partition in partitions:
+ return RuntimeError('Postpone fixed-bucket overwrite
conflict: '
+ 'target partition changed after
snapshot {}'.format(self.baseline))
+ if change.kind == 0 and (partition, change.bucket) in
self.owners:
+ return RuntimeError('Postpone fixed-bucket writer
ownership conflict: '
+ 'bucket changed after snapshot
{}'.format(self.baseline))
+ return None
diff --git a/paimon-python/pypaimon/write/commit/overwrite_changes_provider.py
b/paimon-python/pypaimon/write/commit/overwrite_changes_provider.py
index 39460acc61..94241aaed7 100644
--- a/paimon-python/pypaimon/write/commit/overwrite_changes_provider.py
+++ b/paimon-python/pypaimon/write/commit/overwrite_changes_provider.py
@@ -75,6 +75,7 @@ class OverwriteChangesProvider:
self.full_scan_count += 1
return (FileScanner(self.table, lambda: ([], None),
partition_predicate=self.partition_filter)
+ .with_all_buckets()
.read_manifest_entries(self.manifest_list_manager.read_all(latest_snapshot)))
def _update_cache(self, latest_snapshot: Snapshot) -> bool:
diff --git a/paimon-python/pypaimon/write/file_store_commit.py
b/paimon-python/pypaimon/write/file_store_commit.py
index d7396d367a..7f45040246 100644
--- a/paimon-python/pypaimon/write/file_store_commit.py
+++ b/paimon-python/pypaimon/write/file_store_commit.py
@@ -222,6 +222,14 @@ class FileStoreCommit:
table_rollback = table.catalog_environment.catalog_table_rollback()
self.rollback = CommitRollback(table_rollback) if table_rollback is
not None else None
+ def _set_fixed_bucket_commit_check(self, messages):
+ from pypaimon.write.commit.fixed_bucket_commit_check import
FixedBucketCommitCheck
+
+ self.conflict_detection.fixed_bucket_commit_check = (
+ FixedBucketCommitCheck(messages)
+ if any(message.total_buckets is not None for message in messages)
+ else None)
+
def commit(
self,
commit_messages: List[CommitMessage],
@@ -234,6 +242,7 @@ class FileStoreCommit:
_reject_compact_increment(commit_messages)
check_from_snapshot = _row_id_check_from_messages(commit_messages)
+ self._set_fixed_bucket_commit_check(commit_messages)
# A committer can be reused; an untagged commit clears the prior
baseline.
self.conflict_detection._row_id_check_from_snapshot =
check_from_snapshot
@@ -317,8 +326,9 @@ class FileStoreCommit:
snapshot_properties: Optional[Dict[str, str]] = None):
"""Commit the given commit messages in overwrite mode."""
_reject_compact_increment(commit_messages)
- self.conflict_detection._row_id_check_from_snapshot = (
- _row_id_check_from_messages(commit_messages))
+ check_from_snapshot = _row_id_check_from_messages(commit_messages)
+ self._set_fixed_bucket_commit_check(commit_messages)
+ self.conflict_detection._row_id_check_from_snapshot =
check_from_snapshot
logger.info(
"Ready to overwrite to table %s, number of commit messages: %d",
self.table.identifier,
diff --git a/paimon-python/pypaimon/write/file_store_write.py
b/paimon-python/pypaimon/write/file_store_write.py
index 28e9f8f67d..16f0884516 100644
--- a/paimon-python/pypaimon/write/file_store_write.py
+++ b/paimon-python/pypaimon/write/file_store_write.py
@@ -365,8 +365,11 @@ class FileStoreWrite:
self.max_seq_numbers[partition] = buckets
return buckets
+ def _sequence_read_table(self):
+ return self.table
+
def _load_seq_number_stats(self, partition: Tuple) -> dict:
- read_builder = self.table.new_read_builder()
+ read_builder = self._sequence_read_table().new_read_builder()
predicate_builder = read_builder.new_predicate_builder()
sub_predicates = []
for key, value in zip(self.table.partition_keys, partition):
@@ -388,5 +391,24 @@ class FileStoreWrite:
class PostponeFixedBucketFileStoreWrite(FileStoreWrite):
"""File store write with runtime bucket counts for postpone tables."""
+ def __init__(self, table, commit_user):
+ super().__init__(table, commit_user)
+ snapshot = table.snapshot_manager().get_latest_snapshot()
+ self._check_from_snapshot = snapshot.id if snapshot is not None else 0
+
+ def _sequence_read_table(self):
+ return self.table.copy({'scan.snapshot-id':
str(self._check_from_snapshot)})
+
+ def _load_seq_number_stats(self, partition):
+ if self._check_from_snapshot == 0:
+ return {}
+ return super()._load_seq_number_stats(partition)
+
+ def prepare_commit(self, commit_identifier):
+ messages = super().prepare_commit(commit_identifier)
+ for message in messages:
+ message.check_from_snapshot = self._check_from_snapshot
+ return messages
+
def _configure_data_file_prefix(self, commit_user):
pass
diff --git a/paimon-python/pypaimon/write/native_commit.py
b/paimon-python/pypaimon/write/native_commit.py
index 46de93bcb6..297442a22d 100644
--- a/paimon-python/pypaimon/write/native_commit.py
+++ b/paimon-python/pypaimon/write/native_commit.py
@@ -140,15 +140,13 @@ def create_native_write_table(table):
from pypaimon.filesystem.local_file_io import LocalFileIO
from pypaimon.filesystem.pyarrow_file_io import PyArrowFileIO
from pypaimon.filesystem.resolving_file_io import ResolvingFileIO
- from pypaimon.table.bucket_mode import BucketMode
from pypaimon.table.file_store_table import FileStoreTable
- # Native branch and postpone writes are not supported.
+ # Native writes currently require the main branch.
environment = table.catalog_environment
if (type(table) is not FileStoreTable
or type(environment) is not CatalogEnvironment
or table.current_branch() != 'main'
- or table.bucket_mode() == BucketMode.POSTPONE_MODE
or table.options.query_auth_enabled):
return None
diff --git a/paimon-python/pypaimon/write/native_write.py
b/paimon-python/pypaimon/write/native_write.py
index ca56bd6f7e..d0ca273479 100644
--- a/paimon-python/pypaimon/write/native_write.py
+++ b/paimon-python/pypaimon/write/native_write.py
@@ -60,7 +60,8 @@ def _native_map_layouts_supported(table, schema):
return True
-def create_native_write(table, commit_user, static_partition=None,
stream=False):
+def create_native_write(table, commit_user, static_partition=None,
stream=False,
+ *, fixed_bucket=False, bucket_plan=None):
"""Return a native writer if the table can use the filesystem write
path."""
schema = PyarrowFieldParser.from_paimon_schema(table.table_schema.fields)
sequence_fields = table.options.sequence_field()
@@ -72,7 +73,17 @@ def create_native_write(table, commit_user,
static_partition=None, stream=False)
or (table.options.data_evolution_enabled()
and table.options.data_evolution_row_sidecar_enabled())
or table.bucket_mode() not in (BucketMode.HASH_FIXED,
BucketMode.HASH_DYNAMIC,
- BucketMode.BUCKET_UNAWARE)
+ BucketMode.BUCKET_UNAWARE,
+ BucketMode.CROSS_PARTITION,
+ BucketMode.POSTPONE_MODE)
+ # Postpone bypasses the KV merger. Native retract validation for
+ # other engines and managed BLOB externalization are still missing.
+ or (table.bucket_mode() == BucketMode.POSTPONE_MODE and not
fixed_bucket
+ and (table.options.merge_engine() != MergeEngine.DEDUPLICATE
+ or any(is_blob_file_field(field) for field in
table.table_schema.fields)))
+ or (fixed_bucket and (table.options.deletion_vectors_enabled()
+ or any(is_blob_file_field(field)
+ for field in
table.table_schema.fields)))
or (table.options.deletion_vectors_enabled()
and table.options.merge_engine() in
(MergeEngine.PARTIAL_UPDATE,
MergeEngine.AGGREGATE))
@@ -90,6 +101,14 @@ def create_native_write(table, commit_user,
static_partition=None, stream=False)
native_table = create_native_write_table(table)
if native_table is None:
return None
+ if fixed_bucket:
+ builder =
native_table.new_postpone_fixed_bucket_write_builder()._with_commit_user(commit_user)
+ if bucket_plan is not None:
+ builder = builder.with_bucket_plan(bucket_plan.to_arrow(table))
+ if static_partition is not None:
+ builder = builder.with_overwrite(static_partition)
+ return NativePostponeFixedBucketTableWrite(
+ table, commit_user, static_partition, bucket_plan,
builder.new_write())
if stream:
builder =
native_table.new_stream_write_builder().with_commit_user(commit_user)
else:
@@ -127,13 +146,14 @@ class NativeTableWrite:
'Cannot switch to the Python writer after native data was
written')
self._native_writer.close()
self._native_writer = None
+ self._python_writer = self._new_python_writer()
+ return self._python_writer
+
+ def _new_python_writer(self):
from pypaimon.write.table_write import BatchTableWrite,
StreamTableWrite
if self.stream:
- self._python_writer = StreamTableWrite(self.table,
self.commit_user)
- else:
- self._python_writer = BatchTableWrite(
- self.table, self.commit_user, self.static_partition)
- return self._python_writer
+ return StreamTableWrite(self.table, self.commit_user)
+ return BatchTableWrite(self.table, self.commit_user,
self.static_partition)
def __getattr__(self, name):
if name.startswith('_'):
@@ -222,3 +242,20 @@ class NativeTableWrite:
self.close()
finally:
_abort_commit_messages(self.table, messages)
+
+
+class NativePostponeFixedBucketTableWrite(NativeTableWrite):
+ """Fixed-bucket batches retain the standard native ownership contract.
+
+ A distributed coordinator can use Python planning before receiving data;
+ workers use the same shared plan through the Rust fixed-bucket writer.
+ """
+
+ def __init__(self, table, commit_user, static_partition, bucket_plan,
native_writer):
+ super().__init__(table, commit_user, static_partition, False,
native_writer)
+ self._bucket_plan = bucket_plan
+
+ def _new_python_writer(self):
+ from pypaimon.write.postpone_batch_table_write import
PostponeFixedBucketBatchTableWrite
+ return PostponeFixedBucketBatchTableWrite(
+ self.table, self.commit_user, self.static_partition,
self._bucket_plan)
diff --git a/paimon-python/pypaimon/write/postpone_batch_table_write.py
b/paimon-python/pypaimon/write/postpone_batch_table_write.py
index f4fc24f70b..25227a06bf 100644
--- a/paimon-python/pypaimon/write/postpone_batch_table_write.py
+++ b/paimon-python/pypaimon/write/postpone_batch_table_write.py
@@ -50,6 +50,14 @@ class PostponeFixedBucketWriteBuilder(BatchWriteBuilder):
return self
def new_write(self):
+ # Automatic size estimation still belongs to the Python coordinator.
+ # Native direct writes accept a resolved plan or Java's explicit
default.
+ if (self._bucket_plan is not None
+ or self.table.options.postpone_default_bucket_num() is not
None):
+ native = self._native_write(
+ self.static_partition, fixed_bucket=True,
bucket_plan=self._bucket_plan)
+ if native is not None:
+ return native
return PostponeFixedBucketBatchTableWrite(
self.table,
self.commit_user,
@@ -71,7 +79,8 @@ class PostponeFixedBucketBatchTableWrite(BatchTableWrite):
self._planner = PostponeBucketPlanner(
table,
known_num_buckets=(
- bucket_plan.as_dict() if bucket_plan is not None else None
+ bucket_plan.as_dict() if bucket_plan is not None
+ else {} if static_partition is not None else None
),
)
self._bucket_plan = (
diff --git a/paimon-python/pypaimon/write/postpone_bucket.py
b/paimon-python/pypaimon/write/postpone_bucket.py
index 95b86c6ce1..ca6c9a8100 100644
--- a/paimon-python/pypaimon/write/postpone_bucket.py
+++ b/paimon-python/pypaimon/write/postpone_bucket.py
@@ -17,6 +17,8 @@
from typing import Dict, Tuple
+import pyarrow as pa
+
from pypaimon.schema.data_types import (
ArrayType,
AtomicType,
@@ -178,6 +180,26 @@ class PostponeBucketPlan:
def __init__(self, num_buckets: Dict[Tuple, int]):
self._num_buckets = dict(num_buckets)
+ def to_arrow(self, table) -> pa.RecordBatch:
+ """Encode only the per-partition plan for the native core router."""
+ from pypaimon.schema.data_types import PyarrowFieldParser
+
+ fields = [table.field_dict[name] for name in table.partition_keys]
+ schema = PyarrowFieldParser.from_paimon_schema(fields)
+ schema = schema.append(pa.field('total_buckets', pa.int32(),
nullable=False))
+ partitions = list(self._num_buckets)
+ counts = list(self._num_buckets.values())
+ if any(len(partition) != len(fields) for partition in partitions):
+ raise ValueError('Bucket plan partition arity does not match the
table schema')
+ if any(isinstance(count, bool) or not isinstance(count, int)
+ or not 0 < count <= 2147483647 for count in counts):
+ raise ValueError('Bucket plan total_buckets must be a positive
32-bit integer')
+ # Use positional columns: a partition may itself be named
total_buckets.
+ columns = [pa.array([partition[i] for partition in partitions],
type=field.type)
+ for i, field in enumerate(list(schema)[:-1])]
+ columns.append(pa.array(counts, type=pa.int32()))
+ return pa.RecordBatch.from_arrays(columns, schema=schema)
+
def contains(self, partition: Tuple) -> bool:
return tuple(partition) in self._num_buckets
@@ -216,33 +238,37 @@ class PostponeBucketPlanner:
.format(bucket_function)
)
- self.max_num_buckets = (
- options.postpone_batch_write_fixed_bucket_max_parallelism()
- )
- if self.max_num_buckets <= 0:
- raise ValueError(
- "postpone.batch-write-fixed-bucket.max-parallelism must be "
- "positive, got {}".format(self.max_num_buckets)
+ self.default_bucket_num = options.postpone_default_bucket_num()
+ self.target_row_num_per_bucket = None
+ self.target_size_per_bucket = None
+ if self.default_bucket_num is None:
+ self.max_num_buckets = (
+ options.postpone_batch_write_fixed_bucket_max_parallelism()
)
- self.target_row_num_per_bucket = (
- options.postpone_target_row_num_per_bucket()
- )
- if self.target_row_num_per_bucket is not None:
- if self.target_row_num_per_bucket <= 0:
+ if self.max_num_buckets <= 0:
raise ValueError(
- "postpone.target-row-num-per-bucket must be positive, "
- "got {}".format(self.target_row_num_per_bucket)
+ "postpone.batch-write-fixed-bucket.max-parallelism must be
"
+ "positive, got {}".format(self.max_num_buckets)
)
- self.target_size_per_bucket = None
- else:
- self.target_size_per_bucket = (
- options.postpone_target_size_per_bucket()
+ self.target_row_num_per_bucket = (
+ options.postpone_target_row_num_per_bucket()
)
- if self.target_size_per_bucket <= 0:
- raise ValueError(
- "postpone.target-size-per-bucket must be positive, got "
- "{}".format(self.target_size_per_bucket)
+ if self.target_row_num_per_bucket is not None:
+ if self.target_row_num_per_bucket <= 0:
+ raise ValueError(
+ "postpone.target-row-num-per-bucket must be positive, "
+ "got {}".format(self.target_row_num_per_bucket)
+ )
+ self.target_size_per_bucket = None
+ else:
+ self.target_size_per_bucket = (
+ options.postpone_target_size_per_bucket()
)
+ if self.target_size_per_bucket <= 0:
+ raise ValueError(
+ "postpone.target-size-per-bucket must be positive, got
"
+ "{}".format(self.target_size_per_bucket)
+ )
self._partition_keys = list(table.partition_keys)
self._field_dict = dict(table.field_dict)
@@ -298,7 +324,8 @@ class PostponeBucketPlanner:
stats = {}
fields = [self._field_dict[name] for name in data.schema.names]
columns = [data.column(i) for i in range(len(fields))]
- collect_size = self.target_row_num_per_bucket is None
+ collect_size = (self.default_bucket_num is None
+ and self.target_row_num_per_bucket is None)
for row, partition in enumerate(partitions):
if partition in self._known_num_buckets:
continue
@@ -319,6 +346,9 @@ class PostponeBucketPlanner:
partition = tuple(partition)
if partition in self._known_num_buckets:
continue
+ if self.default_bucket_num is not None:
+ self._known_num_buckets[partition] = self.default_bucket_num
+ continue
postpone_rows = (
self._postpone_row_counts.get(partition, 0)
if include_postpone_rows
diff --git a/paimon-python/pypaimon/write/ray_datasink.py
b/paimon-python/pypaimon/write/ray_datasink.py
index 1fb56bbcac..796f69947e 100644
--- a/paimon-python/pypaimon/write/ray_datasink.py
+++ b/paimon-python/pypaimon/write/ray_datasink.py
@@ -331,7 +331,9 @@ def write_paimon_dataset(
planner = (
postpone_bucket_planner
if postpone_bucket_planner is not None
- else PostponeBucketPlanner(table)
+ else PostponeBucketPlanner(
+ table, known_num_buckets={}
+ if overwrite or static_partition is not None else None)
)
plan = planner.current_plan()
if table.partition_keys or not plan.contains(()):
@@ -344,7 +346,7 @@ def write_paimon_dataset(
overwrite or static_partition is not None
),
)
- return _write_postpone_primary_key_blocks(
+ return _write_primary_key_groups(
dataset,
table,
overwrite=overwrite,
@@ -383,80 +385,6 @@ def write_paimon_dataset(
return datasink.commit_result
-def _write_postpone_primary_key_blocks(
- dataset,
- table,
- *,
- overwrite: bool,
- static_partition: Optional[Dict[str, Any]],
- concurrency: Optional[int],
- ray_remote_args: Optional[Dict[str, Any]],
- bucket_extractor,
- postpone_bucket_plan,
-) -> Optional[PaimonWriteResult]:
- import pickle
-
- from pypaimon.ray.shuffle import (
- _sort_by_partition_bucket_primary_key,
- )
-
- sorted_dataset, routing_columns = (
- _sort_by_partition_bucket_primary_key(
- dataset, table, bucket_extractor
- )
- )
- message_col = "__paimon_commit_messages__"
- error_col = "__paimon_write_error__"
- captured_table = table
-
- def _write_block(batch: pa.Table) -> pa.Table:
- if batch.num_rows == 0:
- return pa.table({
- message_col: pa.array([], type=pa.binary()),
- error_col: pa.array([], type=pa.string()),
- })
-
- rows = batch.drop_columns(routing_columns)
- worker_sink = PaimonDatasink(
- captured_table,
- overwrite=overwrite,
- static_partition=static_partition,
- postpone_bucket_plan=postpone_bucket_plan,
- )
- try:
- commit_messages = worker_sink.write([rows], None)
- error = None
- except Exception:
- commit_messages = []
- error = traceback.format_exc()
- return pa.table({
- message_col: pa.array(
- [pickle.dumps(commit_messages)], type=pa.binary()
- ),
- error_col: pa.array([error], type=pa.string()),
- })
-
- map_kwargs = _ray_map_kwargs(
- sorted_dataset.map_batches,
- concurrency,
- ray_remote_args,
- batch_size=None,
- batch_format="pyarrow",
- zero_copy_batch=True,
- )
-
- results = sorted_dataset.map_batches(_write_block, **map_kwargs)
- coordinator = PaimonDatasink(
- table,
- overwrite=overwrite,
- static_partition=static_partition,
- )
- coordinator.on_write_start()
- return _consume_write_results(
- results, coordinator, message_col, error_col
- )
-
-
def _ray_map_kwargs(method, concurrency, ray_remote_args, **kwargs):
import inspect
diff --git a/paimon-python/pypaimon/write/table_update.py
b/paimon-python/pypaimon/write/table_update.py
index 38589a9215..e90b45195e 100644
--- a/paimon-python/pypaimon/write/table_update.py
+++ b/paimon-python/pypaimon/write/table_update.py
@@ -558,7 +558,7 @@ class TableUpdate:
self.table,
lambda: ([], None),
partition_predicate=partition_filter,
- ).read_manifest_entries(manifest_list_manager.read_all(snapshot))
+
).with_all_buckets().read_manifest_entries(manifest_list_manager.read_all(snapshot))
for entry in data_entries:
message = self._partition_delete_message(
diff --git a/paimon-python/pypaimon/write/write_builder.py
b/paimon-python/pypaimon/write/write_builder.py
index bd518d826d..4cd4f56d64 100644
--- a/paimon-python/pypaimon/write/write_builder.py
+++ b/paimon-python/pypaimon/write/write_builder.py
@@ -58,7 +58,7 @@ class WriteBuilder(ABC):
else:
return str(uuid.uuid4())
- def _native_write(self, static_partition=None, stream=False):
+ def _native_write(self, static_partition=None, stream=False, **kwargs):
from pypaimon.read.merge_engine_support import
check_sequence_field_supported
# Keep invalid configurations outside the native fallback handler and
@@ -69,7 +69,7 @@ class WriteBuilder(ABC):
try:
from pypaimon.write.native_write import create_native_write
return create_native_write(self.table, self.commit_user,
- static_partition, stream)
+ static_partition, stream, **kwargs)
except Exception as error:
# Construction has not written any data; the normal writer is safe.
logger.debug('Native writer preparation failed; using Python: %s',
error)
