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 17bc76450f [python] Fix row-id updates for Arrow columns over 2 GiB
(#9146)
17bc76450f is described below
commit 17bc76450ff3d36990ff258c15d311cdff9ac58e
Author: umi <[email protected]>
AuthorDate: Mon Aug 10 23:06:40 2026 +0800
[python] Fix row-id updates for Arrow columns over 2 GiB (#9146)
---
.../tests/table_update_by_row_id_chunked_test.py | 206 ++++++++++++++++++++
.../pypaimon/write/table_update_by_row_id.py | 210 ++++++++++++++++++---
2 files changed, 393 insertions(+), 23 deletions(-)
diff --git
a/paimon-python/pypaimon/tests/table_update_by_row_id_chunked_test.py
b/paimon-python/pypaimon/tests/table_update_by_row_id_chunked_test.py
new file mode 100644
index 0000000000..c0902c5421
--- /dev/null
+++ b/paimon-python/pypaimon/tests/table_update_by_row_id_chunked_test.py
@@ -0,0 +1,206 @@
+################################################################################
+# 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.
+################################################################################
+
+import unittest
+from unittest import mock
+
+import pyarrow as pa
+import pyarrow.compute as pc
+
+from pypaimon.table.special_fields import SpecialFields
+from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
+
+
+class TableUpdateByRowIdChunkedTest(unittest.TestCase):
+
+ @staticmethod
+ def _updater():
+ updater = TableUpdateByRowId.__new__(TableUpdateByRowId)
+ updater._is_blob_column = lambda _: False
+ return updater
+
+ def test_variant_update_preserves_chunks_and_non_nullable_fields(self):
+ variant_type = pa.struct([
+ pa.field("value", pa.binary(), nullable=False),
+ pa.field("metadata", pa.binary(), nullable=False),
+ ])
+ original_payload = pa.chunked_array([
+ pa.array([
+ {"value": b"v0", "metadata": b"m0"},
+ {"value": b"v1", "metadata": b"m1"},
+ ], type=variant_type),
+ pa.array([
+ {"value": b"v2", "metadata": b"m2"},
+ {"value": b"v3", "metadata": b"m3"},
+ ], type=variant_type),
+ ])
+ update_payload = pa.chunked_array([
+ pa.array([
+ {"value": b"updated-1", "metadata": b"new-1"},
+ ], type=variant_type),
+ pa.array([
+ {"value": b"updated-3", "metadata": b"new-3"},
+ ], type=variant_type),
+ ])
+ original = pa.table({"payload": original_payload})
+ updates = pa.table({
+ SpecialFields.ROW_ID.name: pa.array([101, 103], type=pa.int64()),
+ "payload": update_payload,
+ })
+
+ merged, _ = self._updater()._merge_update_with_original(
+ original, updates, ["payload"], first_row_id=100)
+
+ self.assertEqual(merged.schema.field("payload").type, variant_type)
+ self.assertEqual(merged["payload"].num_chunks, 2)
+ self.assertEqual(merged["payload"].to_pylist(), [
+ {"value": b"v0", "metadata": b"m0"},
+ {"value": b"updated-1", "metadata": b"new-1"},
+ {"value": b"v2", "metadata": b"m2"},
+ {"value": b"updated-3", "metadata": b"new-3"},
+ ])
+
+ def test_update_positions_are_sorted_once_for_all_columns(self):
+ original = pa.table({
+ "left": pa.chunked_array([[1], [2]]),
+ "right": pa.chunked_array([[10], [20]]),
+ })
+ updates = pa.table({
+ SpecialFields.ROW_ID.name: pa.array([1], type=pa.int64()),
+ "left": pa.array([3]),
+ "right": pa.array([30]),
+ })
+
+ with mock.patch("builtins.sorted", wraps=sorted) as sorted_mock:
+ merged, _ = self._updater()._merge_update_with_original(
+ original, updates, ["left", "right"], first_row_id=0)
+
+ self.assertEqual(sorted_mock.call_count, 1)
+ self.assertEqual(merged["left"].to_pylist(), [1, 3])
+ self.assertEqual(merged["right"].to_pylist(), [10, 30])
+
+ def test_update_position_outside_column_range_raises(self):
+ original = pa.table({"payload": pa.array([1, 2])})
+ updates = pa.table({
+ SpecialFields.ROW_ID.name: pa.array([2], type=pa.int64()),
+ "payload": pa.array([3]),
+ })
+
+ with self.assertRaisesRegex(IndexError, "outside column range"):
+ self._updater()._merge_update_with_original(
+ original, updates, ["payload"], first_row_id=0)
+
+ def test_total_offsets_over_int32_remain_in_separate_chunks(self):
+ child_length = 1_100_000_000
+ large_chunk = self._list_chunk([0, child_length])
+ original = pa.table({
+ "payload": pa.chunked_array([large_chunk, large_chunk]),
+ })
+ updates = pa.table({
+ SpecialFields.ROW_ID.name: pa.array([1], type=pa.int64()),
+ "payload": pa.chunked_array([large_chunk]),
+ })
+
+ merged, _ = self._updater()._merge_update_with_original(
+ original, updates, ["payload"], first_row_id=0)
+
+ payload = merged["payload"]
+ self.assertEqual(payload.num_chunks, 2)
+ self.assertEqual(
+ sum(len(chunk.values) for chunk in payload.chunks),
+ 2 * child_length,
+ )
+
+ def test_fallback_splits_before_temporary_concat_overflows(self):
+ original_value_length = 600_000_000
+ replacement_value_length = 1_100_000_000
+ original_chunk = self._list_chunk([
+ 0,
+ original_value_length,
+ 2 * original_value_length,
+ ])
+ replacement_chunk = self._list_chunk([0, replacement_value_length])
+ original = pa.table({
+ "payload": pa.chunked_array([original_chunk]),
+ })
+ updates = pa.table({
+ SpecialFields.ROW_ID.name: pa.array([1], type=pa.int64()),
+ "payload": pa.chunked_array([replacement_chunk]),
+ })
+
+ with mock.patch.object(
+ TableUpdateByRowId,
+ "_chunk_offsets",
+ wraps=TableUpdateByRowId._chunk_offsets,
+ ) as chunk_offsets:
+ merged, _ = self._updater()._merge_update_with_original(
+ original, updates, ["payload"], first_row_id=0)
+
+ payload = merged["payload"]
+ self.assertEqual(chunk_offsets.call_count, 1)
+ self.assertEqual(payload.num_chunks, 2)
+ self.assertEqual(
+ [len(chunk[0].values) for chunk in payload.chunks],
+ [original_value_length, replacement_value_length],
+ )
+
+ def test_replace_with_mask_capacity_error_splits_chunk(self):
+ original = pa.table({
+ "payload": pa.array(
+ [b"a", b"b", b"c", b"d"], type=pa.binary()),
+ })
+ updates = pa.table({
+ SpecialFields.ROW_ID.name:
+ pa.array([1, 3], type=pa.int64()),
+ "payload": pa.array([b"B", b"D"], type=pa.binary()),
+ })
+ real_replace_with_mask = pc.replace_with_mask
+ attempted_lengths = []
+
+ def replace_with_capacity_limit(values, mask, replacements):
+ attempted_lengths.append(len(values))
+ if len(values) > 2:
+ raise pa.lib.ArrowCapacityError(
+ "array cannot contain more than 2147483647 bytes")
+ return real_replace_with_mask(values, mask, replacements)
+
+ with mock.patch.object(
+ pc,
+ "replace_with_mask",
+ side_effect=replace_with_capacity_limit,
+ ):
+ merged, _ = self._updater()._merge_update_with_original(
+ original, updates, ["payload"], first_row_id=0)
+
+ self.assertEqual(attempted_lengths, [4, 2, 2])
+ self.assertEqual(merged["payload"].num_chunks, 2)
+ self.assertEqual(
+ merged["payload"].to_pylist(),
+ [b"a", b"B", b"c", b"D"],
+ )
+
+ @staticmethod
+ def _list_chunk(offsets):
+ return pa.ListArray.from_arrays(
+ pa.array(offsets, type=pa.int32()),
+ pa.nulls(offsets[-1]),
+ )
+
+
+if __name__ == '__main__':
+ unittest.main()
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 c74fbc54fa..508a74b5d0 100644
--- a/paimon-python/pypaimon/write/table_update_by_row_id.py
+++ b/paimon-python/pypaimon/write/table_update_by_row_id.py
@@ -461,15 +461,11 @@ class TableUpdateByRowId:
pa.scalar(first_row_id, type=pa.int64())
).cast(pa.int64())
- # Build a boolean mask: True at positions that need to be updated
- all_indices = pa.array(range(original_data.num_rows), type=pa.int64())
- mask = pc.is_in(all_indices, value_set=relative_indices)
-
# Build the merged table column by column
merged_columns = {}
blob_columns: Dict[str, List[object]] = {}
update_by_col = {
- col_name: update_data[col_name].combine_chunks()
+ col_name: update_data[col_name]
for col_name in column_names
if col_name in update_data.column_names
}
@@ -477,6 +473,7 @@ class TableUpdateByRowId:
int(relative_index.as_py()): idx
for idx, relative_index in enumerate(relative_indices)
}
+ sorted_updates = None
# Caller (_write_by_first_row_id) only enters this method with a
# non-empty group, so update_positions is non-empty here.
blob_row_count = max(update_positions) + 1
@@ -502,34 +499,201 @@ class TableUpdateByRowId:
]
continue
update_col = update_by_col[col_name]
- original_col = original_data[col_name].combine_chunks()
- if update_col.type != original_col.type:
- update_col = self._coerce_column(
- update_col, original_col.type)
- try:
- merged_columns[col_name] = pc.replace_with_mask(
- original_col, mask, update_col)
- except pa.lib.ArrowNotImplementedError:
- n = original_data.num_rows
- combined = pa.concat_arrays(
- [original_col, update_col])
- offset = len(original_col)
- indices = np.arange(n, dtype=np.int64)
- for orig_pos, upd_idx in update_positions.items():
- indices[orig_pos] = offset + upd_idx
- merged_columns[col_name] = combined.take(
- pa.array(indices))
+ original_col = original_data[col_name]
+ if sorted_updates is None:
+ sorted_updates = sorted(update_positions.items())
+ row_count = len(original_col)
+ for position, _ in sorted_updates:
+ if position < 0 or position >= row_count:
+ raise IndexError(
+ f"Update position {position} is outside column "
+ f"range [0, {row_count})")
+ merged_columns[col_name] = self._merge_chunked_column(
+ original_col, update_col, sorted_updates)
merged_table = pa.table(merged_columns) if merged_columns else None
return merged_table, blob_columns
+ @classmethod
+ def _merge_chunked_column(
+ cls,
+ original_col: pa.ChunkedArray,
+ update_col: pa.ChunkedArray,
+ sorted_updates: List[Tuple[int, int]],
+ ) -> pa.ChunkedArray:
+ """Merge updates without flattening a column into one Arrow Array.
+
+ ``binary``, ``string`` and ``list`` use signed 32-bit offsets, so a
+ valid multi-chunk column can exceed 2 GiB while each individual Array
+ remains below the limit. Keep those chunks independent. If merging an
+ individual chunk still overflows (for example, the struct/list
+ fallback temporarily concatenates original and replacement values),
+ split that row range and retry.
+ """
+ update_chunk_offsets = cls._chunk_offsets(update_col)
+ sorted_updates_idx = 0
+ chunk_start_row = 0
+ merged_chunks: List[pa.Array] = []
+
+ for original_chunk in original_col.chunks:
+ chunk_end = chunk_start_row + len(original_chunk)
+ chunk_updates: List[Tuple[int, int]] = []
+ while (
+ sorted_updates_idx < len(sorted_updates)
+ and sorted_updates[sorted_updates_idx][0] < chunk_end
+ ):
+ position, update_index = sorted_updates[sorted_updates_idx]
+ chunk_updates.append((
+ position - chunk_start_row, update_index))
+ sorted_updates_idx += 1
+
+ merged_chunks.extend(cls._merge_chunk_with_updates(
+ original_chunk,
+ update_col,
+ update_chunk_offsets,
+ chunk_updates,
+ ))
+ chunk_start_row = chunk_end
+
+ return pa.chunked_array(merged_chunks, type=original_col.type)
+
+ @classmethod
+ def _merge_chunk_with_updates(
+ cls,
+ original: pa.Array,
+ update_col: pa.ChunkedArray,
+ update_chunk_offsets: List[int],
+ updates: List[Tuple[int, int]],
+ ) -> List[pa.Array]:
+ if not updates:
+ return [original]
+
+ try:
+ replacements = cls._take_from_chunked_array(
+ update_col,
+ update_chunk_offsets,
+ [update_index for _, update_index in updates],
+ )
+ if replacements.type != original.type:
+ replacements = cls._coerce_column(
+ replacements, original.type)
+
+ if len(updates) == len(original):
+ return [replacements]
+
+ mask_values = np.zeros(len(original), dtype=np.bool_)
+ for position, _ in updates:
+ mask_values[position] = True
+ mask = pa.array(mask_values)
+
+ try:
+ return [pc.replace_with_mask(original, mask, replacements)]
+ except pa.lib.ArrowNotImplementedError:
+ combined = pa.concat_arrays([original, replacements])
+ indices = np.arange(len(original), dtype=np.int64)
+ replacement_offset = len(original)
+ for replacement_index, (position, _) in enumerate(updates):
+ indices[position] = replacement_offset + replacement_index
+ return [combined.take(pa.array(indices))]
+ except (pa.lib.ArrowInvalid,
+ pa.lib.ArrowCapacityError) as error:
+ if not cls._is_offset_overflow(error) or len(original) <= 1:
+ raise
+
+ split_at = len(original) // 2
+ left_updates = [
+ update for update in updates if update[0] < split_at
+ ]
+ right_updates = [
+ (position - split_at, update_index)
+ for position, update_index in updates
+ if position >= split_at
+ ]
+ return (
+ cls._merge_chunk_with_updates(
+ original.slice(0, split_at),
+ update_col,
+ update_chunk_offsets,
+ left_updates,
+ )
+ + cls._merge_chunk_with_updates(
+ original.slice(split_at),
+ update_col,
+ update_chunk_offsets,
+ right_updates,
+ )
+ )
+
+ @staticmethod
+ def _chunk_offsets(column: pa.ChunkedArray) -> List[int]:
+ offsets = [0]
+ for chunk in column.chunks:
+ offsets.append(offsets[-1] + len(chunk))
+ return offsets
+
+ @staticmethod
+ def _take_from_chunked_array(
+ column: pa.ChunkedArray,
+ chunk_offsets: List[int],
+ indices: List[int],
+ ) -> pa.Array:
+ """Take values without asking Arrow to combine unrelated chunks."""
+ if not indices:
+ return pa.array([], type=column.type)
+
+ pieces: List[pa.Array] = []
+ current_chunk_index = None
+ current_local_indices: List[int] = []
+
+ def append_piece(chunk_index: int, local_indices: List[int]):
+ chunk = column.chunk(chunk_index)
+ start = local_indices[0]
+ if all(value == start + offset
+ for offset, value in enumerate(local_indices)):
+ pieces.append(chunk.slice(start, len(local_indices)))
+ else:
+ pieces.append(chunk.take(pa.array(
+ local_indices, type=pa.int64())))
+
+ for index in indices:
+ if index < 0 or index >= len(column):
+ raise IndexError(
+ f"Update index {index} is outside column range "
+ f"[0, {len(column)})")
+ chunk_index = bisect.bisect_right(chunk_offsets, index) - 1
+ local_index = index - chunk_offsets[chunk_index]
+ if current_chunk_index is None:
+ current_chunk_index = chunk_index
+ elif chunk_index != current_chunk_index:
+ append_piece(current_chunk_index, current_local_indices)
+ current_chunk_index = chunk_index
+ current_local_indices = []
+ current_local_indices.append(local_index)
+
+ append_piece(current_chunk_index, current_local_indices)
+ if len(pieces) == 1:
+ return pieces[0]
+ return pa.concat_arrays(pieces)
+
+ @staticmethod
+ def _is_offset_overflow(error: pa.lib.ArrowException) -> bool:
+ if isinstance(error, pa.lib.ArrowCapacityError):
+ return True
+ message = str(error).lower()
+ return (
+ "offset overflow" in message
+ or "too large to convert" in message
+ )
+
@staticmethod
def _coerce_column(col: pa.Array, target_type: pa.DataType) -> pa.Array:
try:
return col.cast(target_type)
+ except pa.lib.ArrowInvalid as error:
+ if TableUpdateByRowId._is_offset_overflow(error):
+ raise
except (pa.lib.ArrowNotImplementedError,
- pa.lib.ArrowInvalid,
pa.lib.ArrowTypeError):
pass
pylist = col.to_pylist()
