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 e3bc30c425 [python] Stream scalar update_by_predicate assignments by
file group (#9240)
e3bc30c425 is described below
commit e3bc30c425ad1cc9d05680cb0cf0272ddec71546
Author: XiaoHongbo <[email protected]>
AuthorDate: Sun Aug 16 14:54:49 2026 +0800
[python] Stream scalar update_by_predicate assignments by file group (#9240)
---
.../pypaimon/tests/file_store_commit_test.py | 18 ++++
paimon-python/pypaimon/tests/table_update_test.py | 100 ++++++++++++++++--
paimon-python/pypaimon/write/file_store_commit.py | 59 ++++++-----
paimon-python/pypaimon/write/table_update.py | 112 +++++++++++++++++----
4 files changed, 240 insertions(+), 49 deletions(-)
diff --git a/paimon-python/pypaimon/tests/file_store_commit_test.py
b/paimon-python/pypaimon/tests/file_store_commit_test.py
index 1ce3fb93c4..1bb2852b78 100644
--- a/paimon-python/pypaimon/tests/file_store_commit_test.py
+++ b/paimon-python/pypaimon/tests/file_store_commit_test.py
@@ -34,10 +34,28 @@ from pypaimon.write.file_store_commit import (
ManifestMergeResult,
RetryResult,
RewriteResult,
+ _abort_commit_messages,
_try_replace_manifest_files,
)
+class TestAbortCommitMessages(unittest.TestCase):
+
+ def test_index_path_failure_does_not_escape_abort(self):
+ table = Mock()
+ table.path_factory.side_effect = RuntimeError("path lookup failed")
+ index_file = Mock(file_name="index-file", external_path=None)
+ message = Mock(
+ new_files=[],
+ changelog_files=[],
+ index_adds=[Mock(index_file=index_file)],
+ )
+
+ with self.assertLogs(
+ 'pypaimon.write.file_store_commit', level='WARNING'):
+ _abort_commit_messages(table, [message])
+
+
@patch('pypaimon.write.file_store_commit.ManifestFileManager')
@patch('pypaimon.write.file_store_commit.ManifestListManager')
class TestFileStoreCommit(unittest.TestCase):
diff --git a/paimon-python/pypaimon/tests/table_update_test.py
b/paimon-python/pypaimon/tests/table_update_test.py
index 588825c36f..0d54d78987 100644
--- a/paimon-python/pypaimon/tests/table_update_test.py
+++ b/paimon-python/pypaimon/tests/table_update_test.py
@@ -153,6 +153,14 @@ class _TableUpdateTestBase(DataEvolutionTestBase):
tc.close()
return msgs
+ @staticmethod
+ def _list_table_files(table):
+ return {
+ os.path.relpath(os.path.join(root, name), table.table_path)
+ for root, _dirs, files in os.walk(table.table_path)
+ for name in files
+ }
+
def _do_delete_by_predicate(self, table, predicate):
wb = self._make_write_builder(table)
tu = wb.new_update()
@@ -231,6 +239,90 @@ class _TableUpdateTestBase(DataEvolutionTestBase):
result['name'].to_pylist(),
)
+ def test_update_by_predicate_processes_one_file_group_at_a_time(self):
+ from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
+
+ table = self._create_seeded_table()
+ self._do_update(
+ table,
+ pa.Table.from_pydict({'_ROW_ID': [0], 'age': [26]}),
+ ['age'],
+ )
+ splits = table.new_read_builder().new_scan().plan_for_write().splits()
+ self.assertEqual(1, len(splits))
+
+ group_sizes = []
+ original = TableUpdateByRowId.update_columns
+
+ def capture(updater, data, columns):
+ group_sizes.append(data.num_rows)
+ return original(updater, data, columns)
+
+ with mock.patch.object(TableUpdateByRowId, 'update_columns', capture):
+ self._do_update_by_predicate(table, None, {'city': 'Updated'})
+
+ self.assertEqual([2, 3], group_sizes)
+ self.assertEqual(
+ [26, 30, 35, 40, 45],
+ self._read_all(table)['age'].to_pylist(),
+ )
+ self.assertEqual(
+ ['Updated'] * 5,
+ self._read_all(table)['city'].to_pylist(),
+ )
+
+ def test_predicate_update_aborts_groups_after_later_failure(self):
+ from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
+
+ table = self._create_seeded_table()
+ self._do_update(
+ table,
+ pa.Table.from_pydict({'_ROW_ID': [0], 'age': [26]}),
+ ['age'],
+ )
+ before_files = self._list_table_files(table)
+ calls = 0
+ original = TableUpdateByRowId.update_columns
+
+ def fail_second_group(updater, data, columns):
+ nonlocal calls
+ calls += 1
+ if calls == 2:
+ raise RuntimeError("second group failed")
+ messages = original(updater, data, columns)
+ return messages
+
+ with mock.patch.object(
+ TableUpdateByRowId, 'update_columns', fail_second_group):
+ with self.assertRaisesRegex(RuntimeError, "second group failed"):
+ self._do_update_by_predicate(
+ table,
+ None,
+ {'city': 'Updated'},
+ )
+
+ self.assertEqual(2, calls)
+ self.assertEqual(before_files, self._list_table_files(table))
+
+ def test_array_assignment_spans_file_groups(self):
+ table = self._create_seeded_table()
+ self._do_update(
+ table,
+ pa.Table.from_pydict({'_ROW_ID': [0], 'age': [26]}),
+ ['age'],
+ )
+
+ self._do_update_by_predicate(
+ table,
+ None,
+ {'age': pa.array([101, 102, 103, 104, 105])},
+ )
+
+ self.assertEqual(
+ [101, 102, 103, 104, 105],
+ self._read_all(table)['age'].to_pylist(),
+ )
+
def test_update_by_predicate_no_match_is_noop(self):
table = self._create_seeded_table()
pb = table.new_read_builder().new_predicate_builder()
@@ -1342,14 +1434,6 @@ class TableUpdateBatchTest(_BatchModeMixin,
_TableUpdateTestBase, unittest.TestC
self.assertEqual(before_files, self._list_table_files(table))
- @staticmethod
- def _list_table_files(table):
- return {
- os.path.relpath(os.path.join(root, name), table.table_path)
- for root, _dirs, files in os.walk(table.table_path)
- for name in files
- }
-
class TableUpdateStreamTest(_StreamModeMixin, _TableUpdateTestBase,
unittest.TestCase):
"""All shared update tests under stream (``StreamWriteBuilder``) semantics,
diff --git a/paimon-python/pypaimon/write/file_store_commit.py
b/paimon-python/pypaimon/write/file_store_commit.py
index b1f212283d..7e2b69f97a 100644
--- a/paimon-python/pypaimon/write/file_store_commit.py
+++ b/paimon-python/pypaimon/write/file_store_commit.py
@@ -55,6 +55,41 @@ from pypaimon.write.commit_message import CommitMessage
logger = logging.getLogger(__name__)
+def _abort_commit_messages(table, commit_messages: List[CommitMessage]):
+ """Delete files created by messages known to be uncommitted."""
+ for message in commit_messages:
+ for file in list(message.new_files) + list(message.changelog_files):
+ path = None
+ try:
+ path = file.external_path or file.file_path
+ if path:
+ table.file_io.delete_quietly(str(path))
+ except Exception as error:
+ logger.warning(
+ "Failed to clean up file %s during abort: %s",
+ path,
+ error,
+ )
+ for entry in message.index_adds:
+ file_name = None
+ try:
+ index_file = entry.index_file
+ file_name = index_file.file_name
+ path = (
+ index_file.external_path
+ or table.path_factory()
+ .global_index_path_factory()
+ .to_path(file_name)
+ )
+ table.file_io.delete_quietly(path)
+ except Exception as error:
+ logger.warning(
+ "Failed to clean up index file %s during abort: %s",
+ file_name,
+ error,
+ )
+
+
class CommitResult:
"""Base class for commit results."""
@@ -1041,29 +1076,7 @@ class FileStoreCommit:
def abort(self, commit_messages: List[CommitMessage]):
"""Abort commit and delete files. Uses external_path if available to
ensure proper scheme handling."""
- for message in commit_messages:
- for file in list(message.new_files) +
list(message.changelog_files):
- try:
- path_to_delete = file.external_path if file.external_path
else file.file_path
- if path_to_delete:
- path_str = str(path_to_delete)
- self.table.file_io.delete_quietly(path_str)
- except Exception as e:
- path_to_delete = file.external_path if file.external_path
else file.file_path
- logger.warning(f"Failed to clean up file {path_to_delete}
during abort: {e}")
- for entry in message.index_adds:
- try:
- file_name = entry.index_file.file_name
- index_path = (
- entry.index_file.external_path
- or self.table.path_factory()
- .global_index_path_factory()
- .to_path(file_name)
- )
- self.table.file_io.delete_quietly(index_path)
- except Exception as e:
- logger.warning(
- f"Failed to clean up index file
{entry.index_file.file_name} during abort: {e}")
+ _abort_commit_messages(self.table, commit_messages)
def close(self):
"""Close the FileStoreCommit and release resources."""
diff --git a/paimon-python/pypaimon/write/table_update.py
b/paimon-python/pypaimon/write/table_update.py
index 9f03f482e3..4760a18516 100644
--- a/paimon-python/pypaimon/write/table_update.py
+++ b/paimon-python/pypaimon/write/table_update.py
@@ -30,17 +30,22 @@ from pypaimon.common.options.core_options import (
from pypaimon.common.predicate import Predicate
from pypaimon.common.predicate_builder import PredicateBuilder
from pypaimon.globalindex import Range
+from pypaimon.globalindex.indexed_split import IndexedSplit
from pypaimon.manifest.index_manifest_entry import IndexManifestEntry
from pypaimon.manifest.index_manifest_file import IndexManifestFile
from pypaimon.manifest.manifest_list_manager import ManifestListManager
from pypaimon.manifest.schema.data_file_meta import DataFileMeta
from pypaimon.read.scanner.file_scanner import FileScanner
+from pypaimon.read.scanner.data_evolution_split_generator import (
+ DataEvolutionSplitGenerator,
+)
from pypaimon.read.split import DataSplit
from pypaimon.schema.data_types import PyarrowFieldParser
from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
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.table_delete import TableDeleteByRowId
from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
from pypaimon.write.table_upsert_by_key import TableUpsertByKey
@@ -233,32 +238,103 @@ class TableUpdate:
values, builds an Arrow update table, then delegates to the existing
row-id update path.
"""
+ has_array = any(
+ isinstance(value, (pa.Array, pa.ChunkedArray))
+ for value in assignments.values()
+ )
self._validate_predicate_update(assignments)
scan_table = self._matched_update_scan_table()
- read_builder = scan_table.new_read_builder()
+ read_builder = scan_table.new_read_builder().with_projection(
+ [SpecialFields.ROW_ID.name]
+ )
if predicate is not None:
read_builder.with_filter(predicate)
- read_builder.with_projection(
- list(scan_table.field_names) + [SpecialFields.ROW_ID.name]
- )
- else:
- read_builder.with_projection([SpecialFields.ROW_ID.name])
-
- scan = read_builder.new_scan()
- splits = scan.plan_for_write().splits()
- matched = read_builder.new_read().to_arrow(splits)
- if matched.num_rows == 0:
- return []
- update_table = self._build_predicate_update_table(
- matched[SpecialFields.ROW_ID.name],
- assignments,
- matched.num_rows,
+ 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(
+ snapshot_id, splits
)
- return TableUpdateByRowId(
+ table_read = read_builder.new_read()
+ updater = TableUpdateByRowId(
self.table, self.commit_user, commit_identifier,
- ).update_columns(update_table, list(assignments.keys()))
+ _precomputed_files_info=files_info,
+ )
+ try:
+ if has_array:
+ matched = table_read.to_arrow(splits)
+ if matched.num_rows > 0:
+ update_table = self._build_predicate_update_table(
+ matched[SpecialFields.ROW_ID.name],
+ assignments,
+ matched.num_rows,
+ )
+ updater.update_columns(
+ update_table, list(assignments.keys())
+ )
+ else:
+ for split in self._predicate_update_file_groups(splits):
+ matched = table_read.to_arrow([split], parallelism=1)
+ if matched.num_rows == 0:
+ continue
+ update_table = self._build_predicate_update_table(
+ matched[SpecialFields.ROW_ID.name],
+ assignments,
+ matched.num_rows,
+ )
+ updater.update_columns(
+ update_table, list(assignments.keys())
+ )
+ except Exception:
+ _abort_commit_messages(self.table, updater.commit_messages)
+ raise
+ return updater.commit_messages
+
+ @staticmethod
+ def _predicate_update_file_groups(splits):
+ for split in splits:
+ data_split = (
+ split.data_split()
+ if isinstance(split, IndexedSplit) else split
+ )
+ deletion_files = data_split.data_deletion_files
+ deletion_by_name = (
+ {
+ file.file_name: deletion
+ for file, deletion in zip(
+ data_split.files, deletion_files
+ )
+ }
+ if deletion_files is not None else None
+ )
+ groups = DataEvolutionSplitGenerator._split_by_row_id(
+ data_split.files
+ )
+ for files in groups:
+ group = DataSplit(
+ files=files,
+ partition=data_split.partition,
+ bucket=data_split.bucket,
+ raw_convertible=len(files) == 1,
+ data_deletion_files=(
+ [deletion_by_name[file.file_name] for file in files]
+ if deletion_by_name is not None else None
+ ),
+ snapshot_id=data_split.snapshot_id,
+ )
+ if isinstance(split, IndexedSplit):
+ group_ranges = Range.sort_and_merge_overlap(
+ [file.row_id_range() for file in files], True, True
+ )
+ row_ranges = Range.and_(
+ split.row_ranges(), group_ranges
+ )
+ if not row_ranges:
+ continue
+ group = IndexedSplit(group, row_ranges)
+ yield group
def _matched_update_scan_table(self):
snapshot_manager = self.table.snapshot_manager()
