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 c3930a78bd [python] Expose row-id updater through TableUpdate (#10188)
c3930a78bd is described below
commit c3930a78bdcfcb424c4f425d945b6e512b5a7d92
Author: Jingsong Lee <[email protected]>
AuthorDate: Sat Sep 26 15:09:29 2026 +0800
[python] Expose row-id updater through TableUpdate (#10188)
---
paimon-python/README.md | 21 ++++
paimon-python/pypaimon/multimodal/blob_store.py | 10 +-
.../pypaimon/ray/data_evolution_merge_join.py | 41 ++-----
.../pypaimon/table/data_evolution_merge_into.py | 10 +-
paimon-python/pypaimon/tests/blob_table_test.py | 19 +---
paimon-python/pypaimon/tests/table_update_test.py | 39 ++++++-
.../pypaimon/tests/table_upsert_by_key_test.py | 9 +-
paimon-python/pypaimon/tests/vector_table_test.py | 20 +---
.../write/commit/row_id_conflict_rewriter.py | 15 ++-
paimon-python/pypaimon/write/row_id_file_index.py | 121 +++++++++++++++++++++
paimon-python/pypaimon/write/table_update.py | 39 +++++--
.../pypaimon/write/table_update_by_row_id.py | 111 ++-----------------
.../pypaimon/write/table_upsert_by_key.py | 16 +--
paimon-python/pypaimon/write/write_builder.py | 23 ++++
14 files changed, 286 insertions(+), 208 deletions(-)
diff --git a/paimon-python/README.md b/paimon-python/README.md
index d83b947518..f1654fa0e7 100644
--- a/paimon-python/README.md
+++ b/paimon-python/README.md
@@ -105,6 +105,27 @@ pip3 install dist/*.tar.gz
The command will install the package and core dependencies to your local
Python environment.
+# Row ID column updates
+
+For a batch update of selected columns in a data-evolution table with row
+tracking, pass an Arrow table containing `_ROW_ID` and the columns to update.
+Create the updater and committer from the same builder so they share a commit
+user:
+
+```python
+builder = table.new_batch_write_builder()
+updater = builder.new_update().new_update_by_row_id()
+messages = updater.update_columns(updates, ["name"])
+commit = builder.new_commit()
+try:
+ commit.commit(messages)
+finally:
+ commit.close()
+```
+
+For stream updates, use `table.new_stream_write_builder()` and pass the stream
+commit identifier to `new_update().new_update_by_row_id(commit_identifier)`.
+
# Parquet page-index reads
For row-tracking tables with a Parquet OffsetIndex, PyPaimon can read a
diff --git a/paimon-python/pypaimon/multimodal/blob_store.py
b/paimon-python/pypaimon/multimodal/blob_store.py
index 53831750b5..061402ba87 100644
--- a/paimon-python/pypaimon/multimodal/blob_store.py
+++ b/paimon-python/pypaimon/multimodal/blob_store.py
@@ -24,11 +24,9 @@ from typing import BinaryIO, Dict, Iterable, List, Mapping,
Optional, Sequence
from pypaimon.common.options.core_options import CoreOptions
from pypaimon.common.predicate_builder import PredicateBuilder
from pypaimon.schema.data_types import is_blob_type
-from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
from pypaimon.table.row.blob import Blob, BlobData, BlobDescriptor
from pypaimon.table.row.generic_row import GenericRow
from pypaimon.table.special_fields import SpecialFields
-from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
_RANGE_PATTERN = re.compile(r"^bytes=(\d*)-(\d*)$")
@@ -318,11 +316,9 @@ class BlobStore:
write_builder = self._raw_table.new_batch_write_builder()
table_commit = write_builder.new_commit()
try:
- messages = TableUpdateByRowId(
- self._raw_table,
- write_builder.commit_user,
- BATCH_COMMIT_IDENTIFIER,
- ).update_rows_columns(rows, row_ids_by_row, update_columns)
+ updater = write_builder.new_update().new_update_by_row_id()
+ messages = updater.update_rows_columns(
+ rows, row_ids_by_row, update_columns)
if messages:
table_commit.commit(messages)
finally:
diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py
b/paimon-python/pypaimon/ray/data_evolution_merge_join.py
index 835c9a981b..b86ac7b363 100644
--- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py
+++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py
@@ -345,9 +345,8 @@ def _self_merge_aliases(batch: pa.Table, row_id_name: str)
-> pa.Table:
def _apply_self_merge_update_group(context, file_group, collect_row_ids):
"""Read, transform, and stage one complete first-row-id file group."""
from pypaimon.read.table_read import TableRead
- from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
from pypaimon.write.file_store_commit import _abort_commit_messages
- from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
+ from pypaimon.write.row_id_file_index import RowIdFileIndex
table_read = TableRead(
context.scan_table,
@@ -387,16 +386,11 @@ def _apply_self_merge_update_group(context, file_group,
collect_row_ids):
if collect_row_ids else []
)
- import uuid
- files_info = TableUpdateByRowId._files_info_from_splits(
+ files_info = RowIdFileIndex.from_splits(
context.snapshot_id, [file_group],
)
- updater = TableUpdateByRowId(
- context.table,
- "_self_merge_group_" + uuid.uuid4().hex[:8],
- BATCH_COMMIT_IDENTIFIER,
- _precomputed_files_info=files_info,
- )
+ updater =
context.table.new_batch_write_builder().new_update().new_update_by_row_id(
+ _precomputed_files_info=files_info)
try:
messages = updater.update_columns(updates, context.update_cols)
except Exception:
@@ -698,14 +692,11 @@ def distributed_update_apply(
) -> Tuple[list, int, list]:
import numpy as np
import pickle
- import uuid
import pyarrow.compute as pc
import ray
- from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
from pypaimon.table.special_fields import SpecialFields
- from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
row_id_name = SpecialFields.ROW_ID.name
cols = list(write_update_cols)
@@ -724,11 +715,7 @@ def distributed_update_apply(
table.copy({CoreOptions.SCAN_SNAPSHOT_ID.key(): str(base_snapshot_id)})
if base_snapshot_id is not None else table
)
- planner = TableUpdateByRowId(
- scan_table,
- "_merge_into_planner_" + uuid.uuid4().hex[:8],
- BATCH_COMMIT_IDENTIFIER,
- )
+ planner =
scan_table.new_batch_write_builder().new_update().new_update_by_row_id()
sorted_first_row_ids = list(planner.first_row_ids)
if not sorted_first_row_ids:
return [], 0, []
@@ -843,12 +830,9 @@ def distributed_update_apply(
for_update.column(row_id_name).to_pylist()
if collect_row_ids else []
)
- worker = TableUpdateByRowId(
- captured_table,
- "_merge_into_shard_" + uuid.uuid4().hex[:8],
- BATCH_COMMIT_IDENTIFIER,
- _precomputed_files_info=ray.get(precomputed_info_ref),
- )
+ update = captured_table.new_batch_write_builder().new_update()
+ worker = update.new_update_by_row_id(
+ _precomputed_files_info=ray.get(precomputed_info_ref))
msgs = worker.update_columns(for_update, list(captured_cols))
return pa.Table.from_pydict({
"msgs_blob": [pickle.dumps(msgs)],
@@ -915,16 +899,13 @@ def distributed_read_by_row_id(
is empty. Read-side mirror of ``distributed_update_apply``.
"""
import numpy as np
- import uuid
import ray
from pypaimon.globalindex.indexed_split import IndexedSplit
from pypaimon.read.split import DataSplit
- from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
from pypaimon.table.special_fields import SpecialFields
from pypaimon.utils.range import Range
- from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
row_id_name = SpecialFields.ROW_ID.name
read_cols = list(projection)
@@ -935,11 +916,7 @@ def distributed_read_by_row_id(
empty_out = _read_output_schema(table, read_cols).empty_table()
# The caller pinned the resolved snapshot, including any retained tag
metadata.
- planner = TableUpdateByRowId(
- table,
- "_read_by_row_id_planner_" + uuid.uuid4().hex[:8],
- BATCH_COMMIT_IDENTIFIER,
- )
+ planner =
table.new_batch_write_builder().new_update().new_update_by_row_id()
sorted_first_row_ids = list(planner.first_row_ids)
if not sorted_first_row_ids:
return None
diff --git a/paimon-python/pypaimon/table/data_evolution_merge_into.py
b/paimon-python/pypaimon/table/data_evolution_merge_into.py
index 7439823d08..58c02d4517 100644
--- a/paimon-python/pypaimon/table/data_evolution_merge_into.py
+++ b/paimon-python/pypaimon/table/data_evolution_merge_into.py
@@ -55,7 +55,6 @@ from pypaimon.table.special_fields import SpecialFields
from pypaimon.write.commit_message import CommitMessage
from pypaimon.write.table_write import BatchTableWrite, StreamTableWrite
from pypaimon.write.table_delete import TableDeleteByRowId
-from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
__all__ = [
"merge_into",
@@ -618,11 +617,10 @@ def _prepare_commit_messages(
update_snapshot_table = _copy_at_snapshot(
table, base_snapshot.id if base_snapshot is not None else None
)
- updater = TableUpdateByRowId(
- update_snapshot_table,
- commit_user,
- commit_identifier,
- )
+ from pypaimon.write.write_builder import _new_update_by_row_id
+
+ updater = _new_update_by_row_id(
+ update_snapshot_table, commit_user, commit_identifier)
update_msgs = updater.update_columns(
update_table, list(update_cols_union)
)
diff --git a/paimon-python/pypaimon/tests/blob_table_test.py
b/paimon-python/pypaimon/tests/blob_table_test.py
index 6b4133242c..61d26111f6 100755
--- a/paimon-python/pypaimon/tests/blob_table_test.py
+++ b/paimon-python/pypaimon/tests/blob_table_test.py
@@ -5240,9 +5240,6 @@ class DedicatedFormatWriterTest(unittest.TestCase):
wb.new_commit().commit(w.prepare_commit())
w.close()
- from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
- from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
-
table = self.catalog.get_table(table_name)
rb = table.new_read_builder()
rb = rb.with_projection(['name', '_ROW_ID'])
@@ -5253,11 +5250,10 @@ class DedicatedFormatWriterTest(unittest.TestCase):
'_ROW_ID': source.column('_ROW_ID'),
'name': pa.array(['updated', 'updated'], type=pa.string()),
})
- updater = TableUpdateByRowId(
- table, '_test_', BATCH_COMMIT_IDENTIFIER,
- )
+ update_builder = table.new_batch_write_builder()
+ updater = update_builder.new_update().new_update_by_row_id()
msgs = updater.update_columns(update_data, ['name'])
- table.new_batch_write_builder().new_commit().commit(msgs)
+ update_builder.new_commit().commit(msgs)
table = self.catalog.get_table(table_name)
rb = table.new_read_builder()
@@ -5267,8 +5263,6 @@ class DedicatedFormatWriterTest(unittest.TestCase):
def
test_blob_table_partial_update_non_blob_column_with_rolling_files(self):
from pypaimon.manifest.schema.data_file_meta import DataFileMeta
- from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
- from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
pa_schema = pa.schema([
('id', pa.int32()),
@@ -5322,9 +5316,8 @@ class DedicatedFormatWriterTest(unittest.TestCase):
'_ROW_ID': source.column('_ROW_ID'),
'name': pa.array(['updated'] * source.num_rows, type=pa.string()),
})
- updater = TableUpdateByRowId(
- table, '_test_', BATCH_COMMIT_IDENTIFIER,
- )
+ update_builder = table.new_batch_write_builder()
+ updater = update_builder.new_update().new_update_by_row_id()
msgs = updater.update_columns(update_data, ['name'])
update_normal_files = [
f for msg in msgs for f in msg.new_files
@@ -5335,7 +5328,7 @@ class DedicatedFormatWriterTest(unittest.TestCase):
for file in update_normal_files:
self.assertEqual(file.min_sequence_number, 0)
self.assertEqual(file.max_sequence_number, file.row_count - 1)
- table.new_batch_write_builder().new_commit().commit(msgs)
+ update_builder.new_commit().commit(msgs)
table = self.catalog.get_table(table_name)
rb = table.new_read_builder().with_projection(['id', 'name'])
diff --git a/paimon-python/pypaimon/tests/table_update_test.py
b/paimon-python/pypaimon/tests/table_update_test.py
index 092990ab98..25116b3c3b 100644
--- a/paimon-python/pypaimon/tests/table_update_test.py
+++ b/paimon-python/pypaimon/tests/table_update_test.py
@@ -26,6 +26,7 @@ import pyarrow as pa
import pytest
from pypaimon.read.read_builder import ReadBuilder
+from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
from pypaimon.tests.data_evolution_test_helpers import (
BatchModeMixin,
DataEvolutionTestBase,
@@ -33,6 +34,7 @@ from pypaimon.tests.data_evolution_test_helpers import (
)
from pypaimon.write.table_update import BatchTableUpdate
from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
+from pypaimon.write.write_builder import BatchWriteBuilder, StreamWriteBuilder
# ======================================================================
@@ -40,6 +42,37 @@ from pypaimon.write.table_update_by_row_id import
TableUpdateByRowId
# ======================================================================
[email protected]("stream, commit_identifier", [
+ (False, BATCH_COMMIT_IDENTIFIER),
+ (True, 123),
+])
+def test_write_builder_creates_row_id_updater_with_commit_identity(
+ stream, commit_identifier):
+ table = mock.Mock()
+ table.options.commit_user_prefix.return_value = None
+ files_info = mock.Mock(
+ snapshot_id=42,
+ first_row_ids=[0],
+ first_row_id_index={},
+ valid_row_id_ranges=[],
+ )
+
+ builder = StreamWriteBuilder(table) if stream else BatchWriteBuilder(table)
+ update = builder.new_update()
+ if stream:
+ updater = update.new_update_by_row_id(
+ commit_identifier, _precomputed_files_info=files_info)
+ else:
+ updater = update.new_update_by_row_id(
+ _precomputed_files_info=files_info)
+
+ assert updater.table is table
+ assert updater.commit_user == update.commit_user == builder.commit_user
+ assert updater.commit_identifier == commit_identifier
+ assert updater.snapshot_id == 42
+ table.new_read_builder.assert_not_called()
+
+
def test_batch_row_id_update_batches_reuse_file_index():
table = mock.MagicMock()
table.field_names = ["value"]
@@ -48,8 +81,8 @@ def test_batch_row_id_update_batches_reuse_file_index():
pa.table({"_ROW_ID": [1], "value": [20]}),
]
- with mock.patch(
- "pypaimon.write.table_update.TableUpdateByRowId") as factory:
+ with mock.patch.object(
+ BatchTableUpdate, "_new_row_id_updater") as factory:
updater = factory.return_value
updater.commit_messages = []
@@ -64,7 +97,7 @@ def test_batch_row_id_update_batches_reuse_file_index():
.update_by_arrow_batches_with_row_id(iter(batches))
)
- factory.assert_called_once()
+ factory.assert_called_once_with(BATCH_COMMIT_IDENTIFIER)
assert updater.update_columns.call_count == 2
assert messages == [
(batches[0], ["value"]),
diff --git a/paimon-python/pypaimon/tests/table_upsert_by_key_test.py
b/paimon-python/pypaimon/tests/table_upsert_by_key_test.py
index 564d61f3ed..68d26e86e8 100644
--- a/paimon-python/pypaimon/tests/table_upsert_by_key_test.py
+++ b/paimon-python/pypaimon/tests/table_upsert_by_key_test.py
@@ -675,13 +675,14 @@ class _TableUpsertByKeyTestBase(DataEvolutionTestBase):
)
self._compact_all_data_files(table)
+ from pypaimon.write.row_id_file_index import RowIdFileIndex
from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
- original_build = TableUpdateByRowId._files_info_from_entries
+ original_build = RowIdFileIndex.from_entries
advanced = [False]
def build_after_concurrent_compaction(
- updater_cls, current_table, snapshot_id, entries):
+ index_cls, current_table, snapshot_id, entries):
if not advanced[0]:
advanced[0] = True
self._write_arrow(table, pa.Table.from_pydict({
@@ -698,8 +699,8 @@ class _TableUpsertByKeyTestBase(DataEvolutionTestBase):
'_load_existing_files_info',
side_effect=AssertionError("unexpected snapshot scan"),
), mock.patch.object(
- TableUpdateByRowId,
- '_files_info_from_entries',
+ RowIdFileIndex,
+ 'from_entries',
classmethod(build_after_concurrent_compaction)):
commit = wb.new_commit()
self._apply_commit(commit, messages, commit_identifier)
diff --git a/paimon-python/pypaimon/tests/vector_table_test.py
b/paimon-python/pypaimon/tests/vector_table_test.py
index 96725305fc..574dc747c2 100644
--- a/paimon-python/pypaimon/tests/vector_table_test.py
+++ b/paimon-python/pypaimon/tests/vector_table_test.py
@@ -391,9 +391,6 @@ class VectorTableWriteReadTest(unittest.TestCase):
wb.new_commit().commit(initial_messages)
w.close()
- from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
- from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
-
table = self.catalog.get_table(table_name)
rb = table.new_read_builder()
rb = rb.with_projection(['name', '_ROW_ID'])
@@ -404,11 +401,10 @@ class VectorTableWriteReadTest(unittest.TestCase):
'_ROW_ID': source.column('_ROW_ID'),
'name': pa.array(['updated', 'updated'], type=pa.string()),
})
- updater = TableUpdateByRowId(
- table, '_test_', BATCH_COMMIT_IDENTIFIER,
- )
+ update_builder = table.new_batch_write_builder()
+ updater = update_builder.new_update().new_update_by_row_id()
msgs = updater.update_columns(update_data, ['name'])
- table.new_batch_write_builder().new_commit().commit(msgs)
+ update_builder.new_commit().commit(msgs)
table = self.catalog.get_table(table_name)
rb = table.new_read_builder()
@@ -512,9 +508,6 @@ class VectorTableWriteReadTest(unittest.TestCase):
)
def
test_vector_table_partial_update_non_vector_column_with_rolling_files(self):
- from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
- from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
-
vector_schema = pa.schema([
('id', pa.int32()),
('name', pa.string()),
@@ -567,9 +560,8 @@ class VectorTableWriteReadTest(unittest.TestCase):
'_ROW_ID': source.column('_ROW_ID'),
'name': pa.array(['updated'] * source.num_rows, type=pa.string()),
})
- updater = TableUpdateByRowId(
- table, '_test_', BATCH_COMMIT_IDENTIFIER,
- )
+ update_builder = table.new_batch_write_builder()
+ updater = update_builder.new_update().new_update_by_row_id()
msgs = updater.update_columns(update_data, ['name'])
update_normal_files = [
f for msg in msgs for f in msg.new_files
@@ -579,7 +571,7 @@ class VectorTableWriteReadTest(unittest.TestCase):
for file in update_normal_files:
self.assertEqual(file.min_sequence_number, 0)
self.assertEqual(file.max_sequence_number, file.row_count - 1)
- table.new_batch_write_builder().new_commit().commit(msgs)
+ update_builder.new_commit().commit(msgs)
table = self.catalog.get_table(table_name)
rb = table.new_read_builder().with_projection(['id', 'name'])
diff --git a/paimon-python/pypaimon/write/commit/row_id_conflict_rewriter.py
b/paimon-python/pypaimon/write/commit/row_id_conflict_rewriter.py
index bf361d7d62..898fe909df 100644
--- a/paimon-python/pypaimon/write/commit/row_id_conflict_rewriter.py
+++ b/paimon-python/pypaimon/write/commit/row_id_conflict_rewriter.py
@@ -27,7 +27,7 @@ from pypaimon.table.row.generic_row import GenericRow
from pypaimon.table.special_fields import SpecialFields
from pypaimon.utils.range import Range
from pypaimon.write.commit_message import CommitMessage
-from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
+from pypaimon.write.row_id_file_index import RowIdFileIndex
@dataclass
@@ -103,7 +103,7 @@ class RowIdConflictRewriter:
new_messages = []
try:
- files_info = TableUpdateByRowId._files_info_from_entries(
+ files_info = RowIdFileIndex.from_entries(
self.table,
latest_snapshot.id,
base_entries,
@@ -126,12 +126,11 @@ class RowIdConflictRewriter:
"overlapping staged files for columns {}.".format(
column_names)
)
- updater = TableUpdateByRowId(
- self.table,
- self.commit_user,
- self.commit_identifier,
- _precomputed_files_info=files_info,
- )
+ from pypaimon.write.write_builder import _new_update_by_row_id
+
+ updater = _new_update_by_row_id(
+ self.table, self.commit_user, self.commit_identifier,
+ _precomputed_files_info=files_info)
try:
new_messages.extend(
updater.update_columns(update_data, column_names)
diff --git a/paimon-python/pypaimon/write/row_id_file_index.py
b/paimon-python/pypaimon/write/row_id_file_index.py
new file mode 100644
index 0000000000..52975ae09c
--- /dev/null
+++ b/paimon-python/pypaimon/write/row_id_file_index.py
@@ -0,0 +1,121 @@
+# 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.
+
+"""Snapshot file index used to route row-id updates."""
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Tuple
+
+from pypaimon.manifest.schema.data_file_meta import DataFileMeta
+from pypaimon.manifest.schema.manifest_entry import ManifestEntry
+from pypaimon.read.scanner.data_evolution_split_generator import (
+ DataEvolutionSplitGenerator,
+)
+from pypaimon.read.split import DataSplit
+from pypaimon.utils.range import Range
+
+
+@dataclass(frozen=True)
+class RowIdFileIndex:
+ """Snapshot files keyed by the first row ID of each data file group."""
+
+ snapshot_id: int
+ first_row_ids: List[int]
+ first_row_id_index: Dict[int, Tuple[DataSplit, List[DataFileMeta]]] = (
+ field(default_factory=dict)
+ )
+ valid_row_id_ranges: List[Range] = field(default_factory=list)
+
+ @classmethod
+ def from_table(cls, table) -> "RowIdFileIndex":
+ """Plan the table's write snapshot and index its splits."""
+ scan = table.new_read_builder().new_scan()
+ plan = scan.plan_for_write()
+ snapshot_id = plan.snapshot_id if plan.snapshot_id is not None else -1
+ return cls.from_splits(snapshot_id, plan.splits())
+
+ @classmethod
+ def from_entries(
+ cls, table, snapshot_id: int, entries: List[ManifestEntry]
+ ) -> "RowIdFileIndex":
+ """Index an already resolved snapshot entry set."""
+ splits = DataEvolutionSplitGenerator(
+ table,
+ table.options.source_split_target_size(),
+ table.options.source_split_open_file_cost(),
+ ).create_splits(entries)
+ return cls.from_splits(snapshot_id, splits)
+
+ @classmethod
+ def from_splits(
+ cls, snapshot_id: int, splits: List[DataSplit]
+ ) -> "RowIdFileIndex":
+ index: Dict[int, Tuple[DataSplit, List[DataFileMeta]]] = {}
+ row_id_ranges: List[Range] = []
+ for split in splits:
+ files_with_row_id = [
+ file for file in split.files if file.first_row_id is not None
+ ]
+ data_files = [
+ file for file in files_with_row_id
+ if not DataFileMeta.is_blob_file(file.file_name)
+ ]
+ for file in split.files:
+ if (
+ file.first_row_id is None
+ or DataFileMeta.is_blob_file(file.file_name)
+ ):
+ continue
+ row_id_ranges.append(file.row_id_range())
+ for file in data_files:
+ target_files = [
+ target_file
+ for target_file in files_with_row_id
+ if cls._overlaps(
+ file.row_id_range(), target_file.row_id_range()
+ )
+ ]
+
+ entry = index.get(file.first_row_id)
+ if entry is None:
+ index[file.first_row_id] = (split, target_files)
+ else:
+ existing_files = entry[1]
+ existing_names = {
+ existing.file_name for existing in existing_files
+ }
+ existing_files.extend(
+ target_file
+ for target_file in target_files
+ if target_file.file_name not in existing_names
+ )
+
+ if row_id_ranges:
+ merged = Range.sort_and_merge_overlap(row_id_ranges, True, True)
+ else:
+ merged = []
+
+ return cls(
+ snapshot_id=snapshot_id,
+ first_row_ids=sorted(index.keys()),
+ first_row_id_index=index,
+ valid_row_id_ranges=merged,
+ )
+
+ @staticmethod
+ def _overlaps(left: Range, right: Range) -> bool:
+ return left.from_ <= right.to and right.from_ <= left.to
diff --git a/paimon-python/pypaimon/write/table_update.py
b/paimon-python/pypaimon/write/table_update.py
index 922b1bc8c3..494127d791 100644
--- a/paimon-python/pypaimon/write/table_update.py
+++ b/paimon-python/pypaimon/write/table_update.py
@@ -46,6 +46,7 @@ from pypaimon.snapshot.time_travel_util import SCAN_KEYS,
TimeTravelUtil
from pypaimon.table.special_fields import SpecialFields
from pypaimon.write.commit_message import CommitMessage
from pypaimon.write.file_store_commit import _abort_commit_messages
+from pypaimon.write.row_id_file_index import RowIdFileIndex
from pypaimon.write.table_delete import TableDeleteByRowId
from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
from pypaimon.write.table_upsert_by_key import TableUpsertByKey
@@ -125,6 +126,16 @@ class TableUpdate:
self.update_cols = None
self.projection = None
+ def _new_row_id_updater(
+ self, commit_identifier: int, _precomputed_files_info=None
+ ) -> TableUpdateByRowId:
+ return TableUpdateByRowId(
+ self.table,
+ self.commit_user,
+ commit_identifier,
+ _precomputed_files_info=_precomputed_files_info,
+ )
+
def with_update_type(self, update_cols: List[str]):
update_cols = list(dict.fromkeys(update_cols))
for col in update_cols:
@@ -163,9 +174,8 @@ class TableUpdate:
cols = self.update_cols if self.update_cols is not None else [
c for c in table.column_names if c != SpecialFields.ROW_ID.name
]
- return TableUpdateByRowId(
- self.table, self.commit_user, commit_identifier,
- ).update_columns(table, cols)
+ return self._new_row_id_updater(commit_identifier).update_columns(
+ table, cols)
def _update_by_arrow_batches_with_row_id(
self, tables: Iterable[pa.Table], commit_identifier: int
@@ -178,8 +188,7 @@ class TableUpdate:
if c != SpecialFields.ROW_ID.name
]
if updater is None:
- updater = TableUpdateByRowId(
- self.table, self.commit_user, commit_identifier)
+ updater = self._new_row_id_updater(commit_identifier)
updater.update_columns(table, cols)
return [] if updater is None else updater.commit_messages
except Exception:
@@ -283,14 +292,11 @@ class TableUpdate:
plan = read_builder.new_scan().plan_for_write()
splits = plan.splits()
snapshot_id = plan.snapshot_id if plan.snapshot_id is not None else -1
- files_info = TableUpdateByRowId._files_info_from_splits(
+ files_info = RowIdFileIndex.from_splits(
snapshot_id, splits
)
table_read = read_builder.new_read()
- updater = TableUpdateByRowId(
- self.table, self.commit_user, commit_identifier,
- _precomputed_files_info=files_info,
- )
+ updater = self._new_row_id_updater(commit_identifier, files_info)
try:
if has_array:
matched = table_read.to_arrow(splits)
@@ -641,6 +647,12 @@ class BatchTableUpdate(TableUpdate):
"""Batch-mode table update; commit messages always use
:data:`BATCH_COMMIT_IDENTIFIER`."""
+ def new_update_by_row_id(
+ self, _precomputed_files_info=None) -> TableUpdateByRowId:
+ """Create a row-id updater using this batch update's commit user."""
+ return self._new_row_id_updater(
+ BATCH_COMMIT_IDENTIFIER, _precomputed_files_info)
+
def update_by_arrow_with_row_id(self, table: pa.Table) ->
List[CommitMessage]:
"""Apply column updates keyed by ``_ROW_ID`` to existing rows."""
return self._update_by_arrow_with_row_id(table,
BATCH_COMMIT_IDENTIFIER)
@@ -720,6 +732,13 @@ class StreamTableUpdate(TableUpdate):
"""Stream-mode table update; the same instance may drive many rounds,
each tagged with its own ``commit_identifier``."""
+ def new_update_by_row_id(
+ self, commit_identifier: int,
+ _precomputed_files_info=None) -> TableUpdateByRowId:
+ """Create a row-id updater for a stream commit identifier."""
+ return self._new_row_id_updater(
+ commit_identifier, _precomputed_files_info)
+
def update_by_arrow_with_row_id(
self, table: pa.Table, commit_identifier: int
) -> List[CommitMessage]:
diff --git a/paimon-python/pypaimon/write/table_update_by_row_id.py
b/paimon-python/pypaimon/write/table_update_by_row_id.py
index 9e6275fd4a..a8af846c9c 100644
--- a/paimon-python/pypaimon/write/table_update_by_row_id.py
+++ b/paimon-python/pypaimon/write/table_update_by_row_id.py
@@ -17,7 +17,6 @@
import bisect
import uuid
-from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Set, Tuple
import numpy as np
@@ -26,11 +25,7 @@ import pyarrow.compute as pc
from pypaimon.common.options.core_options import ChangelogProducer, CoreOptions
from pypaimon.manifest.schema.data_file_meta import DataFileMeta
-from pypaimon.manifest.schema.manifest_entry import ManifestEntry
from pypaimon.manifest.schema.simple_stats import SimpleStats
-from pypaimon.read.scanner.data_evolution_split_generator import (
- DataEvolutionSplitGenerator,
-)
from pypaimon.read.split import DataSplit
from pypaimon.read.table_read import TableRead
from pypaimon.schema.data_types import (
@@ -43,9 +38,9 @@ from pypaimon.schema.data_types import (
from pypaimon.table.row.blob import Blob
from pypaimon.table.row.generic_row import GenericRow
from pypaimon.table.special_fields import SpecialFields
-from pypaimon.utils.range import Range
from pypaimon.write.commit_message import CommitMessage
from pypaimon.write.file_store_write import FileStoreWrite
+from pypaimon.write.row_id_file_index import RowIdFileIndex
from pypaimon.write.row_utils import (
require_columns,
row_to_named_values,
@@ -61,21 +56,6 @@ _ARROW_MAJOR = int(pa.__version__.split('.')[0])
_DEFAULT_PARQUET_BLOCK_SIZE = 128 * 1024 * 1024
-@dataclass(frozen=True)
-class _FilesInfo:
- """Snapshot view of target data files keyed by first_row_id.
-
- Built once per merge by the driver and broadcast to workers so each task
- avoids re-scanning the manifest.
- """
- snapshot_id: int
- first_row_ids: List[int]
- first_row_id_index: Dict[int, Tuple[DataSplit, List[DataFileMeta]]] = (
- field(default_factory=dict)
- )
- valid_row_id_ranges: List[Range] = field(default_factory=list)
-
-
class _RowIdUpdateFileWriter:
"""Write one plain-Parquet update file for a row-id file group."""
@@ -223,7 +203,7 @@ class TableUpdateByRowId:
def __init__(
self, table, commit_user: str, commit_identifier: int,
- _precomputed_files_info: Optional[_FilesInfo] = None,
+ _precomputed_files_info: Optional[RowIdFileIndex] = None,
):
from pypaimon.table.file_store_table import FileStoreTable
@@ -240,101 +220,24 @@ class TableUpdateByRowId:
self.commit_messages: List[CommitMessage] = []
self._updated_first_row_ids_by_column: Dict[str, Set[int]] = {}
- def _snapshot_files_info(self) -> _FilesInfo:
+ def _snapshot_files_info(self) -> RowIdFileIndex:
"""Return the already loaded snapshot file index for broadcast."""
- return _FilesInfo(
+ return RowIdFileIndex(
snapshot_id=self.snapshot_id,
first_row_ids=self.first_row_ids,
first_row_id_index=self._first_row_id_index,
valid_row_id_ranges=self.valid_row_id_ranges,
)
- def _load_existing_files_info(self) -> _FilesInfo:
+ def _load_existing_files_info(self) -> RowIdFileIndex:
"""Scan the latest snapshot once and index files by ``first_row_id``.
- Returns a :class:`_FilesInfo` whose ``first_row_id_index`` maps each
+ Returns a :class:`RowIdFileIndex` whose ``first_row_id_index`` maps
each
``first_row_id`` to the owning split and the list of files with that
id (a single id may belong to multiple files when data evolution has
split a logical row range).
"""
- scan = self.table.new_read_builder().new_scan()
- plan = scan.plan_for_write()
- snapshot_id = plan.snapshot_id if plan.snapshot_id is not None else -1
- return self._files_info_from_splits(snapshot_id, plan.splits())
-
- @classmethod
- def _files_info_from_entries(
- cls,
- table,
- snapshot_id: int,
- entries: List[ManifestEntry],
- ) -> _FilesInfo:
- """Build a file index from an already resolved snapshot entry set."""
- splits = DataEvolutionSplitGenerator(
- table,
- table.options.source_split_target_size(),
- table.options.source_split_open_file_cost(),
- ).create_splits(entries)
- return cls._files_info_from_splits(snapshot_id, splits)
-
- @classmethod
- def _files_info_from_splits(
- cls, snapshot_id: int, splits: List[DataSplit]) -> _FilesInfo:
- index: Dict[int, Tuple[DataSplit, List[DataFileMeta]]] = {}
- row_id_ranges: List[Range] = []
- for split in splits:
- files_with_row_id = [
- file for file in split.files if file.first_row_id is not None
- ]
- data_files = [
- file for file in files_with_row_id
- if not DataFileMeta.is_blob_file(file.file_name)
- ]
- for file in split.files:
- if (
- file.first_row_id is None
- or DataFileMeta.is_blob_file(file.file_name)
- ):
- continue
- row_id_ranges.append(file.row_id_range())
- for file in data_files:
- target_files = [
- target_file
- for target_file in files_with_row_id
- if cls._overlaps(
- file.row_id_range(), target_file.row_id_range()
- )
- ]
-
- entry = index.get(file.first_row_id)
- if entry is None:
- index[file.first_row_id] = (split, target_files)
- else:
- existing_files = entry[1]
- existing_names = {
- existing.file_name for existing in existing_files
- }
- existing_files.extend(
- target_file
- for target_file in target_files
- if target_file.file_name not in existing_names
- )
-
- if row_id_ranges:
- merged = Range.sort_and_merge_overlap(row_id_ranges, True, True)
- else:
- merged = []
-
- return _FilesInfo(
- snapshot_id=snapshot_id,
- first_row_ids=sorted(index.keys()),
- first_row_id_index=index,
- valid_row_id_ranges=merged,
- )
-
- @staticmethod
- def _overlaps(left: Range, right: Range) -> bool:
- return left.from_ <= right.to and right.from_ <= left.to
+ return RowIdFileIndex.from_table(self.table)
def update_columns(self, data: pa.Table, column_names: List[str]) ->
List[CommitMessage]:
"""
diff --git a/paimon-python/pypaimon/write/table_upsert_by_key.py
b/paimon-python/pypaimon/write/table_upsert_by_key.py
index 4981b8187f..aebbdf6176 100644
--- a/paimon-python/pypaimon/write/table_upsert_by_key.py
+++ b/paimon-python/pypaimon/write/table_upsert_by_key.py
@@ -26,7 +26,6 @@ from pypaimon.table.row.internal_row import InternalRow
from pypaimon.table.special_fields import SpecialFields
from pypaimon.write.commit_message import CommitMessage
from pypaimon.write.row_utils import require_columns, row_to_named_values
-from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
from pypaimon.write.table_write import StreamTableWrite
# Composite key is represented as a tuple of values
@@ -54,6 +53,12 @@ class TableUpsertByKey:
self.commit_user = commit_user
self.commit_identifier = commit_identifier
+ def _new_row_id_updater(self):
+ from pypaimon.write.write_builder import _new_update_by_row_id
+
+ return _new_update_by_row_id(
+ self.table, self.commit_user, self.commit_identifier)
+
def upsert(self, data: pa.Table, upsert_keys: List[str],
update_cols: Optional[List[str]] = None) -> List[CommitMessage]:
"""
@@ -200,9 +205,7 @@ class TableUpsertByKey:
)
for _, values_by_name in matched_items:
require_columns(values_by_name, cols_to_update,
"upsert_by_key")
- commit_messages.extend(TableUpdateByRowId(
- self.table, self.commit_user, self.commit_identifier,
- ).update_rows_columns(
+
commit_messages.extend(self._new_row_id_updater().update_rows_columns(
[row for row, _ in matched_items],
matched_row_ids,
cols_to_update,
@@ -551,9 +554,8 @@ class TableUpsertByKey:
)
cols_to_update = list(update_cols) if update_cols else
list(self.table.field_names)
- return TableUpdateByRowId(
- self.table, self.commit_user, self.commit_identifier,
- ).update_columns(update_data, cols_to_update)
+ return self._new_row_id_updater().update_columns(
+ update_data, cols_to_update)
def _do_appends(
self,
diff --git a/paimon-python/pypaimon/write/write_builder.py
b/paimon-python/pypaimon/write/write_builder.py
index 0f79308ab9..bfe76659f9 100644
--- a/paimon-python/pypaimon/write/write_builder.py
+++ b/paimon-python/pypaimon/write/write_builder.py
@@ -46,6 +46,11 @@ class WriteBuilder(ABC):
def new_commit(self) -> TableCommit:
"""Returns a table commit."""
+ def _with_commit_user(self, commit_user: str):
+ """Reuse the identity of an existing write or commit operation."""
+ self.commit_user = commit_user
+ return self
+
def _create_commit_user(self):
commit_user_prefix = self.table.options.commit_user_prefix()
if commit_user_prefix is not None:
@@ -111,3 +116,21 @@ class StreamWriteBuilder(WriteBuilder):
def new_commit(self) -> StreamTableCommit:
commit = StreamTableCommit(self.table, self.commit_user)
return commit
+
+
+def _new_update_by_row_id(
+ table, commit_user: str, commit_identifier: int,
+ _precomputed_files_info=None):
+ """Build an internal updater with an operation's existing commit
identity."""
+ from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
+
+ if commit_identifier == BATCH_COMMIT_IDENTIFIER:
+ return (table.new_batch_write_builder()
+ ._with_commit_user(commit_user)
+ .new_update()
+ .new_update_by_row_id(_precomputed_files_info))
+ return (table.new_stream_write_builder()
+ ._with_commit_user(commit_user)
+ .new_update()
+ .new_update_by_row_id(
+ commit_identifier, _precomputed_files_info))
