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 d21cef37c3 [python] Support callable assignments in predicate updates 
(#9239)
d21cef37c3 is described below

commit d21cef37c3b9e1964b039e45619f0b327ec04ea2
Author: XiaoHongbo <[email protected]>
AuthorDate: Sun Aug 16 16:01:10 2026 +0800

    [python] Support callable assignments in predicate updates (#9239)
---
 docs/docs/pypaimon/data-evolution.md               |  16 +-
 paimon-python/pypaimon/tests/table_update_test.py  | 287 ++++++++++++++++++++-
 paimon-python/pypaimon/write/table_update.py       |  96 +++++--
 .../pypaimon/write/table_update_by_row_id.py       |   7 +-
 4 files changed, 365 insertions(+), 41 deletions(-)

diff --git a/docs/docs/pypaimon/data-evolution.md 
b/docs/docs/pypaimon/data-evolution.md
index 279e942357..c8cce6df9f 100644
--- a/docs/docs/pypaimon/data-evolution.md
+++ b/docs/docs/pypaimon/data-evolution.md
@@ -96,8 +96,10 @@ table_commit.close()
 ## Update Columns By Predicate
 
 You can use `update_by_predicate` for SQL-like `UPDATE ... SET ... WHERE ...`
-operations. The `Predicate` identifies rows to update, and the assignment map
-contains literal values for updated columns.
+operations. Assignments may be literals or callables. Callables require 
explicit
+`read_columns` and may run in multiple bounded batches. They must be 
deterministic,
+side-effect-free, and row-local, and return one Arrow value per input row. 
Inputs
+are read from the same pinned snapshot used to plan the update.
 When global indexes are available, `update_by_predicate` discovers matching
 `_ROW_ID` values with `scalar-index.search-mode=full` on the configured
 point-in-time scan snapshot or, if none is configured, the latest snapshot.
@@ -133,11 +135,15 @@ commit.commit(write.prepare_commit())
 write.close()
 commit.close()
 
-# UPDATE users_update SET age = 99 WHERE id IN (1, 3)
+# UPDATE users_update SET age = age + 1 WHERE id IN (1, 3)
 write_builder = table.new_batch_write_builder()
 table_update = write_builder.new_update()
 predicate = table_update.new_predicate_builder().is_in('id', [1, 3])
-messages = table_update.update_by_predicate(predicate, {'age': 99})
+messages = table_update.update_by_predicate(
+    predicate,
+    {'age': lambda rows: pa.compute.add(rows['age'], 1)},
+    read_columns=['age'],
+)
 
 commit = write_builder.new_commit()
 commit.commit(messages)
@@ -624,7 +630,7 @@ The API mapping is:
 | --- | --- |
 | `write.prepare_commit()` | `write.prepare_commit(commit_identifier)` |
 | `update.update_by_arrow_with_row_id(table)` | 
`update.update_by_arrow_with_row_id(table, commit_identifier)` |
-| `update.update_by_predicate(predicate, assignments)` | 
`update.update_by_predicate(predicate, assignments, commit_identifier)` |
+| `update.update_by_predicate(predicate, assignments, read_columns=...)` | 
`update.update_by_predicate(predicate, assignments, commit_identifier, 
read_columns=...)` |
 | `update.delete_by_predicate(predicate)` | 
`update.delete_by_predicate(predicate, commit_identifier)` |
 | `update.delete_by_row_id(row_ids)` | `update.delete_by_row_id(row_ids, 
commit_identifier)` |
 | `update.upsert_by_arrow_with_key(table, keys)` | 
`update.upsert_by_arrow_with_key(table, keys, commit_identifier)` |
diff --git a/paimon-python/pypaimon/tests/table_update_test.py 
b/paimon-python/pypaimon/tests/table_update_test.py
index 0d54d78987..ac641b65dd 100644
--- a/paimon-python/pypaimon/tests/table_update_test.py
+++ b/paimon-python/pypaimon/tests/table_update_test.py
@@ -25,6 +25,7 @@ from unittest import mock
 import pyarrow as pa
 import pytest
 
+from pypaimon.read.read_builder import ReadBuilder
 from pypaimon.tests.data_evolution_test_helpers import (
     BatchModeMixin,
     DataEvolutionTestBase,
@@ -54,7 +55,8 @@ class _TableUpdateTestBase(DataEvolutionTestBase):
         raise NotImplementedError
 
     def _apply_update_by_predicate(
-            self, table_update, predicate, assignments, cid):
+            self, table_update, predicate, assignments, cid,
+            read_columns=None):
         raise NotImplementedError
 
     def _apply_delete_by_predicate(self, table_update, predicate, cid):
@@ -138,7 +140,8 @@ class _TableUpdateTestBase(DataEvolutionTestBase):
         tc.close()
         return msgs
 
-    def _do_update_by_predicate(self, table, predicate, assignments):
+    def _do_update_by_predicate(
+            self, table, predicate, assignments, read_columns=None):
         wb = self._make_write_builder(table)
         tu = wb.new_update()
         cid = self._next_commit_id()
@@ -147,6 +150,7 @@ class _TableUpdateTestBase(DataEvolutionTestBase):
             predicate,
             assignments,
             cid,
+            read_columns,
         )
         tc = wb.new_commit()
         self._apply_commit(tc, msgs, cid)
@@ -304,6 +308,29 @@ class _TableUpdateTestBase(DataEvolutionTestBase):
         self.assertEqual(2, calls)
         self.assertEqual(before_files, self._list_table_files(table))
 
+    def test_callable_update_aborts_groups_after_later_failure(self):
+        table = self._create_seeded_table()
+        before_files = self._list_table_files(table)
+        calls = 0
+
+        def fail_second_group(rows):
+            nonlocal calls
+            calls += 1
+            if calls == 2:
+                raise RuntimeError("second callback failed")
+            return pa.compute.add(rows['age'], 1)
+
+        with self.assertRaisesRegex(RuntimeError, "second callback failed"):
+            self._do_update_by_predicate(
+                table,
+                None,
+                {'age': fail_second_group},
+                read_columns=['age'],
+            )
+
+        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(
@@ -338,6 +365,30 @@ class _TableUpdateTestBase(DataEvolutionTestBase):
             self._read_all(table)['age'].to_pylist(),
         )
 
+    def test_literal_predicate_update_projects_only_row_id(self):
+        table = self._create_seeded_table()
+        pb = table.new_read_builder().new_predicate_builder()
+        projections = []
+        with_projection = ReadBuilder.with_projection
+
+        def capture_projection(builder, projection):
+            projections.append(projection)
+            return with_projection(builder, projection)
+
+        with mock.patch.object(
+                ReadBuilder, 'with_projection', capture_projection):
+            self._do_update_by_predicate(
+                table,
+                pb.greater_or_equal('age', 35),
+                {'city': 'Updated'},
+            )
+
+        self.assertEqual([['_ROW_ID']], projections)
+        self.assertEqual(
+            ['NYC', 'LA', 'Updated', 'Updated', 'Updated'],
+            self._read_all(table)['city'].to_pylist(),
+        )
+
     def test_update_by_predicate_accepts_array_chunked_and_scalar_values(self):
         table = self._create_seeded_table()
         pb = table.new_read_builder().new_predicate_builder()
@@ -375,12 +426,43 @@ class _TableUpdateTestBase(DataEvolutionTestBase):
             rows,
         )
 
+    def test_update_by_predicate_accepts_callable_assignments(self):
+        table = self._create_seeded_table()
+        pb = table.new_read_builder().new_predicate_builder()
+
+        def increment_age(rows):
+            self.assertEqual(
+                ['age', 'city', '_ROW_ID'], rows.column_names
+            )
+            return pa.compute.add(rows['age'], 1)
+
+        self._do_update_by_predicate(
+            table,
+            pb.greater_or_equal('age', 35),
+            {
+                'age': increment_age,
+                'city': lambda rows: pa.compute.utf8_upper(rows['city']),
+            },
+            read_columns=['age', 'city'],
+        )
+
+        result = self._read_all(table).sort_by('id')
+        self.assertEqual([25, 30, 36, 41, 46], result['age'].to_pylist())
+        self.assertEqual(
+            ['NYC', 'LA', 'CHICAGO', 'HOUSTON', 'PHOENIX'],
+            result['city'].to_pylist(),
+        )
+
     def test_update_by_predicate_updates_all_rows_when_predicate_is_none(self):
         table = self._create_seeded_table()
         self._do_update_by_predicate(
             table,
             None,
-            {'age': pa.scalar(7, type=pa.int64()), 'city': None},
+            {
+                'age': pa.scalar(7, type=pa.int64()),
+                'city': None,
+                'name': 'UPDATED',
+            },
         )
 
         result = self._read_all(table)
@@ -388,10 +470,105 @@ class _TableUpdateTestBase(DataEvolutionTestBase):
         self.assertEqual([None, None, None, None, None],
                          result['city'].to_pylist())
         self.assertEqual(
-            ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
+            ['UPDATED'] * 5,
             result['name'].to_pylist(),
         )
 
+    def test_update_by_predicate_streams_callable_by_file_group(self):
+        table = self._create_seeded_table()
+        group_sizes = []
+
+        def increment_age(rows):
+            group_sizes.append(rows.num_rows)
+            return pa.compute.add(rows['age'], 1)
+
+        self._do_update_by_predicate(
+            table,
+            None,
+            {'age': increment_age},
+            read_columns=['age'],
+        )
+
+        self.assertEqual([2, 3], group_sizes)
+        self.assertEqual(
+            [26, 31, 36, 41, 46],
+            self._read_all(table)['age'].to_pylist(),
+        )
+
+    def test_update_by_predicate_requires_callable_read_columns(self):
+        table = self._create_seeded_table()
+        pb = table.new_read_builder().new_predicate_builder()
+        with self.assertRaisesRegex(
+                ValueError, "Callable assignments require read_columns"):
+            self._do_update_by_predicate(
+                table,
+                pb.equal('id', 1),
+                {'age': lambda rows: rows['age']},
+            )
+
+    def test_update_by_predicate_rejects_unused_read_columns(self):
+        table = self._create_seeded_table()
+        pb = table.new_read_builder().new_predicate_builder()
+        with self.assertRaisesRegex(
+                ValueError, "read_columns requires a callable assignment"):
+            self._do_update_by_predicate(
+                table,
+                pb.equal('id', 1),
+                {'age': 26},
+                read_columns=['age'],
+            )
+
+    def test_update_by_predicate_rejects_callable_with_array_assignment(self):
+        table = self._create_seeded_table()
+        with self.assertRaisesRegex(
+                ValueError, "cannot be combined with Arrow array"):
+            self._do_update_by_predicate(
+                table,
+                None,
+                {
+                    'age': lambda rows: rows['age'],
+                    'city': pa.array(['A', 'B', 'C', 'D', 'E']),
+                },
+                read_columns=['age'],
+            )
+
+    def test_no_match_does_not_invoke_callable(self):
+        table = self._create_seeded_table()
+        pb = table.new_read_builder().new_predicate_builder()
+
+        def unexpected(_rows):
+            self.fail("Callable should not run without matched rows.")
+
+        messages = self._do_update_by_predicate(
+            table,
+            pb.greater_than('age', 100),
+            {'age': unexpected},
+            read_columns=['age'],
+        )
+
+        self.assertEqual([], messages)
+
+    def test_callable_assignment_rejects_python_list(self):
+        table_schema = pa.schema([
+            ('id', pa.int32()),
+            ('values', pa.list_(pa.int32())),
+        ])
+        table = self._create_table(pa_schema=table_schema)
+        self._write_arrow(table, pa.Table.from_pydict({
+            'id': [1, 2],
+            'values': [[0], [0]],
+        }, schema=table_schema))
+        pb = table.new_read_builder().new_predicate_builder()
+
+        with self.assertRaisesRegex(
+                ValueError, "must return a pyarrow.Array"):
+            self._do_update_by_predicate(
+                table,
+                pb.greater_or_equal('id', 1),
+                {'values': lambda _rows: [1, 2]},
+                read_columns=['id'],
+            )
+
     def 
test_update_by_predicate_rejects_assignment_array_length_mismatch(self):
         table = self._create_seeded_table()
         pb = table.new_read_builder().new_predicate_builder()
@@ -1354,8 +1531,11 @@ class _BatchModeMixin(BatchModeMixin):
         return table_update.update_by_arrow_with_row_id(data)
 
     def _apply_update_by_predicate(
-            self, table_update, predicate, assignments, cid):
-        return table_update.update_by_predicate(predicate, assignments)
+            self, table_update, predicate, assignments, cid,
+            read_columns=None):
+        return table_update.update_by_predicate(
+            predicate, assignments, read_columns
+        )
 
     def _apply_delete_by_predicate(self, table_update, predicate, cid):
         return table_update.delete_by_predicate(predicate)
@@ -1369,11 +1549,13 @@ class _StreamModeMixin(StreamModeMixin):
         return table_update.update_by_arrow_with_row_id(data, cid)
 
     def _apply_update_by_predicate(
-            self, table_update, predicate, assignments, cid):
+            self, table_update, predicate, assignments, cid,
+            read_columns=None):
         return table_update.update_by_predicate(
             predicate,
             assignments,
             cid,
+            read_columns,
         )
 
     def _apply_delete_by_predicate(self, table_update, predicate, cid):
@@ -1390,6 +1572,97 @@ class _StreamModeMixin(StreamModeMixin):
 class TableUpdateBatchTest(_BatchModeMixin, _TableUpdateTestBase, 
unittest.TestCase):
     """All shared update tests under batch (``BatchWriteBuilder``) 
semantics."""
 
+    def test_callable_output_preserves_large_offset_chunks(self):
+        from pypaimon.write.table_update import TableUpdate
+        from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
+
+        chunks = [
+            pa.ListArray.from_arrays(
+                pa.array([0, 1_100_000_000], type=pa.int32()),
+                pa.nulls(1_100_000_000),
+            )
+            for _ in range(2)
+        ]
+        value = pa.chunked_array(chunks)
+
+        result = TableUpdate._assignment_to_array(value, value.type, 2)
+
+        self.assertEqual(2, result.num_chunks)
+        self.assertEqual(2_200_000_000, sum(
+            len(chunk.values) for chunk in result.chunks
+        ))
+
+        updater = TableUpdateByRowId.__new__(TableUpdateByRowId)
+        updater.table = mock.Mock(field_names=['payload'])
+        updater.commit_messages = []
+        updates = pa.Table.from_arrays(
+            [pa.array([1, 0], type=pa.int64()), result],
+            names=['_ROW_ID', 'payload'],
+        )
+        with mock.patch.object(
+                updater, '_calculate_first_row_id',
+                side_effect=lambda data: data) as calculate:
+            with mock.patch.object(updater, '_write_by_first_row_id'):
+                updater.update_columns(updates, ['payload'])
+
+        routed = calculate.call_args[0][0]
+        self.assertEqual([1, 0], routed['_ROW_ID'].to_pylist())
+        self.assertEqual(2, routed['payload'].num_chunks)
+
+    def 
test_callable_predicate_update_allows_concurrent_read_column_update(self):
+        table = self._create_seeded_table()
+        pb = table.new_read_builder().new_predicate_builder()
+        wb = self._make_write_builder(table)
+        messages = wb.new_update().update_by_predicate(
+            pb.equal('id', 1),
+            {
+                'name': lambda rows: pa.array([
+                    'age-%d' % value for value in rows['age'].to_pylist()
+                ]),
+            },
+            read_columns=['age'],
+        )
+
+        self._do_update(table, pa.Table.from_pydict({
+            '_ROW_ID': pa.array([0], type=pa.int64()),
+            'age': pa.array([100], type=pa.int32()),
+        }), ['age'])
+
+        commit = wb.new_commit()
+        commit.commit(messages)
+        commit.close()
+
+        result = self._read_all(table).sort_by('id')
+        self.assertEqual(100, result['age'][0].as_py())
+        self.assertEqual('age-25', result['name'][0].as_py())
+
+    def test_callable_predicate_update_conflicts_with_concurrent_update(self):
+        table = self._create_seeded_table()
+        pb = table.new_read_builder().new_predicate_builder()
+        wb = self._make_write_builder(table)
+        update = wb.new_update()
+
+        def increment_age(rows):
+            self.assertEqual([25], rows['age'].to_pylist())
+            self._do_update(table, pa.Table.from_pydict({
+                '_ROW_ID': pa.array([0], type=pa.int64()),
+                'age': pa.array([100], type=pa.int32()),
+            }), ['age'])
+            return pa.compute.add(rows['age'], 1)
+
+        messages = update.update_by_predicate(
+            pb.equal('id', 1),
+            {'age': increment_age},
+            read_columns=['age'],
+        )
+        commit = wb.new_commit()
+        with self.assertRaisesRegex(RuntimeError, "multiple 'MERGE INTO'"):
+            commit.commit(messages)
+        commit.close()
+
+        result = self._read_all(table).sort_by('id')
+        self.assertEqual([100, 30, 35, 40, 45], result['age'].to_pylist())
+
     def test_update_by_row_id_aborts_files_after_prepare_commit_failure(self):
         from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
 
diff --git a/paimon-python/pypaimon/write/table_update.py 
b/paimon-python/pypaimon/write/table_update.py
index 4760a18516..175af0367e 100644
--- a/paimon-python/pypaimon/write/table_update.py
+++ b/paimon-python/pypaimon/write/table_update.py
@@ -230,26 +230,34 @@ class TableUpdate:
             predicate: Optional[Predicate],
             assignments: Mapping[str, Any],
             commit_identifier: int,
+            read_columns: Optional[Sequence[str]] = None,
     ) -> List[CommitMessage]:
         """Shared implementation for SQL-like ``UPDATE ... WHERE ...``.
 
-        ``predicate`` identifies the target rows. ``assignments`` maps target
-        column names to literal values. The method reads matching ``_ROW_ID``
-        values, builds an Arrow update table, then delegates to the existing
-        row-id update path.
+        ``predicate`` identifies the target rows. Assignment values may be
+        literals or callables receiving one matched logical file group as an
+        Arrow table.
         """
+        has_callable = any(callable(value) for value in assignments.values())
         has_array = any(
             isinstance(value, (pa.Array, pa.ChunkedArray))
             for value in assignments.values()
         )
-        self._validate_predicate_update(assignments)
+        read_columns = tuple(read_columns or ())
+        self._validate_predicate_update(
+            assignments, read_columns, has_callable, has_array
+        )
 
         scan_table = self._matched_update_scan_table()
-        read_builder = scan_table.new_read_builder().with_projection(
-            [SpecialFields.ROW_ID.name]
-        )
+        read_builder = scan_table.new_read_builder()
         if predicate is not None:
             read_builder.with_filter(predicate)
+        if has_callable:
+            projection = list(dict.fromkeys(read_columns))
+            projection.append(SpecialFields.ROW_ID.name)
+            read_builder.with_projection(projection)
+        else:
+            read_builder.with_projection([SpecialFields.ROW_ID.name])
 
         plan = read_builder.new_scan().plan_for_write()
         splits = plan.splits()
@@ -267,9 +275,8 @@ class TableUpdate:
                 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,
+                        matched,
                     )
                     updater.update_columns(
                         update_table, list(assignments.keys())
@@ -280,9 +287,8 @@ class TableUpdate:
                     if matched.num_rows == 0:
                         continue
                     update_table = self._build_predicate_update_table(
-                        matched[SpecialFields.ROW_ID.name],
                         assignments,
-                        matched.num_rows,
+                        matched,
                     )
                     updater.update_columns(
                         update_table, list(assignments.keys())
@@ -363,7 +369,13 @@ class TableUpdate:
 
         return self.table.copy(dynamic_options)
 
-    def _validate_predicate_update(self, assignments: Mapping[str, Any]):
+    def _validate_predicate_update(
+            self,
+            assignments: Mapping[str, Any],
+            read_columns: Optional[Sequence[str]],
+            has_callable: bool,
+            has_array: bool,
+    ):
         if not self.table.options.data_evolution_enabled():
             raise ValueError(
                 "update_by_predicate requires "
@@ -376,6 +388,25 @@ class TableUpdate:
             )
         if not assignments:
             raise ValueError("assignments must not be empty.")
+        if read_columns and not has_callable:
+            raise ValueError(
+                "read_columns requires a callable assignment."
+            )
+        if has_callable:
+            if has_array:
+                raise ValueError(
+                    "Callable assignments cannot be combined with Arrow "
+                    "array assignments."
+                )
+            if not read_columns:
+                raise ValueError(
+                    "Callable assignments require read_columns."
+                )
+            for col in read_columns:
+                if col not in self.table.field_names:
+                    raise ValueError(
+                        f"Read column {col} is not in table schema."
+                    )
 
         partition_keys = set(self.table.partition_keys)
         for col in assignments:
@@ -389,19 +420,27 @@ class TableUpdate:
 
     def _build_predicate_update_table(
             self,
-            row_ids,
             assignments: Mapping[str, Any],
-            row_count: int,
+            matched: pa.Table,
     ) -> pa.Table:
         table_schema = PyarrowFieldParser.from_paimon_schema(
             self.table.table_schema.fields
         )
-        arrays = [row_ids]
+        arrays = [matched[SpecialFields.ROW_ID.name]]
         fields = [pa.field(SpecialFields.ROW_ID.name, pa.int64())]
         for col, value in assignments.items():
+            if callable(value):
+                value = value(matched)
+                if not isinstance(value, (pa.Array, pa.ChunkedArray)):
+                    raise ValueError(
+                        f"Callable assignment for {col} must return a "
+                        "pyarrow.Array or pyarrow.ChunkedArray."
+                    )
             target_field = table_schema.field(col)
             arrays.append(
-                self._assignment_to_array(value, target_field.type, row_count)
+                self._assignment_to_array(
+                    value, target_field.type, matched.num_rows
+                )
             )
             fields.append(target_field)
         return pa.Table.from_arrays(arrays, schema=pa.schema(fields))
@@ -410,7 +449,7 @@ class TableUpdate:
     def _assignment_to_array(
             value: Any, data_type: pa.DataType, row_count: int):
         if isinstance(value, pa.ChunkedArray):
-            array = value.combine_chunks()
+            array = value
         elif isinstance(value, pa.Array):
             array = value
         else:
@@ -424,7 +463,13 @@ class TableUpdate:
                 f"{len(array)} != {row_count}."
             )
         if array.type != data_type:
-            array = array.cast(data_type)
+            if isinstance(array, pa.ChunkedArray):
+                array = pa.chunked_array(
+                    [chunk.cast(data_type) for chunk in array.chunks],
+                    type=data_type,
+                )
+            else:
+                array = array.cast(data_type)
         return array
 
     def _delete_by_predicate(
@@ -599,10 +644,14 @@ class BatchTableUpdate(TableUpdate):
             self,
             predicate: Optional[Predicate],
             assignments: Mapping[str, Any],
+            read_columns: Optional[Sequence[str]] = None,
     ) -> List[CommitMessage]:
-        """Update rows matching ``predicate`` with literal assignments."""
+        """Update rows using literal or Arrow callable assignments."""
         return self._update_by_predicate(
-            predicate, assignments, BATCH_COMMIT_IDENTIFIER
+            predicate,
+            assignments,
+            BATCH_COMMIT_IDENTIFIER,
+            read_columns,
         )
 
     def delete_by_predicate(
@@ -674,11 +723,12 @@ class StreamTableUpdate(TableUpdate):
             predicate: Optional[Predicate],
             assignments: Mapping[str, Any],
             commit_identifier: int,
+            read_columns: Optional[Sequence[str]] = None,
     ) -> List[CommitMessage]:
-        """Update rows matching ``predicate`` with literal assignments,
+        """Update rows using literal or Arrow callable assignments,
         tagging the produced commit messages with ``commit_identifier``."""
         return self._update_by_predicate(
-            predicate, assignments, commit_identifier
+            predicate, assignments, commit_identifier, read_columns
         )
 
     def delete_by_predicate(
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 508a74b5d0..0a68c682e2 100644
--- a/paimon-python/pypaimon/write/table_update_by_row_id.py
+++ b/paimon-python/pypaimon/write/table_update_by_row_id.py
@@ -218,12 +218,7 @@ class TableUpdateByRowId:
             if col_name not in self.table.field_names:
                 raise ValueError(f"Column {col_name} not found in table 
schema")
 
-        sort_keys = [(SpecialFields.ROW_ID.name, "ascending")]
-        if hasattr(data, "sort_by"):
-            sorted_data = data.sort_by(sort_keys)
-        else:
-            sorted_data = data.take(pc.sort_indices(data, sort_keys=sort_keys))
-        data_with_first_row_id = self._calculate_first_row_id(sorted_data)
+        data_with_first_row_id = self._calculate_first_row_id(data)
         self._write_by_first_row_id(data_with_first_row_id, column_names)
 
         return self.commit_messages

Reply via email to