discivigour commented on code in PR #9146:
URL: https://github.com/apache/paimon/pull/9146#discussion_r3750158063
##########
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:
Fixed.
--
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]
