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 3f297583a9 [Python] Keep grouped bucket writes on Native (#10412)
3f297583a9 is described below
commit 3f297583a936d52150224ea66c420e44c80672e3
Author: Jingsong Lee <[email protected]>
AuthorDate: Wed Oct 7 16:39:01 2026 +0800
[Python] Keep grouped bucket writes on Native (#10412)
---
paimon-python/conftest.py | 21 +-
paimon-python/pypaimon/daft/daft_datasink.py | 47 +--
paimon-python/pypaimon/daft/daft_paimon.py | 13 -
paimon-python/pypaimon/index/dynamic_bucket.py | 37 +-
paimon-python/pypaimon/table/file_store_table.py | 7 +-
paimon-python/pypaimon/tests/arrow_schema_test.py | 2 +-
.../tests/daft/daft_pk_distributed_write_test.py | 57 +++
.../tests/native_precomputed_bucket_test.py | 420 +++++++++++++++++++++
.../pypaimon/tests/write/dynamic_bucket_test.py | 88 ++---
paimon-python/pypaimon/write/file_store_write.py | 13 +-
paimon-python/pypaimon/write/native_write.py | 44 ++-
.../pypaimon/write/postpone_batch_table_write.py | 8 +-
paimon-python/pypaimon/write/row_key_extractor.py | 61 +--
paimon-python/pypaimon/write/table_write.py | 114 ++----
paimon-python/pypaimon/write/write_builder.py | 20 +-
paimon-python/pypaimon/write/writer/data_writer.py | 11 +-
16 files changed, 681 insertions(+), 282 deletions(-)
diff --git a/paimon-python/conftest.py b/paimon-python/conftest.py
index 7aa11d1fc6..4288d9b997 100644
--- a/paimon-python/conftest.py
+++ b/paimon-python/conftest.py
@@ -130,17 +130,20 @@ def pytest_configure(config):
if _native_write_enabled():
from pypaimon.write.native_write import NativeTableWrite
- original_write = NativeTableWrite.write_arrow_batch
+ def track_write(method):
+ original_write = getattr(NativeTableWrite, method)
- def tracked_write(self, data):
- global _native_write_count
- native = self._native_writer is not None
- result = original_write(self, data)
- if native and data.num_rows and _force_native_write_for_test:
- _native_write_count += 1
- return result
+ def tracked_write(self, data, *args, **kwargs):
+ global _native_write_count
+ native = self._native_writer is not None
+ result = original_write(self, data, *args, **kwargs)
+ if native and data.num_rows and _force_native_write_for_test:
+ _native_write_count += 1
+ return result
- NativeTableWrite.write_arrow_batch = tracked_write
+ setattr(NativeTableWrite, method, tracked_write)
+
+ track_write('write_arrow_batch')
if _native_commit_enabled():
from pypaimon.write.table_commit import TableCommit
diff --git a/paimon-python/pypaimon/daft/daft_datasink.py
b/paimon-python/pypaimon/daft/daft_datasink.py
index 1376eebe2b..9bf0dbe8ad 100644
--- a/paimon-python/pypaimon/daft/daft_datasink.py
+++ b/paimon-python/pypaimon/daft/daft_datasink.py
@@ -519,7 +519,6 @@ def make_dynamic_bucket_assignment_udf(
num_assigners: int,
bucket_column: str,
key_hash_column: str,
- new_mapping_column: str,
base_snapshot_column: str,
ignore_existing: bool = False,
base_snapshot_id: int | None = None,
@@ -549,9 +548,7 @@ def make_dynamic_bucket_assignment_udf(
field.name: DataType.from_arrow_type(field.type)
for field in output_fields
}
- return_fields[key_hash_column] = DataType.int32()
return_fields[bucket_column] = DataType.int32()
- return_fields[new_mapping_column] = DataType.bool()
return_fields[base_snapshot_column] = DataType.int64()
@daft.func.batch(
@@ -566,9 +563,7 @@ def make_dynamic_bucket_assignment_udf(
pa.array([], type=pa.struct([
pa.field(field.name, field.type) for field in output_fields
] + [
- pa.field(key_hash_column, pa.int32()),
pa.field(bucket_column, pa.int32()),
- pa.field(new_mapping_column, pa.bool_()),
pa.field(base_snapshot_column, pa.int64()),
]))
)
@@ -598,25 +593,21 @@ def make_dynamic_bucket_assignment_udf(
base_snapshot_id=base_snapshot_id,
)
buckets = []
- new_mappings = []
offset = 0
for batch in arrow_table.to_batches():
batch_hashes = key_hash_values[offset:offset + batch.num_rows]
- _, batch_buckets, batch_new_mappings = (
+ _, batch_buckets, _ = (
extractor.extract_partition_bucket_status_from_hashes_batch(
batch, batch_hashes
)
)
buckets.extend(batch_buckets)
- new_mappings.extend(batch_new_mappings)
offset += batch.num_rows
arrays = [
arrow_table.column(field.name).combine_chunks()
for field in output_fields
]
- arrays.append(pa.array(key_hash_values, type=pa.int32()))
arrays.append(pa.array(buckets, type=pa.int32()))
- arrays.append(pa.array(new_mappings, type=pa.bool_()))
arrays.append(pa.array(
[extractor.base_snapshot_id] * arrow_table.num_rows,
type=pa.int64(),
@@ -625,9 +616,7 @@ def make_dynamic_bucket_assignment_udf(
pa.StructArray.from_arrays(
arrays,
names=[field.name for field in output_fields] + [
- key_hash_column,
bucket_column,
- new_mapping_column,
base_snapshot_column,
],
)
@@ -667,8 +656,6 @@ def make_group_write_udf(
def _write_group(*columns) -> Series:
data_columns = columns[:data_column_count]
bucket = None
- key_hash_values = None
- new_mapping_values = None
group_base_snapshot_id = base_snapshot_id
if precomputed_bucket:
bucket_values = columns[data_column_count].to_pylist()
@@ -680,11 +667,7 @@ def make_group_write_udf(
"A Paimon write group contained multiple buckets"
)
if dynamic_bucket:
- key_hash_values = columns[data_column_count + 1].to_pylist()
- new_mapping_values = columns[data_column_count + 2].to_pylist()
- base_snapshot_values = columns[
- data_column_count + 3
- ].to_pylist()
+ base_snapshot_values = columns[data_column_count +
1].to_pylist()
observed_base_snapshot_id = base_snapshot_values[0]
if any(
value != observed_base_snapshot_id
@@ -714,33 +697,13 @@ def make_group_write_udf(
write_builder.commit_user = commit_user
if mode == "overwrite":
write_builder.overwrite({})
- table_write = write_builder.new_write()
if precomputed_bucket and dynamic_bucket:
- table_write.with_dynamic_bucket_index(
- ignore_existing=mode == "overwrite",
- base_snapshot_id=group_base_snapshot_id,
- )
+ write_builder.with_restore_snapshot(group_base_snapshot_id)
+ table_write = write_builder.new_write()
commit_messages = []
try:
if precomputed_bucket:
- offset = 0
- for batch in arrow_table.to_batches():
- batch_hashes = None
- batch_new_mappings = None
- if key_hash_values is not None:
- batch_hashes = key_hash_values[
- offset:offset + batch.num_rows
- ]
- batch_new_mappings = new_mapping_values[
- offset:offset + batch.num_rows
- ]
- table_write.write_arrow_batch_to_bucket(
- batch,
- bucket,
- batch_hashes,
- batch_new_mappings,
- )
- offset += batch.num_rows
+ table_write.write_arrow(arrow_table, bucket=bucket)
else:
table_write.write_arrow(arrow_table)
commit_messages = table_write.prepare_commit()
diff --git a/paimon-python/pypaimon/daft/daft_paimon.py
b/paimon-python/pypaimon/daft/daft_paimon.py
index b294ede3cd..765c009a33 100644
--- a/paimon-python/pypaimon/daft/daft_paimon.py
+++ b/paimon-python/pypaimon/daft/daft_paimon.py
@@ -324,22 +324,12 @@ def _write_table(
],
"__paimon_key_hash__",
)
- new_mapping_column = _pick_internal_column(
- df.column_names + [
- commit_column,
- assigner_column,
- bucket_column,
- key_hash_column,
- ],
- "__paimon_new_mapping__",
- )
base_snapshot_column = _pick_internal_column(
df.column_names + [
commit_column,
assigner_column,
bucket_column,
key_hash_column,
- new_mapping_column,
],
"__paimon_base_snapshot__",
)
@@ -360,7 +350,6 @@ def _write_table(
num_assigners,
bucket_column,
key_hash_column,
- new_mapping_column,
base_snapshot_column,
ignore_existing=mode == "overwrite",
base_snapshot_id=base_snapshot_id,
@@ -388,8 +377,6 @@ def _write_table(
group_write(
*data_columns,
daft.col(bucket_column),
- daft.col(key_hash_column),
- daft.col(new_mapping_column),
daft.col(base_snapshot_column),
).alias(commit_column)
)
diff --git a/paimon-python/pypaimon/index/dynamic_bucket.py
b/paimon-python/pypaimon/index/dynamic_bucket.py
index a2838a9494..707b03e724 100644
--- a/paimon-python/pypaimon/index/dynamic_bucket.py
+++ b/paimon-python/pypaimon/index/dynamic_bucket.py
@@ -22,7 +22,7 @@ import random
import struct
import uuid
from dataclasses import dataclass
-from typing import Dict, Iterator, List, Optional, Set, Tuple
+from typing import Dict, Iterator, List, Set, Tuple
from pypaimon.index.index_file_handler import IndexFileHandler
from pypaimon.index.index_file_meta import IndexFileMeta
@@ -31,7 +31,6 @@ from pypaimon.table.row.generic_row import GenericRow
HASH_INDEX = "HASH"
_ADD = 0
-_DELETE = 1
SHORT_MAX_VALUE = 32767
MAX_DYNAMIC_BUCKETS = SHORT_MAX_VALUE + 1
_SNAPSHOT_UNSET = object()
@@ -425,7 +424,7 @@ class DynamicBucketIndexMaintainer:
self.snapshot.id if self.snapshot is not None else 0
)
self._states: Dict[
- Tuple[Tuple, int], Tuple[Set[int], Optional[IndexManifestEntry],
bool]
+ Tuple[Tuple, int], Tuple[Set[int], bool]
] = {}
self._new_paths: List[str] = []
@@ -435,33 +434,27 @@ class DynamicBucketIndexMaintainer:
key = (partition, bucket)
state = self._states.get(key)
if state is None:
- hashes, old_entry = self._load_bucket(partition, bucket)
- state = (hashes, old_entry, False)
- hashes, old_entry, modified = state
+ state = (self._load_bucket(partition, bucket), False)
+ hashes, modified = state
previous_size = len(hashes)
hashes.add(to_signed_int32(key_hash))
- self._states[key] = (hashes, old_entry, modified or len(hashes) !=
previous_size)
+ self._states[key] = (hashes, modified or len(hashes) != previous_size)
def prepare_commit(self) -> Dict[Tuple[Tuple, int],
DynamicBucketIndexChanges]:
changes: Dict[Tuple[Tuple, int], DynamicBucketIndexChanges] = {}
prepared_states = {}
- for key, (hashes, old_entry, modified) in self._states.items():
+ for key, (hashes, modified) in self._states.items():
if not modified:
continue
partition, bucket = key
new_entry = self._write_index(partition, bucket, hashes)
changes[key] = DynamicBucketIndexChanges(
additions=[new_entry],
- deletions=[
- IndexManifestEntry(
- kind=_DELETE,
- partition=old_entry.partition,
- bucket=old_entry.bucket,
- index_file=old_entry.index_file,
- )
- ] if old_entry is not None else [],
+ # Java sends only the full new HASH file. BucketedCombiner
+ # replaces the previous entry by partition and bucket.
+ deletions=[],
)
- prepared_states[key] = (hashes, new_entry, False)
+ prepared_states[key] = (hashes, False)
# A failed bucket must not make a retry forget earlier additions.
self._states.update(prepared_states)
return changes
@@ -477,7 +470,7 @@ class DynamicBucketIndexMaintainer:
def _load_bucket(
self, partition: Tuple, bucket: int
- ) -> Tuple[Set[int], Optional[IndexManifestEntry]]:
+ ) -> Set[int]:
"""Load a complete bucket index for full-file replacement.
This matches Java's HASH index lifecycle. A modified bucket therefore
@@ -485,10 +478,10 @@ class DynamicBucketIndexMaintainer:
format is introduced.
"""
if self.ignore_existing:
- return set(), None
+ return set()
snapshot = self.snapshot
if snapshot is None:
- return set(), None
+ return set()
entries = IndexFileHandler(self.table).scan(
snapshot,
lambda entry: entry.index_file.index_type == HASH_INDEX
@@ -502,8 +495,8 @@ class DynamicBucketIndexMaintainer:
f"{[entry.index_file.file_name for entry in entries]}"
)
if not entries:
- return set(), None
- return _read_hashes(self.table, entries[0]), entries[0]
+ return set()
+ return _read_hashes(self.table, entries[0])
def _write_index(
self, partition: Tuple, bucket: int, hashes: Set[int]
diff --git a/paimon-python/pypaimon/table/file_store_table.py
b/paimon-python/pypaimon/table/file_store_table.py
index 26bad0deda..76f2d019a1 100644
--- a/paimon-python/pypaimon/table/file_store_table.py
+++ b/paimon-python/pypaimon/table/file_store_table.py
@@ -520,7 +520,12 @@ class FileStoreTable(Table):
def _copy_with_snapshot(self, snapshot):
"""Keep one resolved read view, including tag metadata and empty
tables."""
from pypaimon.snapshot.time_travel_util import SCAN_KEYS
- options = {key: None for key in SCAN_KEYS if key in
self.table_schema.options}
+ selectors = SCAN_KEYS + [
+ CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP.key(),
+ CoreOptions.SCAN_FILE_CREATION_TIME_MILLIS.key(),
+ CoreOptions.SCAN_CREATION_TIME_MILLIS.key(),
+ ]
+ options = {key: None for key in selectors if key in
self.table_schema.options}
options[CoreOptions.SCAN_MODE.key()] = "from-snapshot" if snapshot is
not None else "default"
if snapshot is not None:
options[CoreOptions.SCAN_SNAPSHOT_ID.key()] = str(snapshot.id)
diff --git a/paimon-python/pypaimon/tests/arrow_schema_test.py
b/paimon-python/pypaimon/tests/arrow_schema_test.py
index 5abf03e75a..274973a5cf 100644
--- a/paimon-python/pypaimon/tests/arrow_schema_test.py
+++ b/paimon-python/pypaimon/tests/arrow_schema_test.py
@@ -102,7 +102,7 @@ class ArrowSchemaTest(unittest.TestCase):
if entry in ('batch', 'postpone_batch'):
writer.write_arrow_batch(source.to_batches()[0])
elif entry == 'bucket':
- writer.write_arrow_batch_to_bucket(source.to_batches()[0],
0)
+ writer.write_arrow(source.to_batches()[0], bucket=0)
else:
writer.write_arrow(source)
commit.commit(writer.prepare_commit())
diff --git
a/paimon-python/pypaimon/tests/daft/daft_pk_distributed_write_test.py
b/paimon-python/pypaimon/tests/daft/daft_pk_distributed_write_test.py
index c96ad3c20c..e087277d0e 100644
--- a/paimon-python/pypaimon/tests/daft/daft_pk_distributed_write_test.py
+++ b/paimon-python/pypaimon/tests/daft/daft_pk_distributed_write_test.py
@@ -312,6 +312,63 @@ def
test_hash_dynamic_pk_restores_bucket_mapping_across_commits(catalog):
)
[email protected]_plan
[email protected]('bucket', ['4', '-1'])
[email protected]('partitioned', [False, True])
[email protected]('overwrite', [False, True])
+def test_grouped_pk_writes_use_native_without_python_fallback(
+ catalog, monkeypatch, bucket, partitioned, overwrite):
+ from pypaimon.write.native_write import NativeTableWrite
+
+ def reject_fallback(self):
+ raise AssertionError('Grouped Daft write fell back to Python')
+
+ monkeypatch.setattr(NativeTableWrite, '_switch_to_python', reject_fallback)
+ table = _create_pk_table(catalog, 'native_grouped',
partitioned=partitioned, options={
+ 'bucket': bucket, 'dynamic-bucket.target-row-num': '1',
+ 'write.native.enabled': 'true', 'read.native.enabled': 'true',
+ })
+ initial = {'id': [1, 2], 'value': ['old', 'keep']}
+ updated = {'id': [1, 3], 'value': ['new', 'added']}
+ if partitioned:
+ initial['dt'] = ['a', 'b']
+ updated['dt'] = ['a', 'c']
+ _write_table(daft.from_pydict(initial).into_partitions(2),
table).to_pydict()
+ _write_table(daft.from_pydict(updated).into_partitions(2), table,
+ mode='overwrite' if overwrite else 'append').to_pydict()
+ # Partitioned tables overwrite only the incoming partitions by default.
+ expected = updated if overwrite and not partitioned else {
+ 'id': [1, 2, 3], 'value': ['new', 'keep', 'added']}
+ if partitioned:
+ expected['dt'] = ['a', 'b', 'c']
+ assert _read_table(table).sort('id').to_pydict() == expected
+ if bucket == '-1':
+ indexes = [entry for entry in IndexFileHandler(table).scan(
+ table.snapshot_manager().get_latest_snapshot())
+ if entry.index_file.index_type == 'HASH']
+ assert sum(entry.index_file.row_count for entry in indexes) ==
len(expected['id'])
+ assert len({(tuple(entry.partition.values), entry.bucket) for entry in
indexes}) == len(indexes)
+
+
[email protected]('native', [False, pytest.param(True,
marks=pytest.mark.native_plan)])
[email protected]('ignore_option,kind', [
+ ('ignore-delete', '-D'), ('ignore-update-before', '-U')])
+def test_grouped_dynamic_row_kind_filter_keeps_surviving_key_indexed(catalog,
native, ignore_option, kind):
+ options = {'bucket': '-1', 'file.format': 'parquet', 'rowkind.field': 'op',
+ ignore_option: 'true', 'write.native.enabled':
str(native).lower()}
+ catalog.create_table('test_db.filtered_group',
pypaimon.Schema.from_pyarrow_schema(
+ pa.schema([('id', pa.int64()), ('value', pa.string()), ('op',
pa.string())]),
+ primary_keys=['id'], options=options), False)
+ table = catalog.get_table('test_db.filtered_group')
+ _write_table(daft.from_pydict({'id': [1, 1], 'value': ['ignored',
'retained'],
+ 'op': [kind, '+I']}), table).to_pydict()
+ assert _read_table(table).to_pydict() == {'id': [1], 'value':
['retained'], 'op': ['+I']}
+ indexes = [entry for entry in IndexFileHandler(table).scan(
+ table.snapshot_manager().get_latest_snapshot()) if
entry.index_file.index_type == 'HASH']
+ assert len(indexes) == 1
+ assert indexes[0].index_file.row_count == 1
+
+
def test_hash_dynamic_pins_one_driver_snapshot_for_all_group_udfs(
catalog, monkeypatch
):
diff --git a/paimon-python/pypaimon/tests/native_precomputed_bucket_test.py
b/paimon-python/pypaimon/tests/native_precomputed_bucket_test.py
new file mode 100644
index 0000000000..8b62c6b221
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_precomputed_bucket_test.py
@@ -0,0 +1,420 @@
+# 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 writes of the partition/bucket groups produced by Daft."""
+
+from unittest.mock import patch
+
+import pyarrow as pa
+import pytest
+
+from pypaimon import CatalogFactory, Schema
+from pypaimon.index.index_file_handler import IndexFileHandler
+from pypaimon.write.native_write import NativeTableWrite
+
+
+pytestmark = pytest.mark.native_plan
+
+
+def _table(tmp_path, bucket, partitioned=False, options=None,
primary_key=True):
+ catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_database('default', True)
+ catalog.create_table('default.t', Schema.from_pyarrow_schema(
+ pa.schema([('id', pa.int64()), ('value', pa.string()), ('pt',
pa.string())]),
+ primary_keys=(['id', 'pt'] if partitioned else ['id']) if primary_key
else [],
+ partition_keys=['pt'] if partitioned else [],
+ options={'bucket': str(bucket), 'file.format': 'parquet',
+ 'write.native.enabled': 'true', 'read.native.enabled': 'true',
+ 'dynamic-bucket.target-row-num': '1', **(options or {})}),
False)
+ return catalog.get_table('default.t')
+
+
+def _batch(ids, values, partitions=None):
+ return pa.record_batch([pa.array(ids, pa.int64()), pa.array(values,
pa.string()),
+ pa.array(partitions or ['a'] * len(ids),
pa.string())],
+ names=['id', 'value', 'pt'])
+
+
+def _rows(table):
+ builder = table.new_read_builder()
+ return sorted(builder.new_read().to_arrow(
+ builder.new_scan().plan().splits()).to_pylist(), key=lambda row:
(row['pt'], row['id']))
+
+
+def _indexes(table):
+ snapshot = table.snapshot_manager().get_latest_snapshot()
+ return [entry for entry in IndexFileHandler(table).scan(snapshot)
+ if entry.index_file.index_type == 'HASH']
+
+
[email protected]('streaming', [False, True])
+def test_native_precomputed_fixed_bucket_after_regular_write(tmp_path,
streaming):
+ table = _table(tmp_path, 4)
+ builder = (table.new_stream_write_builder() if streaming
+ else table.new_batch_write_builder())
+ writer = builder.new_write()
+ assert isinstance(writer, NativeTableWrite)
+ try:
+ with patch.object(writer, '_switch_to_python',
side_effect=AssertionError('Python fallback')):
+ writer.write_arrow_batch(_batch([1], ['regular']))
+ writer.write_arrow(_batch([2, 2], ['old', 'new']), bucket=3)
+ messages = writer.prepare_commit(7) if streaming else
writer.prepare_commit()
+ assert any(message.bucket == 3 for message in messages)
+ commit = builder.new_commit()
+ try:
+ commit.commit(messages, 7) if streaming else
commit.commit(messages)
+ finally:
+ commit.close()
+ finally:
+ writer.close()
+ assert _rows(table) == [{'id': 1, 'value': 'regular', 'pt': 'a'},
+ {'id': 2, 'value': 'new', 'pt': 'a'}]
+
+
[email protected]('partitioned', [False, True])
+def test_native_dynamic_index_keeps_upstream_bucket_on_restart(tmp_path,
partitioned):
+ table = _table(tmp_path, -1, partitioned)
+ for ids, values in [([1, 2], ['one', 'two']), ([2, 3], ['updated',
'three'])]:
+ builder = table.new_batch_write_builder()
+ writer = builder.new_write()
+ assert isinstance(writer, NativeTableWrite)
+ try:
+ with patch.object(writer, '_switch_to_python',
side_effect=AssertionError('Python fallback')):
+ writer.write_arrow(_batch(ids, values), bucket=17)
+ messages = writer.prepare_commit()
+ assert {message.bucket for message in messages} == {17}
+ builder.new_commit().commit(messages)
+ finally:
+ writer.close()
+ indexes = _indexes(table)
+ assert len(indexes) == 1
+ assert indexes[0].bucket == 17
+ assert indexes[0].index_file.row_count == 3
+ assert _rows(table) == [{'id': 1, 'value': 'one', 'pt': 'a'},
+ {'id': 2, 'value': 'updated', 'pt': 'a'},
+ {'id': 3, 'value': 'three', 'pt': 'a'}]
+
+
[email protected]('streaming', [False, True])
[email protected]('partitioned', [False, True])
+def test_native_computes_java_hashes_for_multiple_batches(tmp_path, streaming,
partitioned):
+ table = _table(tmp_path, -1, partitioned)
+ builder = (table.new_stream_write_builder() if streaming
+ else table.new_batch_write_builder())
+ builder.with_restore_snapshot(0)
+ writer = builder.new_write()
+ first = _batch([1, 2], ['one', 'two'])
+ try:
+ with patch.object(writer, '_switch_to_python',
side_effect=AssertionError('Python fallback')):
+ writer.write_arrow(first, bucket=12)
+ # A second batch can refer to this writer's uncommitted mappings.
+ updated = _batch([2, 3], ['updated', 'three'])
+ writer.write_arrow(updated, bucket=12)
+ messages = writer.prepare_commit(8) if streaming else
writer.prepare_commit()
+ assert len(messages) == 1
+ assert messages[0].index_adds[0].index_file.row_count == 3
+ commit = builder.new_commit()
+ try:
+ commit.commit(messages, 8) if streaming else
commit.commit(messages)
+ finally:
+ commit.close()
+ finally:
+ writer.close()
+ assert _rows(table) == [{'id': 1, 'value': 'one', 'pt': 'a'},
+ {'id': 2, 'value': 'updated', 'pt': 'a'},
+ {'id': 3, 'value': 'three', 'pt': 'a'}]
+ # Restored hashes and indexes remain readable by Python's assigner.
+ from pypaimon.write.row_key_extractor import DynamicBucketRowKeyExtractor
+
+ extractor = DynamicBucketRowKeyExtractor(table.table_schema, table=table)
+ assert extractor.extract_partition_bucket_batch(_batch([1, 2, 3], ['a',
'b', 'c']))[1] == [12, 12, 12]
+
+
[email protected]('bucket', [4, -1])
[email protected]('invalid', ['bucket', 'partition'])
+def test_invalid_precomputed_group_cannot_stage_half_a_partition(tmp_path,
bucket, invalid):
+ table = _table(tmp_path, bucket, partitioned=True)
+ writer = table.new_batch_write_builder().new_write()
+ batch = _batch([1, 2], ['a', 'b'])
+ if invalid == 'partition':
+ batch = _batch([1, 2], ['a', 'b'], ['a', 'b'])
+ try:
+ with patch.object(writer, '_switch_to_python',
side_effect=AssertionError('Python fallback')):
+ with pytest.raises(ValueError):
+ writer.write_arrow(batch, bucket=-1 if invalid == 'bucket'
else 0)
+ assert writer.prepare_commit() == []
+ finally:
+ writer.close()
+ assert table.snapshot_manager().get_latest_snapshot() is None
+
+
+def test_precomputed_append_projects_full_or_selected_schema(tmp_path):
+ table = _table(tmp_path, 4, partitioned=True, primary_key=False,
+ options={'bucket-key': 'id'})
+ builder = table.new_batch_write_builder()
+ writer = builder.new_write()
+ try:
+ with patch.object(writer, '_switch_to_python',
side_effect=AssertionError('Python fallback')):
+ writer.with_write_type(['pt', 'id'])
+ writer.write_arrow(_batch([1], ['omitted']), bucket=3)
+ writer.write_arrow(_batch([2], ['omitted']).select(['pt', 'id']),
bucket=3)
+ messages = writer.prepare_commit()
+ assert {message.bucket for message in messages} == {3}
+ assert all(file.write_cols == ['pt', 'id'] for message in messages for
file in message.new_files)
+ builder.new_commit().commit(messages)
+ finally:
+ writer.close()
+ assert _rows(table) == [{'id': 1, 'value': None, 'pt': 'a'},
+ {'id': 2, 'value': None, 'pt': 'a'}]
+
+
[email protected]('scan_mode', ['default', 'latest-full',
'incremental'])
+def
test_pinned_dynamic_configuration_survives_early_python_selection(tmp_path,
scan_mode):
+ options = {'scan.mode': scan_mode}
+ if scan_mode == 'incremental':
+ options['incremental-between-timestamp'] = '0,9999999999999'
+ table = _table(tmp_path, -1, options=options)
+ initial = table.new_batch_write_builder()
+ first = initial.new_write()
+ first.write_arrow(_batch([1], ['old']), bucket=0)
+ initial.new_commit().commit(first.prepare_commit())
+ first.close()
+ pinned_id = table.snapshot_manager().get_latest_snapshot().id
+ writer =
table.new_batch_write_builder().with_restore_snapshot(pinned_id).new_write()
+ try:
+ later = table.new_batch_write_builder()
+ second = later.new_write()
+ second.write_arrow(_batch([2], ['later']), bucket=0)
+ later.new_commit().commit(second.prepare_commit())
+ second.close()
+ # This existing advanced API still selects Python before any write.
+ writer.with_blob_consumer(lambda field, blob: blob)
+ assert writer._python_writer.row_key_extractor.base_snapshot_id ==
pinned_id
+ writer.write_arrow(_batch([1], ['restored']), bucket=0)
+ messages = writer.prepare_commit()
+ assert messages[0].new_files[0].min_sequence_number == 1
+ # Never submit this historical restore after another bucket owner's
write.
+ table.new_batch_write_builder().new_commit().abort(messages)
+ finally:
+ writer.close()
+
+
[email protected]('streaming', [False, True])
+def test_builder_snapshot_configuration_is_copied_to_each_writer(tmp_path,
streaming):
+ table = _table(tmp_path, -1)
+ builder = (table.new_stream_write_builder() if streaming
+ else table.new_batch_write_builder())
+ assert builder.with_restore_snapshot(0) is builder
+ writer = builder.new_write()
+ try:
+ writer.write_arrow(_batch([], []), bucket=0)
+ # Later builder configuration must not alter a writer already created.
+ builder.with_restore_snapshot(99)
+ writer.write_arrow(_batch([1], ['one']), bucket=0)
+ messages = writer.prepare_commit(1) if streaming else
writer.prepare_commit()
+ assert messages[0].index_adds[0].index_file.row_count == 1
+ finally:
+ writer.close()
+
+
+def
test_native_hash_mapping_conflict_poisoned_writer_preserves_prepared_files(tmp_path):
+ table = _table(tmp_path, -1)
+ builder = table.new_stream_write_builder()
+ writer = builder.new_write()
+ try:
+ writer.write_arrow(_batch([1], ['published']), bucket=7)
+ messages = writer.prepare_commit(1)
+ builder.new_commit().commit(messages, 1)
+ paths = [file.file_path for message in messages for file in
message.new_files]
+ writer.write_arrow(_batch([2], ['pending']), bucket=7)
+ with pytest.raises(ValueError, match='belongs to bucket 7'):
+ writer.write_arrow(_batch([2], ['conflict']), bucket=8)
+ with pytest.raises(ValueError, match='cannot be reused'):
+ writer.prepare_commit(2)
+ assert all(table.file_io.exists(path) for path in paths)
+ finally:
+ writer.close()
+ assert _rows(table) == [{'id': 1, 'value': 'published', 'pt': 'a'}]
+
+
[email protected]('publish_first', [False, True])
+def test_native_checkpoint_sequence_advances_with_pinned_index_base(tmp_path,
publish_first):
+ table = _table(tmp_path, -1)
+ builder = table.new_stream_write_builder().with_restore_snapshot(0)
+ writer = builder.new_write()
+ commit = builder.new_commit()
+ try:
+ writer.write_arrow(_batch([2, 1], ['other', 'old']), bucket=7)
+ first = writer.prepare_commit(1)
+ if publish_first:
+ commit.commit(first, 1)
+ writer.write_arrow(_batch([1], ['new']), bucket=7)
+ second = writer.prepare_commit(2)
+ assert second[0].new_files[0].min_sequence_number == 2
+ if not publish_first:
+ commit.commit(first, 1)
+ commit.commit(second, 2)
+ finally:
+ writer.close()
+ commit.close()
+ assert _rows(table) == [{'id': 1, 'value': 'new', 'pt': 'a'},
+ {'id': 2, 'value': 'other', 'pt': 'a'}]
+
+
+def test_restore_snapshot_pins_sequences_and_default_restores_latest(tmp_path):
+ table = _table(tmp_path, -1)
+ first_builder = table.new_batch_write_builder()
+ first = first_builder.new_write()
+ first.write_arrow(_batch([1], ['old']), bucket=7)
+ first_builder.new_commit().commit(first.prepare_commit())
+ first.close()
+ builder = table.new_batch_write_builder().with_restore_snapshot(1)
+ pinned = builder.new_write()
+ try:
+ later_builder = table.new_batch_write_builder()
+ later = later_builder.new_write()
+ later.write_arrow(_batch([1, 1, 1], ['b', 'c', 'd']), bucket=7)
+ later_builder.new_commit().commit(later.prepare_commit())
+ later.close()
+ pinned.write_arrow(_batch([1], ['new']), bucket=7)
+ messages = pinned.prepare_commit()
+ assert messages[0].new_files[0].min_sequence_number == 1
+ # These messages were never submitted. Preserve the newer bucket
owner's
+ # data and verify a default writer still restores the latest snapshot.
+ builder.new_commit().abort(messages)
+ finally:
+ pinned.close()
+ assert _rows(table) == [{'id': 1, 'value': 'd', 'pt': 'a'}]
+ current_builder = table.new_batch_write_builder()
+ current = current_builder.new_write()
+ try:
+ current.write_arrow(_batch([1], ['new']), bucket=7)
+ messages = current.prepare_commit()
+ assert messages[0].new_files[0].min_sequence_number == 4
+ current_builder.new_commit().commit(messages)
+ finally:
+ current.close()
+ assert _rows(table) == [{'id': 1, 'value': 'new', 'pt': 'a'}]
+
+
[email protected]('streaming', [False, True])
[email protected]('bucket_mode', [4, -1])
+def test_arrow_table_keeps_explicit_zero_bucket_across_chunks(tmp_path,
streaming, bucket_mode):
+ table = _table(tmp_path, bucket_mode)
+ builder = (table.new_stream_write_builder() if streaming
+ else table.new_batch_write_builder())
+ writer = builder.new_write()
+ data = pa.Table.from_batches([_batch([1, 2], ['old', 'other']),
_batch([1], ['new'])])
+ assert len(data.to_batches()) == 2
+ try:
+ with patch.object(writer, '_switch_to_python',
side_effect=AssertionError('Python fallback')):
+ writer.write_arrow(data, bucket=0)
+ messages = writer.prepare_commit(1) if streaming else
writer.prepare_commit()
+ assert {message.bucket for message in messages} == {0}
+ commit = builder.new_commit()
+ try:
+ commit.commit(messages, 1) if streaming else
commit.commit(messages)
+ finally:
+ commit.close()
+ finally:
+ writer.close()
+ assert _rows(table) == [{'id': 1, 'value': 'new', 'pt': 'a'},
+ {'id': 2, 'value': 'other', 'pt': 'a'}]
+
+
[email protected]('streaming', [False, True])
+def
test_builder_rejects_invalid_index_snapshot_before_native_selection(tmp_path,
streaming):
+ table = _table(tmp_path, -1)
+ builder = (table.new_stream_write_builder() if streaming
+ else table.new_batch_write_builder())
+ for value in [-1, True, None, '1']:
+ with pytest.raises(ValueError, match='nonnegative integer'):
+ builder.with_restore_snapshot(value)
+ fixed = _table(tmp_path / 'fixed', 4)
+ with pytest.raises(ValueError, match='HASH_DYNAMIC'):
+ fixed.new_batch_write_builder().with_restore_snapshot(0)
+ assert table.snapshot_manager().get_latest_snapshot() is None
+
+
[email protected]('python_fallback', [False, True])
+def test_empty_restore_ignores_existing_sequences_and_hashes(tmp_path,
python_fallback):
+ table = _table(tmp_path, -1)
+ initial = table.new_batch_write_builder()
+ first = initial.new_write()
+ first.write_arrow(_batch([1, 2], ['one', 'two']), bucket=0)
+ initial.new_commit().commit(first.prepare_commit())
+ first.close()
+ builder = table.new_batch_write_builder().with_restore_snapshot(0)
+ writer = builder.new_write()
+ try:
+ if python_fallback:
+ writer.with_blob_consumer(lambda field, blob: blob)
+ writer.write_arrow(_batch([3], ['empty-base']), bucket=0)
+ messages = writer.prepare_commit()
+ assert messages[0].new_files[0].min_sequence_number == 0
+ assert messages[0].index_adds[0].index_file.row_count == 1
+ # Explicit abort is safe: these messages have never been submitted.
+ builder.new_commit().abort(messages)
+ finally:
+ writer.close()
+ assert _rows(table) == [{'id': 1, 'value': 'one', 'pt': 'a'},
+ {'id': 2, 'value': 'two', 'pt': 'a'}]
+
+
[email protected]('overwrite', [False, True])
+def test_empty_restore_preserves_python_write_authorization(tmp_path,
overwrite):
+ from pypaimon.catalog.catalog_exception import TableNoPermissionException
+ from pypaimon.catalog.table_query_auth import TableQueryAuthResult
+
+ table = _table(tmp_path, -1)
+ builder = table.new_batch_write_builder().with_restore_snapshot(0)
+ if overwrite:
+ builder.overwrite()
+ writer = builder.new_write()
+ # Exercise the ordinary writer selected before any Native data is staged.
+ writer.with_blob_consumer(lambda field, blob: blob)
+ auth = TableQueryAuthResult(None, {'value': '{"name":"NULL"}'})
+ try:
+ with patch.object(type(table.catalog_environment), 'table_query_auth',
+ return_value=lambda _: auth):
+ with pytest.raises(TableNoPermissionException):
+ writer.write_arrow(_batch([1], ['forbidden']), bucket=0)
+ finally:
+ writer.close()
+ assert table.snapshot_manager().get_latest_snapshot() is None
+
+
+def test_python_pk_file_sequence_range_matches_deduplicated_rows(tmp_path):
+ import pyarrow.parquet as pq
+
+ table = _table(tmp_path, -1)
+ builder = table.new_batch_write_builder().with_restore_snapshot(0)
+ writer = builder.new_write()
+ try:
+ writer.with_blob_consumer(lambda field, blob: blob)
+ writer.write_arrow(_batch([1, 1], ['old', 'new']), bucket=0)
+ messages = writer.prepare_commit()
+ file = messages[0].new_files[0]
+ sequences = pq.read_table(file.file_path,
columns=['_SEQUENCE_NUMBER']).column(0).to_pylist()
+ assert sequences == [1]
+ assert file.row_count == 1
+ assert file.min_sequence_number == min(sequences)
+ assert file.max_sequence_number == max(sequences)
+ builder.new_commit().commit(messages)
+ finally:
+ writer.close()
+ assert _rows(table) == [{'id': 1, 'value': 'new', 'pt': 'a'}]
diff --git a/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py
b/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py
index 218532bbbf..a8b71f6b73 100644
--- a/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py
+++ b/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py
@@ -69,13 +69,17 @@ class DynamicBucketTest(unittest.TestCase):
return catalog.get_table(f'default.{name}')
@staticmethod
- def _prepare_indexed_write(table, ids):
+ def _prepare_indexed_write(table, ids, bucket=None):
builder = table.new_batch_write_builder()
- writer = builder.new_write().with_dynamic_bucket_index()
- writer.write_arrow_batch(pa.RecordBatch.from_pydict({
+ writer = builder.new_write()
+ batch = pa.RecordBatch.from_pydict({
'id': ids,
'value': [f'v-{value}' for value in ids],
- }))
+ })
+ if bucket is None:
+ writer.write_arrow_batch(batch)
+ else:
+ writer.write_arrow(batch, bucket=bucket)
return writer, builder.new_commit(), writer.prepare_commit()
@staticmethod
@@ -538,8 +542,7 @@ class DynamicBucketTest(unittest.TestCase):
table = self._create_table(root, 'implicit_hash_replace')
self._commit_arrow(table, [1], ['one'])
writer, commit, messages = self._prepare_indexed_write(table, [2])
- self.assertEqual(1, len(messages[0].index_deletes))
- messages[0].index_deletes.clear()
+ self.assertEqual([], messages[0].index_deletes)
commit.commit(messages)
indexes = self._hash_indexes(table)
self.assertEqual(1, len(indexes))
@@ -547,7 +550,7 @@ class DynamicBucketTest(unittest.TestCase):
writer.close()
commit.close()
- def test_concurrent_hash_index_replacement_conflicts(self):
+ def test_sequential_hash_index_replacements_keep_one_complete_index(self):
with tempfile.TemporaryDirectory() as root:
table = self._create_table(root, 'concurrent_replace')
seed_writer, seed_commit, seed_messages = (
@@ -560,24 +563,17 @@ class DynamicBucketTest(unittest.TestCase):
writer1, commit1, messages1 = self._prepare_indexed_write(
table, [2]
)
+ commit1.commit(messages1)
writer2, commit2, messages2 = self._prepare_indexed_write(
table, [3]
)
- old_index = messages1[0].index_deletes[0].index_file.file_name
- self.assertEqual(
- old_index,
- messages2[0].index_deletes[0].index_file.file_name,
- )
-
- commit1.commit(messages1)
- with self.assertRaisesRegex(
- RuntimeError, 'HASH index conflict detected'
- ):
- commit2.commit(messages2)
+ self.assertEqual([], messages1[0].index_deletes)
+ self.assertEqual([], messages2[0].index_deletes)
+ commit2.commit(messages2)
indexes = self._hash_indexes(table)
self.assertEqual(1, len(indexes))
- self.assertEqual(2, indexes[0].index_file.row_count)
+ self.assertEqual(3, indexes[0].index_file.row_count)
writer1.close()
writer2.close()
commit1.close()
@@ -598,20 +594,12 @@ class DynamicBucketTest(unittest.TestCase):
seed_writer.close()
seed_commit.close()
- with patch(
- 'pypaimon.index.dynamic_bucket.random.choice',
- return_value=0,
- ):
- writer1, commit1, messages1 = self._prepare_indexed_write(
- table, [3]
- )
- with patch(
- 'pypaimon.index.dynamic_bucket.random.choice',
- return_value=1,
- ):
- writer2, commit2, messages2 = self._prepare_indexed_write(
- table, [4]
- )
+ writer1, commit1, messages1 = self._prepare_indexed_write(
+ table, [3], bucket=0
+ )
+ writer2, commit2, messages2 = self._prepare_indexed_write(
+ table, [4], bucket=1
+ )
self.assertNotEqual(messages1[0].bucket, messages2[0].bucket)
commit1.commit(messages1)
@@ -662,11 +650,13 @@ class DynamicBucketTest(unittest.TestCase):
stale_writer.close()
stale_commit.close()
- def test_retry_then_hash_index_conflict_preserves_prepared_files(self):
+ def
test_retry_after_disjoint_hash_index_commit_preserves_prepared_files(self):
with tempfile.TemporaryDirectory() as root:
- table = self._create_table(root, 'retry_hash_conflict')
+ table = self._create_table(root, 'retry_hash_commit',
target_row_num=1)
self._commit_arrow(table, [0], ['seed'])
- writer, commit, messages = self._prepare_indexed_write(table, [1])
+ # The pending upsert owns bucket 0. A new key is assigned bucket 1,
+ # so the competing commit obeys Java's single owner per bucket.
+ writer, commit, messages = self._prepare_indexed_write(table, [0])
prepared_paths = [
file.file_path
for message in messages
@@ -680,19 +670,22 @@ class DynamicBucketTest(unittest.TestCase):
for entry in message.index_adds
]
original_snapshot_commit = commit.file_store_commit.snapshot_commit
+ snapshot_commit = original_snapshot_commit.commit
calls = 0
def lose_first_compare_and_set(
base_snapshot_uuid, snapshot, statistics):
nonlocal calls
calls += 1
- concurrent_writer, concurrent_commit, concurrent_messages = (
- self._prepare_indexed_write(table, [2])
- )
- concurrent_commit.commit(concurrent_messages)
- concurrent_writer.close()
- concurrent_commit.close()
- return False
+ if calls == 1:
+ concurrent_writer, concurrent_commit, concurrent_messages
= (
+ self._prepare_indexed_write(table, [2])
+ )
+ concurrent_commit.commit(concurrent_messages)
+ concurrent_writer.close()
+ concurrent_commit.close()
+ return False
+ return snapshot_commit(base_snapshot_uuid, snapshot,
statistics)
with patch.object(
original_snapshot_commit,
@@ -702,15 +695,14 @@ class DynamicBucketTest(unittest.TestCase):
commit.file_store_commit,
'_commit_retry_wait',
):
- with self.assertRaisesRegex(
- RuntimeError, 'HASH index conflict detected'
- ):
- commit.commit(messages)
+ commit.commit(messages)
- self.assertEqual(1, calls)
+ self.assertEqual(2, calls)
self.assertTrue(all(
table.file_io.exists(path) for path in prepared_paths
))
+ self.assertEqual({'id': [0, 2], 'value': ['v-0', 'v-2']},
+ self._read_arrow(table).sort_by('id').to_pydict())
writer.close()
commit.close()
diff --git a/paimon-python/pypaimon/write/file_store_write.py
b/paimon-python/pypaimon/write/file_store_write.py
index 599e301eeb..16290ff864 100644
--- a/paimon-python/pypaimon/write/file_store_write.py
+++ b/paimon-python/pypaimon/write/file_store_write.py
@@ -57,6 +57,7 @@ class FileStoreWrite:
self.data_writers: Dict[Tuple, DataWriter] = {}
self._runtime_total_buckets: Dict[Tuple, int] = {}
self.max_seq_numbers: dict = {}
+ self.restore_snapshot_id = None
self.write_cols = None
self.blob_consumer = None
self.blob_uri_reader_factory = None
@@ -173,7 +174,8 @@ class FileStoreWrite:
"unset it or write with Java/Flink/Spark.")
def max_seq_number():
- return self._seq_number_stats(partition).get(bucket, 1)
+ default = -1 if self.table.is_primary_key_table else 1
+ return self._seq_number_stats(partition).get(bucket, default)
# Dedicated Blob files are an append-table layout. PK tables require
# managed packs and references attached to their key-value Parquet
files.
@@ -375,6 +377,15 @@ class FileStoreWrite:
return buckets
def _sequence_read_table(self):
+ if self.restore_snapshot_id is not None:
+ snapshot = None
+ if self.restore_snapshot_id != 0:
+ snapshot =
self.table.snapshot_manager().get_snapshot_by_id(self.restore_snapshot_id)
+ if snapshot is None:
+ raise ValueError("Snapshot id '{}' doesn't
exist".format(self.restore_snapshot_id))
+ # Replace inherited selectors and mode, including an explicit empty
+ # view. plan_for_write still validates row filters and column
masks.
+ return self.table._copy_with_snapshot(snapshot)
return self.table
def _load_seq_number_stats(self, partition: Tuple) -> dict:
diff --git a/paimon-python/pypaimon/write/native_write.py
b/paimon-python/pypaimon/write/native_write.py
index 317940226d..1a53d25f6e 100644
--- a/paimon-python/pypaimon/write/native_write.py
+++ b/paimon-python/pypaimon/write/native_write.py
@@ -48,7 +48,7 @@ def _native_partition_types_supported(schema, partition_keys):
def create_native_write(table, commit_user, static_partition=None,
stream=False,
- *, fixed_bucket=False, bucket_plan=None):
+ *, fixed_bucket=False, bucket_plan=None,
restore_snapshot_id=None):
"""Return a native writer if the table can use the filesystem write
path."""
schema = PyarrowFieldParser.from_paimon_schema(table.table_schema.fields)
if (not native_write_available()
@@ -83,8 +83,10 @@ def create_native_write(table, commit_user,
static_partition=None, stream=False,
builder =
native_table.new_batch_write_builder()._with_commit_user(commit_user)
if static_partition is not None:
builder = builder.with_overwrite(static_partition)
+ if restore_snapshot_id is not None:
+ builder.with_restore_snapshot(restore_snapshot_id)
return NativeTableWrite(table, commit_user, static_partition, stream,
- builder.new_write())
+ builder.new_write(), restore_snapshot_id)
class NativeTableWrite:
@@ -95,7 +97,8 @@ class NativeTableWrite:
logical write across two writers, so it is rejected.
"""
- def __init__(self, table, commit_user, static_partition, stream,
native_writer):
+ def __init__(self, table, commit_user, static_partition, stream,
native_writer,
+ restore_snapshot_id=None):
self.table = table
self.commit_user = commit_user
self.static_partition = static_partition
@@ -106,6 +109,7 @@ class NativeTableWrite:
self._table_schema =
PyarrowFieldParser.from_paimon_schema(table.table_schema.fields)
self._schema = self._table_schema
self._write_cols = None
+ self._restore_snapshot_id = restore_snapshot_id
def _switch_to_python(self):
if self._python_writer is not None:
@@ -121,9 +125,11 @@ class NativeTableWrite:
def _new_python_writer(self):
from pypaimon.write.table_write import BatchTableWrite,
StreamTableWrite
if self.stream:
- writer = StreamTableWrite(self.table, self.commit_user)
+ writer = StreamTableWrite(self.table, self.commit_user,
+
restore_snapshot_id=self._restore_snapshot_id)
else:
- writer = BatchTableWrite(self.table, self.commit_user,
self.static_partition)
+ writer = BatchTableWrite(self.table, self.commit_user,
self.static_partition,
+
restore_snapshot_id=self._restore_snapshot_id)
if self._write_cols is not None:
writer.with_write_type(self._write_cols)
return writer
@@ -148,17 +154,26 @@ class NativeTableWrite:
self._write_cols = names
return self
- def write_arrow(self, data):
+ def write_arrow(self, data, bucket=None):
if self._python_writer is not None:
- return self._python_writer.write_arrow(data)
+ return self._python_writer.write_arrow(data, bucket)
+ data = self._prepare_native_arrow_data(data)
if isinstance(data, pa.RecordBatch):
- return self.write_arrow_batch(data)
+ return self.write_arrow_batch(data, bucket)
for batch in data.to_batches():
- self.write_arrow_batch(batch)
+ self.write_arrow_batch(batch, bucket)
- def write_arrow_batch(self, data):
+ def write_arrow_batch(self, data, bucket=None):
if self._python_writer is not None:
- return self._python_writer.write_arrow_batch(data)
+ return self._python_writer.write_arrow_batch(data, bucket)
+ data = self._prepare_native_arrow_data(data)
+ if data.num_rows:
+ # A failed native write may already have produced files. Never
+ # retry that batch through Python after this point.
+ self._written = True
+ self._native_writer.write_arrow(data, bucket)
+
+ def _prepare_native_arrow_data(self, data):
if self._write_cols is not None and arrow_schemas_compatible(
data.schema, self._table_schema,
check_top_level_nullability=False,
allow_binary_compatibility=True):
@@ -172,12 +187,7 @@ class NativeTableWrite:
"Input schema isn't consistent with table schema and write
cols. "
f"Input schema is: {data.schema} Table schema is:
{self._schema} "
f"Write cols is: {self._write_cols}")
- data = normalize_arrow_strings(data)
- if data.num_rows:
- # A failed native write may already have produced files. Never
- # retry that batch through Python after this point.
- self._written = True
- self._native_writer.write_arrow(data)
+ return normalize_arrow_strings(data)
def write_pandas(self, dataframe):
if self._python_writer is not None:
diff --git a/paimon-python/pypaimon/write/postpone_batch_table_write.py
b/paimon-python/pypaimon/write/postpone_batch_table_write.py
index 25227a06bf..f00f15b544 100644
--- a/paimon-python/pypaimon/write/postpone_batch_table_write.py
+++ b/paimon-python/pypaimon/write/postpone_batch_table_write.py
@@ -126,14 +126,18 @@ class PostponeFixedBucketBatchTableWrite(BatchTableWrite):
for partition in self._planner.input_partition_stats(data)
)
- def write_arrow(self, table: pa.Table):
+ def write_arrow(self, table: pa.Table, bucket=None):
+ if bucket is not None:
+ raise ValueError('Precomputed bucket writes require HASH_FIXED or
HASH_DYNAMIC tables')
table = self._prepare_arrow_data(table)
if not self._buffer_input(table):
return super().write_arrow(table)
self._pending_inputs.extend(
("batch", batch) for batch in table.to_batches())
- def write_arrow_batch(self, data: pa.RecordBatch):
+ def write_arrow_batch(self, data: pa.RecordBatch, bucket=None):
+ if bucket is not None:
+ raise ValueError('Precomputed bucket writes require HASH_FIXED or
HASH_DYNAMIC tables')
data = self._prepare_arrow_data(data)
if not self._buffer_input(data):
return super().write_arrow_batch(data)
diff --git a/paimon-python/pypaimon/write/row_key_extractor.py
b/paimon-python/pypaimon/write/row_key_extractor.py
index 1cb629b7b2..53c62bf52a 100644
--- a/paimon-python/pypaimon/write/row_key_extractor.py
+++ b/paimon-python/pypaimon/write/row_key_extractor.py
@@ -492,6 +492,8 @@ class DynamicBucketRowKeyExtractor(RowKeyExtractor):
snapshot = table.snapshot_manager().get_snapshot_by_id(
base_snapshot_id
)
+ if snapshot is None:
+ raise ValueError("Snapshot id '{}' doesn't
exist".format(base_snapshot_id))
self.base_snapshot_id = snapshot.id if snapshot is not None else 0
self._assigner = HashBucketAssigner(
table=table,
@@ -662,58 +664,25 @@ class DynamicBucketRowKeyExtractor(RowKeyExtractor):
def notify_precomputed_bucket_batch(
self, data: pa.RecordBatch, bucket: int
) -> Optional[Tuple]:
- """Maintain HASH indexes for rows assigned by an upstream
coordinator."""
- partitions, _, key_hashes = self.extract_hashes_batch(data)
- return self.notify_precomputed_bucket_hashes_batch(
- data, bucket, key_hashes, partitions=partitions
- )
-
- def notify_precomputed_bucket_hashes_batch(
- self,
- data: pa.RecordBatch,
- bucket: int,
- key_hashes: List[int],
- partitions: Optional[List[Tuple]] = None,
- new_mappings: Optional[List[bool]] = None,
- ) -> Optional[Tuple]:
- """Maintain HASH indexes using hashes carried through the shuffle."""
+ """Maintain HASH indexes for every row in an upstream bucket."""
from pypaimon.index.dynamic_bucket import to_signed_int32
if self._index_maintainer is None:
- raise RuntimeError(
- "Precomputed dynamic buckets require a persistent table
extractor"
- )
- if len(key_hashes) != data.num_rows:
- raise ValueError(
- "Precomputed key hash count {} does not match row count
{}".format(
- len(key_hashes), data.num_rows
- )
- )
- if partitions is None:
- partitions = self._extract_partitions_batch(data)
- if new_mappings is not None and len(new_mappings) != data.num_rows:
- raise ValueError(
- "Precomputed new-mapping count {} does not match row count
{}".format(
- len(new_mappings), data.num_rows
- )
- )
+ raise RuntimeError('Precomputed dynamic buckets require a
persistent table extractor')
+ partitions, _, key_hashes = self.extract_hashes_batch(data)
if not partitions:
return None
partition = tuple(partitions[0])
- if new_mappings is None:
- new_mappings = [True] * data.num_rows
- for actual_partition, key_hash, is_new in zip(
- partitions, key_hashes, new_mappings
- ):
- if tuple(actual_partition) != tuple(partition):
- raise RuntimeError(
- "A precomputed dynamic-bucket group contained multiple "
- f"partitions: expected {partition}, got {actual_partition}"
- )
- if is_new:
- self._index_maintainer.notify_new_record(
- partition, bucket, to_signed_int32(key_hash)
- )
+ if any(tuple(actual_partition) != partition for actual_partition in
partitions):
+ raise RuntimeError(
+ "A precomputed dynamic-bucket group contained multiple "
+ f"partitions: expected {partition}, got {partitions}"
+ )
+ # Java's maintainer observes every written key after RowKind filtering.
+ for key_hash in key_hashes:
+ self._index_maintainer.notify_new_record(
+ partition, bucket, to_signed_int32(key_hash)
+ )
return partition
def release_prepared(self) -> None:
diff --git a/paimon-python/pypaimon/write/table_write.py
b/paimon-python/pypaimon/write/table_write.py
index e982248dbd..04e7b4df01 100644
--- a/paimon-python/pypaimon/write/table_write.py
+++ b/paimon-python/pypaimon/write/table_write.py
@@ -36,14 +36,19 @@ if TYPE_CHECKING:
class TableWrite:
- def __init__(self, table, commit_user, static_partition: Optional[dict] =
None):
+ def __init__(self, table, commit_user, static_partition: Optional[dict] =
None,
+ *, restore_snapshot_id=None):
from pypaimon.table.file_store_table import FileStoreTable
self.table: FileStoreTable = table
self.table_pyarrow_schema =
PyarrowFieldParser.from_paimon_schema(self.table.table_schema.fields)
self.commit_user = commit_user
self.static_partition = static_partition
+ self.restore_snapshot_id = restore_snapshot_id
self.file_store_write = self._create_file_store_write(commit_user)
+ self.file_store_write.restore_snapshot_id = (
+ 0 if static_partition is not None and restore_snapshot_id is not
None
+ else restore_snapshot_id)
if static_partition is not None:
# An overwrite replaces state, not an input changelog. Java's
# overwrite commit does not publish changelog manifests; avoid
@@ -55,17 +60,27 @@ class TableWrite:
return FileStoreWrite(self.table, commit_user)
def _create_row_key_extractor(self, static_partition):
+ if self.restore_snapshot_id is not None:
+ from pypaimon.write.row_key_extractor import
DynamicBucketRowKeyExtractor
+ return DynamicBucketRowKeyExtractor(
+ self.table.table_schema, table=self.table,
+ ignore_existing=static_partition is not None,
+ base_snapshot_id=self.restore_snapshot_id)
return self.table.create_row_key_extractor(
ignore_existing=static_partition is not None
)
- def write_arrow(self, table: pa.Table):
- table = self._prepare_arrow_data(table)
- batches_iterator = table.to_batches()
- for batch in batches_iterator:
- self.write_arrow_batch(batch)
-
- def write_arrow_batch(self, data: pa.RecordBatch):
+ def write_arrow(self, data, bucket: Optional[int] = None):
+ """Write Arrow data, optionally to a bucket assigned upstream."""
+ self._validate_pyarrow_schema(data.schema)
+ if isinstance(data, pa.RecordBatch):
+ return self.write_arrow_batch(data, bucket)
+ for batch in data.to_batches():
+ self.write_arrow_batch(batch, bucket)
+
+ def write_arrow_batch(self, data: pa.RecordBatch, bucket: Optional[int] =
None):
+ if bucket is not None:
+ return self._write_precomputed_bucket(data, bucket)
data = self._prepare_arrow_data(data)
for partition, bucket, row_indices in \
@@ -102,40 +117,8 @@ class TableWrite:
def _write_partition_bucket_batch(self, partition, bucket, data):
self.file_store_write.write(partition, bucket, data)
- def with_dynamic_bucket_index(
- self,
- ignore_existing: bool = False,
- base_snapshot_id: Optional[int] = None,
- ):
- """Enable persistent HASH-index maintenance for coordinated writes."""
- from pypaimon.table.bucket_mode import BucketMode
- from pypaimon.write.row_key_extractor import
DynamicBucketRowKeyExtractor
-
- if self.table.bucket_mode() != BucketMode.HASH_DYNAMIC:
- raise ValueError(
- "Dynamic bucket index maintenance is only valid for "
- "HASH_DYNAMIC tables"
- )
- if self.file_store_write.data_writers:
- raise RuntimeError(
- "Dynamic bucket index maintenance must be enabled before
writing"
- )
- self.row_key_extractor = DynamicBucketRowKeyExtractor(
- self.table.table_schema,
- table=self.table,
- ignore_existing=ignore_existing,
- base_snapshot_id=base_snapshot_id,
- )
- return self
-
- def write_arrow_batch_to_bucket(
- self,
- data: pa.RecordBatch,
- bucket: int,
- key_hashes: Optional[List[int]] = None,
- new_mappings: Optional[List[bool]] = None,
- ):
- """Write one complete group whose bucket was computed upstream."""
+ def _write_precomputed_bucket(self, data, bucket):
+ """Write one partition/bucket group using normal index maintenance."""
from pypaimon.table.bucket_mode import BucketMode
from pypaimon.write.row_key_extractor import
DynamicBucketRowKeyExtractor
@@ -148,42 +131,20 @@ class TableWrite:
if not isinstance(self.row_key_extractor,
DynamicBucketRowKeyExtractor):
if bucket_mode == BucketMode.HASH_DYNAMIC:
raise RuntimeError("Dynamic bucket extractor is not
configured")
+ from pypaimon.index.dynamic_bucket import MAX_DYNAMIC_BUCKETS
+ upper_bound = MAX_DYNAMIC_BUCKETS if bucket_mode ==
BucketMode.HASH_DYNAMIC else self.table.options.bucket()
+ if not 0 <= bucket < upper_bound:
+ raise ValueError('Bucket id must be between 0 and
{}'.format(upper_bound - 1))
data = self._prepare_arrow_data(data)
if bucket_mode == BucketMode.HASH_DYNAMIC:
- if key_hashes is None:
- partition =
self.row_key_extractor.notify_precomputed_bucket_batch(
- data, bucket
- )
- else:
- partition = (
- self.row_key_extractor
- .notify_precomputed_bucket_hashes_batch(
- data,
- bucket,
- key_hashes,
- new_mappings=new_mappings,
- )
- )
+ partition =
self.row_key_extractor.notify_precomputed_bucket_batch(data, bucket)
else:
- if key_hashes is not None:
- raise ValueError(
- "Precomputed key hashes are only valid for HASH_DYNAMIC
tables"
- )
- if new_mappings is not None:
- raise ValueError(
- "Precomputed new-mapping flags are only valid for "
- "HASH_DYNAMIC tables"
- )
partitions = self.row_key_extractor.extract_partitions_batch(data)
if not partitions:
return
partition = tuple(partitions[0])
- for actual_partition in partitions[1:]:
- if tuple(actual_partition) != partition:
- raise RuntimeError(
- "A precomputed fixed-bucket group contained multiple "
- f"partitions: expected {partition}, got
{actual_partition}"
- )
+ if any(tuple(actual) != partition for actual in partitions):
+ raise ValueError('A precomputed bucket group contained
multiple partitions')
if partition is None:
return
self._write_partition_bucket_batch(partition, bucket, data)
@@ -387,8 +348,10 @@ class TableWrite:
class BatchTableWrite(TableWrite):
- def __init__(self, table, commit_user, static_partition: Optional[dict] =
None):
- super().__init__(table, commit_user, static_partition)
+ def __init__(self, table, commit_user, static_partition: Optional[dict] =
None,
+ *, restore_snapshot_id=None):
+ super().__init__(table, commit_user, static_partition,
+ restore_snapshot_id=restore_snapshot_id)
self.batch_committed = False
def prepare_commit(self) -> List[CommitMessage]:
@@ -400,8 +363,9 @@ class BatchTableWrite(TableWrite):
class StreamTableWrite(TableWrite):
- def __init__(self, table, commit_user):
- super().__init__(table, commit_user, None)
+ def __init__(self, table, commit_user, *, restore_snapshot_id=None):
+ super().__init__(table, commit_user, None,
+ restore_snapshot_id=restore_snapshot_id)
def prepare_commit(self, commit_identifier) -> List[CommitMessage]:
return self._prepare_commit(commit_identifier)
diff --git a/paimon-python/pypaimon/write/write_builder.py
b/paimon-python/pypaimon/write/write_builder.py
index 4cd4f56d64..ba03648b11 100644
--- a/paimon-python/pypaimon/write/write_builder.py
+++ b/paimon-python/pypaimon/write/write_builder.py
@@ -36,6 +36,7 @@ class WriteBuilder(ABC):
self.table: FileStoreTable = table
self.commit_user = self._create_commit_user()
+ self.restore_snapshot_id = None
def new_write(self) -> TableWrite:
"""Returns a table write."""
@@ -46,6 +47,16 @@ class WriteBuilder(ABC):
def new_commit(self) -> TableCommit:
"""Returns a table commit."""
+ def with_restore_snapshot(self, snapshot_id: int):
+ """Restore dynamic-bucket data and HASH state from a snapshot; 0 is
empty."""
+ from pypaimon.table.bucket_mode import BucketMode
+ if self.table.bucket_mode() != BucketMode.HASH_DYNAMIC:
+ raise ValueError('Restore snapshots are only valid for
HASH_DYNAMIC tables')
+ if isinstance(snapshot_id, bool) or not isinstance(snapshot_id, int)
or snapshot_id < 0:
+ raise ValueError('Restore snapshot id must be a nonnegative
integer')
+ self.restore_snapshot_id = snapshot_id
+ return self
+
def _with_commit_user(self, commit_user: str):
"""Reuse the identity of an existing write or commit operation."""
self.commit_user = commit_user
@@ -69,7 +80,8 @@ 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, **kwargs)
+ static_partition, stream,
+
restore_snapshot_id=self.restore_snapshot_id, **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)
@@ -88,7 +100,8 @@ class BatchWriteBuilder(WriteBuilder):
def new_write(self) -> BatchTableWrite:
return (self._native_write(self.static_partition)
- or BatchTableWrite(self.table, self.commit_user,
self.static_partition))
+ or BatchTableWrite(self.table, self.commit_user,
self.static_partition,
+
restore_snapshot_id=self.restore_snapshot_id))
def new_update(self) -> BatchTableUpdate:
return BatchTableUpdate(self.table, self.commit_user)
@@ -102,7 +115,8 @@ class StreamWriteBuilder(WriteBuilder):
def new_write(self) -> StreamTableWrite:
return (self._native_write(stream=True)
- or StreamTableWrite(self.table, self.commit_user))
+ or StreamTableWrite(self.table, self.commit_user,
+
restore_snapshot_id=self.restore_snapshot_id))
def new_update(self) -> StreamTableUpdate:
return StreamTableUpdate(self.table, self.commit_user)
diff --git a/paimon-python/pypaimon/write/writer/data_writer.py
b/paimon-python/pypaimon/write/writer/data_writer.py
index e8f178b974..8c7e7fa45a 100644
--- a/paimon-python/pypaimon/write/writer/data_writer.py
+++ b/paimon-python/pypaimon/write/writer/data_writer.py
@@ -341,8 +341,15 @@ class DataWriter(ABC):
# Read the range without advancing it: the advance belongs with the
# append below, so a retried flush derives the same range.
- min_seq = self.sequence_generator.start
- max_seq = self.sequence_generator.current
+ if self.table.is_primary_key_table and '_SEQUENCE_NUMBER' in
data.schema.names:
+ # PK files can contain a deduplicated or split subset of the
+ # buffer. Java records the range of the rows actually written.
+ sequences = data.column('_SEQUENCE_NUMBER')
+ min_seq = pc.min(sequences).as_py()
+ max_seq = pc.max(sequences).as_py()
+ else:
+ min_seq = self.sequence_generator.start
+ max_seq = self.sequence_generator.current
creation_time = Timestamp.now()
data_meta = self._create_data_file_meta(
file_name=file_name,
