JingsongLi commented on code in PR #9146:
URL: https://github.com/apache/paimon/pull/9146#discussion_r3749696010
##########
paimon-python/pypaimon/write/table_update_by_row_id.py:
##########
@@ -493,34 +490,198 @@ def _merge_update_with_original(
]
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 as error:
Review Comment:
Could we also catch `pa.lib.ArrowCapacityError` here? For binary/string
arrays, `replace_with_mask` uses Arrow’s binary builder, whose 32-bit buffer
overflow is surfaced as `ArrowCapacityError` (`array cannot contain more than
... bytes`), not `ArrowInvalid` (including PyArrow 6). Therefore, when a
replacement pushes one chunk over 2 GiB, the exception bypasses the recursive
split even though each half would be representable. Please treat capacity
errors as splittable offset overflows and add a regression test for this path.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
