Copilot commented on code in PR #10143:
URL: https://github.com/apache/paimon/pull/10143#discussion_r4082213930
##########
paimon-python/pypaimon/data/map_shared_shredding.py:
##########
@@ -377,105 +382,191 @@ def assemble_shared_shredding_selected_keys(
field_names = [field.name for field in column.type]
if not field_names or field_names[0] != _FIELD_MAPPING:
raise ValueError(
- "Shared-shredding physical struct must start with {}".format(
- _FIELD_MAPPING))
+ "Shared-shredding physical struct must start with
{}".format(_FIELD_MAPPING))
physical_columns = {}
- overflow = None
+ overflow: Optional[pa.MapArray] = None
for position, field_name in enumerate(field_names[1:], 1):
+ child = column.field(position)
if field_name == _OVERFLOW:
- overflow = column.field(position)
+ if not isinstance(child, pa.MapArray):
+ raise TypeError("Shared-shredding overflow field must be a
map")
+ overflow = child
elif field_name.startswith(_PHYSICAL_COLUMN_PREFIX):
try:
- physical_columns[int(
- field_name[len(_PHYSICAL_COLUMN_PREFIX):])] = column.field(
- position)
+
physical_columns[int(field_name[len(_PHYSICAL_COLUMN_PREFIX):])] = child
except ValueError:
- raise ValueError(
- "Unexpected shared-shredding physical field: {}".format(
- field_name))
+ raise ValueError("Unexpected shared-shredding physical field:
{}".format(field_name))
else:
- raise ValueError(
- "Unexpected shared-shredding physical field: {}".format(
- field_name))
-
- mapping = column.field(0).to_pylist()
- null_rows = column.is_null().to_pylist()
- overflow_offsets = overflow_keys = overflow_values = overflow_nulls = None
- if overflow is not None:
- overflow_offsets, overflow_start, overflow_end = _normalized_offsets(
- overflow)
- overflow_keys = overflow.keys.slice(
- overflow_start, overflow_end - overflow_start).to_pylist()
- overflow_values = _restore_orc_temporal_values(
- overflow.items.slice(
- overflow_start, overflow_end - overflow_start), value_type)
- overflow_nulls = overflow.is_null().to_pylist()
+ raise ValueError("Unexpected shared-shredding physical field:
{}".format(field_name))
- sources = []
- source_by_column = {}
- for physical_index in sorted(physical_columns):
- source_by_column[physical_index] = len(sources)
- sources.append(_restore_orc_temporal_values(
- physical_columns[physical_index], value_type))
- overflow_source = None
- if overflow_values is not None:
- overflow_source = len(sources)
- sources.append(overflow_values)
+ mapping_column = column.field(0)
+ requested_ids = [id_by_name[key] for key in selected_keys if key in
id_by_name]
+ candidate_count = sum(len(field_to_columns.get(field_id, ())) for field_id
in requested_ids)
Review Comment:
`requested_ids` preserves duplicates if `selected_keys` contains duplicates,
which can inflate `candidate_count` and incorrectly disable the vectorized path
(hurting performance without changing correctness). Consider de-duplicating for
the gating calculation (e.g., compute candidate_count over
`set(requested_ids)`), while still iterating `selected_keys` normally to
preserve output schema/order.
##########
paimon-python/pypaimon/data/map_shared_shredding.py:
##########
@@ -377,105 +382,191 @@ def assemble_shared_shredding_selected_keys(
field_names = [field.name for field in column.type]
if not field_names or field_names[0] != _FIELD_MAPPING:
raise ValueError(
- "Shared-shredding physical struct must start with {}".format(
- _FIELD_MAPPING))
+ "Shared-shredding physical struct must start with
{}".format(_FIELD_MAPPING))
physical_columns = {}
- overflow = None
+ overflow: Optional[pa.MapArray] = None
for position, field_name in enumerate(field_names[1:], 1):
+ child = column.field(position)
if field_name == _OVERFLOW:
- overflow = column.field(position)
+ if not isinstance(child, pa.MapArray):
+ raise TypeError("Shared-shredding overflow field must be a
map")
+ overflow = child
elif field_name.startswith(_PHYSICAL_COLUMN_PREFIX):
try:
- physical_columns[int(
- field_name[len(_PHYSICAL_COLUMN_PREFIX):])] = column.field(
- position)
+
physical_columns[int(field_name[len(_PHYSICAL_COLUMN_PREFIX):])] = child
except ValueError:
- raise ValueError(
- "Unexpected shared-shredding physical field: {}".format(
- field_name))
+ raise ValueError("Unexpected shared-shredding physical field:
{}".format(field_name))
else:
- raise ValueError(
- "Unexpected shared-shredding physical field: {}".format(
- field_name))
-
- mapping = column.field(0).to_pylist()
- null_rows = column.is_null().to_pylist()
- overflow_offsets = overflow_keys = overflow_values = overflow_nulls = None
- if overflow is not None:
- overflow_offsets, overflow_start, overflow_end = _normalized_offsets(
- overflow)
- overflow_keys = overflow.keys.slice(
- overflow_start, overflow_end - overflow_start).to_pylist()
- overflow_values = _restore_orc_temporal_values(
- overflow.items.slice(
- overflow_start, overflow_end - overflow_start), value_type)
- overflow_nulls = overflow.is_null().to_pylist()
+ raise ValueError("Unexpected shared-shredding physical field:
{}".format(field_name))
- sources = []
- source_by_column = {}
- for physical_index in sorted(physical_columns):
- source_by_column[physical_index] = len(sources)
- sources.append(_restore_orc_temporal_values(
- physical_columns[physical_index], value_type))
- overflow_source = None
- if overflow_values is not None:
- overflow_source = len(sources)
- sources.append(overflow_values)
+ mapping_column = column.field(0)
+ requested_ids = [id_by_name[key] for key in selected_keys if key in
id_by_name]
+ candidate_count = sum(len(field_to_columns.get(field_id, ())) for field_id
in requested_ids)
+ mapping_rows: _FieldMappings = []
+ null_rows = []
+ if (isinstance(mapping_column, (pa.ListArray, pa.LargeListArray))
+ and _use_vectorized_selected_keys(value_type, len(column),
candidate_count)):
+ vectorized = True
+ valid_rows = column.is_valid()
+ if requested_ids:
+ _validate_selected_key_mapping(mapping_column, valid_rows,
num_columns)
+ # Sliced offsets still address the original values. Mask null parents
+ # before gathering, since their child lists may be empty.
+ mapping_offsets = pc.if_else(valid_rows, mapping_column.offsets[:-1],
None)
+ mapping_values = mapping_column.values
+ row_indices = pa.array(range(len(column)), type=pa.int64())
+ else:
+ vectorized = False
+ mapping_rows = column.field(0).to_pylist()
+ null_rows = column.is_null().to_pylist()
- source_bases = []
+ # Use one value pool, with a direct physical-column -> starting-offset map.
value_arrays = []
+ column_bases = {}
next_base = 0
- for source in sources:
- source_bases.append(next_base)
- value_arrays.append(source)
- next_base += len(source)
- value_pool = (
- pa.concat_arrays(value_arrays)
- if value_arrays else pa.array([], type=value_type)
- )
+ for physical_index in sorted(physical_columns):
+ values =
_restore_orc_temporal_values(physical_columns[physical_index], value_type)
+ column_bases[physical_index] = next_base
+ value_arrays.append(values)
+ next_base += len(values)
+
+ overflow_base = None
+ overflow_offsets: List[int] = []
+ overflow_keys: List[Optional[int]] = []
+ overflow_nulls: List[Optional[bool]] = []
+ if overflow is not None:
+ # Unlike mapping values, overflow values are sliced into the pool, so
+ # their offsets must be relative to that slice.
+ overflow_offsets, start, end = _normalized_offsets(overflow)
+ overflow_keys = overflow.keys.slice(start, end - start).to_pylist()
+ overflow_nulls = overflow.is_null().to_pylist()
+ overflow_base = next_base
+ value_arrays.append(_restore_orc_temporal_values(
+ overflow.items.slice(start, end - start), value_type))
+ value_pool = (pa.concat_arrays(value_arrays)
+ if value_arrays else pa.array([], type=value_type))
children = []
for key in selected_keys:
field_id = id_by_name.get(key)
- indices = []
- candidate_columns = (
- field_to_columns.get(field_id, ()) if field_id is not None else ())
- for row in range(len(column)):
- selected = None
- if not null_rows[row] and field_id is not None:
- row_mapping = mapping[row]
- if row_mapping is None or len(row_mapping) != num_columns:
- raise ValueError(
- "Shared-shredding field mapping length must equal
{}".format(
- num_columns))
- for physical_index in candidate_columns:
- if row_mapping[physical_index] == field_id:
- source = source_by_column.get(physical_index)
- if source is None:
- raise ValueError(
- "Missing shared-shredding physical column
{}".format(
- physical_index))
- selected = source_bases[source] + row
+ if field_id is None:
+ children.append(pc.take(value_pool, pa.nulls(len(column),
type=pa.int64())))
+ continue
+ candidate_columns = field_to_columns.get(field_id, ())
+ selected_rows: List[Optional[int]]
+ if vectorized:
+ indices = _selected_key_indices_vectorized(
+ mapping_values, mapping_offsets, row_indices,
+ field_id, candidate_columns, column_bases)
+ if field_id not in overflow_set or overflow_base is None:
+ children.append(pc.take(value_pool, indices))
+ continue
+ selected_rows = indices.to_pylist()
+ else:
+ selected_rows = _selected_key_indices(
+ mapping_rows, null_rows, num_columns,
+ field_id, candidate_columns, column_bases)
+
+ if field_id in overflow_set and overflow_base is not None:
+ # Both paths use the same first-match overflow lookup. A physical
+ # match already has an index, even when its value is null.
+ if not null_rows:
+ null_rows = column.is_null().to_pylist()
+ for row, selected in enumerate(selected_rows):
+ if selected is not None or null_rows[row] or
overflow_nulls[row]:
+ continue
+ for item_index in range(overflow_offsets[row],
overflow_offsets[row + 1]):
+ if overflow_keys[item_index] == field_id:
+ selected_rows[row] = overflow_base + item_index
break
- if (selected is None
- and field_id in overflow_set
- and overflow_source is not None
- and not overflow_nulls[row]):
- for item_index in range(
- overflow_offsets[row], overflow_offsets[row + 1]):
- if overflow_keys[item_index] == field_id:
- selected = (
- source_bases[overflow_source] + item_index)
- break
- indices.append(selected)
- children.append(pc.take(
- value_pool, pa.array(indices, type=pa.int64())))
+ children.append(pc.take(value_pool, pa.array(selected_rows,
type=pa.int64())))
fields = [pa.field(key, value_type) for key in selected_keys]
mask = column.is_null() if column.null_count else None
return pa.StructArray.from_arrays(children, fields=fields, mask=mask)
+def _use_vectorized_selected_keys(value_type: pa.DataType, rows: int,
candidate_count: int) -> bool:
+ # Conservative limits from local benchmarks, not universal crossover
points.
+ # Count key/column pairs: shared columns still need a match for each key.
+ if rows < _VECTORIZED_MIN_ROWS or candidate_count >
_VECTORIZED_MAX_CANDIDATE_PAIRS:
+ return False
+ # Other types retain the original path pending performance validation.
+ return (pa.types.is_string(value_type)
+ or pa.types.is_binary(value_type)
+ or pa.types.is_boolean(value_type)
+ or pa.types.is_signed_integer(value_type)
+ or pa.types.is_floating(value_type)
+ or pa.types.is_decimal128(value_type)
+ or pa.types.is_date32(value_type)
+ or pa.types.is_time32(value_type)
+ or pa.types.is_timestamp(value_type))
+
+
+def _validate_selected_key_mapping(
+ mapping: Union[pa.ListArray, pa.LargeListArray],
+ valid_rows: pa.BooleanArray,
+ num_columns: int) -> None:
+ lengths = pc.list_value_length(mapping)
+ invalid_lengths = pc.not_equal(lengths, pa.scalar(num_columns))
+ invalid_rows = pc.and_(valid_rows, pc.fill_null(invalid_lengths,
pa.scalar(True)))
+ if pc.any(invalid_rows).as_py():
+ raise ValueError("Shared-shredding field mapping length must equal
{}".format(num_columns))
+
+
+def _selected_key_indices(
+ mapping_rows: _FieldMappings,
+ null_rows: List[Optional[bool]],
+ num_columns: int,
+ field_id: int,
+ candidate_columns: Sequence[int],
+ column_bases: Dict[int, int]) -> List[Optional[int]]:
+ """Find the first physical match in each row, as an index into the value
pool."""
+ indices: List[Optional[int]] = []
+ for row, row_mapping in enumerate(mapping_rows):
+ selected = None
+ if not null_rows[row]:
+ if row_mapping is None or len(row_mapping) != num_columns:
+ raise ValueError(
+ "Shared-shredding field mapping length must equal
{}".format(num_columns))
+ for physical_index in candidate_columns:
+ if row_mapping[physical_index] == field_id:
+ base = column_bases.get(physical_index)
+ if base is None:
+ raise ValueError(
+ "Missing shared-shredding physical column
{}".format(physical_index))
+ selected = base + row
+ break
+ indices.append(selected)
+ return indices
+
+
+def _selected_key_indices_vectorized(
+ mapping_values: pa.Array,
+ mapping_offsets: pa.Array,
+ row_indices: pa.Array,
+ field_id: int,
+ candidate_columns: Sequence[int],
+ column_bases: Dict[int, int]) -> pa.Array:
+ """Find the same first-match indices using Arrow operations across rows."""
+ indices = pa.nulls(len(row_indices), type=pa.int64())
+ field_id_scalar = pa.scalar(field_id)
+ false_scalar = pa.scalar(False)
+ for physical_index in candidate_columns:
+ mapped_ids = pc.take(mapping_values, pc.add(mapping_offsets,
physical_index))
+ matches = pc.equal(mapped_ids, field_id_scalar)
+ # A null mapping entry is not a match; later candidates cannot replace
+ # an index chosen by an earlier candidate.
+ matches = pc.and_(pc.fill_null(matches, false_scalar),
pc.is_null(indices))
+ base = column_bases.get(physical_index)
+ if base is None:
+ if pc.any(matches).as_py():
+ raise ValueError("Missing shared-shredding physical column
{}".format(physical_index))
+ continue
+ indices = pc.if_else(matches, pc.add(row_indices, base), indices)
Review Comment:
The vectorized gather path assumes every `physical_index` in
`candidate_columns` is within the mapping width (`0 <= physical_index <
num_columns`). If metadata ever contains an out-of-range index,
`pc.take(mapping_values, mapping_offsets + physical_index)` can read the wrong
element (potentially from the next row’s flattened mapping) instead of failing,
causing incorrect key selection. Add an explicit bounds validation for
`candidate_columns` (either in `assemble_shared_shredding_selected_keys` before
calling the vectorized helper, or by passing `num_columns` into
`_selected_key_indices_vectorized` and validating there) and raise a clear
`ValueError` when violated.
--
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]
