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 f5efd61d60 [python] Vectorize selected-key reconstruction for
shared-shredding MAPs (#10143)
f5efd61d60 is described below
commit f5efd61d6036ac0ca8243e34d0093c4781b9e1c1
Author: zhigang <[email protected]>
AuthorDate: Thu Sep 24 10:58:46 2026 +0800
[python] Vectorize selected-key reconstruction for shared-shredding MAPs
(#10143)
---
.../pypaimon/data/map_shared_shredding.py | 259 ++++++++++++++-------
.../format_pyarrow_shared_shredding_map_test.py | 195 +++++++++++++++-
2 files changed, 372 insertions(+), 82 deletions(-)
diff --git a/paimon-python/pypaimon/data/map_shared_shredding.py
b/paimon-python/pypaimon/data/map_shared_shredding.py
index 1fe0a61627..42994543fd 100644
--- a/paimon-python/pypaimon/data/map_shared_shredding.py
+++ b/paimon-python/pypaimon/data/map_shared_shredding.py
@@ -19,7 +19,7 @@
import json
import struct
from copy import copy
-from typing import Dict, List
+from typing import Dict, List, Optional, Sequence, Set, Tuple, Union, cast
import pyarrow as pa
import pyarrow.compute as pc
@@ -39,6 +39,11 @@ _OVERFLOW = "__overflow"
_PHYSICAL_COLUMN_PREFIX = "__col_"
_SELECTED_KEYS_PREFIX = "__PAIMON_MAP_SELECTED_KEYS:"
_SELECTED_KEYS_DELIMITER = ";"
+_VECTORIZED_MIN_ROWS = 1024
+_VECTORIZED_MAX_CANDIDATE_PAIRS = 16
+
+_SelectedKeyMetadata = Tuple[Dict[int, str], Dict[int, List[int]], Set[int],
int]
+_FieldMappings = List[Optional[List[Optional[int]]]]
def is_shared_shredding(field: pa.Field) -> bool:
@@ -379,7 +384,7 @@ def assemble_shared_shredding_selected_keys(
column: pa.StructArray,
selected_keys: List[str],
value_type: pa.DataType,
- metadata) -> pa.StructArray:
+ metadata: _SelectedKeyMetadata) -> pa.StructArray:
"""Materialize selected MAP values from a pruned physical struct."""
if not pa.types.is_struct(column.type):
raise TypeError("Shared-shredding MAP must be stored as a struct")
@@ -389,105 +394,196 @@ 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))
+ 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()
-
- 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, ())
+ for physical_index in candidate_columns:
+ if not 0 <= physical_index < num_columns:
+ raise ValueError(
+ "Shared-shredding physical column {} is out of range for
{} columns".format(
+ physical_index, num_columns))
+ 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)
+ return indices
+
+
def assemble_normal_map_selected_keys(
column: pa.MapArray,
selected_keys: List[str],
@@ -615,7 +711,7 @@ def _append_entry(keys, entry_sources, entry_positions,
selected_indices,
selected_indices[source].append(source_index)
-def _normalized_offsets(column):
+def _normalized_offsets(column) -> Tuple[List[int], int, int]:
offsets_array = getattr(column, "offsets", None)
if offsets_array is None:
offsets_array = pa.Array.from_buffers(
@@ -624,7 +720,8 @@ def _normalized_offsets(column):
[None, column.buffers()[1]],
offset=column.offset,
)
- offsets = offsets_array.to_pylist()
+ # Arrow list offsets are non-null integers, even for null parent rows.
+ offsets = cast(List[int], offsets_array.to_pylist())
start = offsets[0]
normalized = [value - start for value in offsets]
return normalized, start, offsets[-1]
diff --git
a/paimon-python/pypaimon/tests/format_pyarrow_shared_shredding_map_test.py
b/paimon-python/pypaimon/tests/format_pyarrow_shared_shredding_map_test.py
index 4fd2922032..9bddd6d7ec 100644
--- a/paimon-python/pypaimon/tests/format_pyarrow_shared_shredding_map_test.py
+++ b/paimon-python/pypaimon/tests/format_pyarrow_shared_shredding_map_test.py
@@ -22,6 +22,7 @@ import shutil
import struct
import tempfile
import unittest
+from typing import Dict, List, Optional
from unittest import mock
import pyarrow as pa
@@ -29,7 +30,11 @@ import pyarrow.fs as pafs
import pyarrow.orc as orc
import pyarrow.parquet as pq
-from pypaimon.data.map_shared_shredding import map_selected_keys_field
+from pypaimon.data import map_shared_shredding
+from pypaimon.data.map_shared_shredding import (
+ assemble_shared_shredding_selected_keys,
+ map_selected_keys_field,
+)
from pypaimon.read.reader.format_pyarrow_reader import FormatPyArrowReader
from pypaimon.schema.data_types import (
ArrayType,
@@ -73,6 +78,194 @@ def _metadata(compression):
class SharedShreddingMapReaderTest(unittest.TestCase):
+ def test_selected_keys_preserve_sliced_mapping_and_overflow(self):
+ # Exercise both sides of the batch/candidate limits with the same
+ # sliced data, including null parents, moving keys and duplicate
overflow.
+ # Keep explicit boundaries to catch accidental changes to the policy.
+ cases = [
+ (pa.string(), 6, 2, False),
+ (pa.string(), 1023, 8, False),
+ (pa.string(), 1024, 8, True),
+ (pa.string(), 1025, 8, True),
+ (pa.string(), 1024, 9, False),
+ (pa.struct([("v", pa.int64())]), 1024, 8, False),
+ (pa.int64(), 1024, 8, True),
+ (pa.float64(), 1024, 8, True),
+ (pa.bool_(), 1024, 8, True),
+ (pa.binary(), 1024, 8, True),
+ ]
+ for value_type, rows, width, vectorized in cases:
+ with self.subTest(value_type=value_type, rows=rows, width=width):
+
+ def value(number: int):
+ if pa.types.is_string(value_type):
+ return str(number)
+ if pa.types.is_binary(value_type):
+ return number.to_bytes(2, "little")
+ if pa.types.is_integer(value_type):
+ return number
+ if pa.types.is_floating(value_type):
+ return number / 8.0
+ if pa.types.is_boolean(value_type):
+ return bool(number % 2)
+ return {"v": number}
+
+ mappings = [[0, 1], [], [-1, -1], [0, 1], [-1, -1], [1, 0]]
+ mappings = [m + [-1] * (width - 2) if m else [] for m in
mappings]
+ repeats = (rows + 5) // 6
+ # Prefix is excluded by the slice but shifts both child
offsets.
+ mapping = [[0, 1] + [-1] * (width - 2)] + mappings * repeats
+ first = [value(99)] + [
+ value(1),
+ None,
+ None,
+ None,
+ None,
+ value(15),
+ ] * repeats
+ second = [value(199)] + [
+ value(11),
+ None,
+ None,
+ None,
+ None,
+ value(5),
+ ] * repeats
+ overflow = [[(0, value(999))]] + [
+ [],
+ [],
+ [(0, value(2)), (0, value(3)), (1, value(12))],
+ [(0, value(4)), (1, value(14))],
+ [],
+ None,
+ ] * repeats
+ physical = pa.StructArray.from_arrays(
+ [
+ pa.array(mapping, type=pa.list_(pa.int32())),
+ pa.array(first, type=value_type),
+ pa.array(second, type=value_type),
+ ]
+ + [
+ pa.nulls(len(mapping), type=value_type)
+ for _ in range(width - 2)
+ ]
+ + [pa.array(overflow, type=pa.map_(pa.int32(),
value_type))],
+ names=["__field_mapping"]
+ + ["__col_%d" % i for i in range(width)]
+ + ["__overflow"],
+ mask=pa.array(
+ [False] + [False, True, False, False, False, False] *
repeats
+ ),
+ )
+ expected = [
+ {"k": value(1), "other": value(11), "missing": None},
+ None,
+ {"k": value(2), "other": value(12), "missing": None},
+ {"k": None, "other": None, "missing": None},
+ {"k": None, "other": None, "missing": None},
+ {"k": value(5), "other": value(15), "missing": None},
+ ] * repeats
+ with mock.patch.object(
+ map_shared_shredding.pc,
+ "list_value_length",
+ wraps=map_shared_shredding.pc.list_value_length,
+ ) as vectorized_lengths:
+ selected = assemble_shared_shredding_selected_keys(
+ physical.slice(1, rows),
+ ["k", "other", "missing"],
+ value_type,
+ (
+ {0: "k", 1: "other"},
+ {0: list(range(width)), 1: list(range(min(width,
8)))},
+ {0, 1},
+ width,
+ ),
+ )
+ self.assertEqual(vectorized_lengths.call_count,
int(vectorized))
+ expected_array = pa.array(expected[:rows], type=selected.type)
+ assert isinstance(expected_array, pa.StructArray)
+ self.assertTrue(selected.equals(expected_array))
+
+ def test_selected_keys_mapping_boundaries(self):
+ for rows in (1023, 1024):
+ for case in (
+ "tail_null", "null_id", "short_mapping", "unknown_key",
+ "missing_column", "fixed_size_list", "large_list",
+ ):
+ with self.subTest(rows=rows, case=case):
+ mappings: List[List[Optional[int]]] = [[0, 1] for _ in
range(rows + 1)]
+ mask = [False] * (rows + 1)
+ expected: List[Optional[Dict[str, Optional[str]]]] = [
+ {"k": str(i)} for i in range(1, rows + 1)
+ ]
+ selected_keys = ["k"]
+ mapping_type = pa.list_(pa.int32())
+ if case == "tail_null":
+ mappings[-1] = []
+ mask[-1] = True
+ expected[-1] = None
+ elif case == "null_id":
+ # A later null candidate must not erase an earlier
match.
+ mappings[7] = [0, None]
+ elif case in ("short_mapping", "unknown_key"):
+ mappings[7] = []
+ if case == "unknown_key":
+ selected_keys = ["missing"]
+ expected = [{"missing": None}] * rows
+ elif case == "missing_column":
+ mappings[7] = [1, 0]
+ elif case == "fixed_size_list":
+ mapping_type = pa.list_(pa.int32(), 2)
+ elif case == "large_list":
+ mapping_type = pa.large_list(pa.int32())
+ arrays = [
+ pa.array(mappings, type=mapping_type),
+ pa.array([str(i) for i in range(rows + 1)]),
+ ]
+ names = ["__field_mapping", "__col_0"]
+ if case != "missing_column":
+ arrays.append(pa.nulls(rows + 1, type=pa.string()))
+ names.append("__col_1")
+ physical = pa.StructArray.from_arrays(
+ arrays, names=names, mask=pa.array(mask)
+ ).slice(1)
+ metadata = ({0: "k", 1: "other"}, {0: [0, 1]}, set(), 2)
+ if case in ("short_mapping", "missing_column"):
+ message = (
+ "field mapping length must equal 2"
+ if case == "short_mapping"
+ else "Missing shared-shredding physical column 1"
+ )
+ with self.assertRaisesRegex(ValueError, message):
+ assemble_shared_shredding_selected_keys(
+ physical, selected_keys, pa.string(), metadata
+ )
+ else:
+ selected = assemble_shared_shredding_selected_keys(
+ physical, selected_keys, pa.string(), metadata
+ )
+ expected_array = pa.array(expected, type=selected.type)
+ assert isinstance(expected_array, pa.StructArray)
+ self.assertTrue(selected.equals(expected_array))
+
+ def test_selected_keys_reject_out_of_range_candidates(self):
+ for rows in (1023, 1024):
+ for physical_index in (-1, 1):
+ with self.subTest(rows=rows, physical_index=physical_index):
+ # Keep adjacent rows in the backing array so an invalid
+ # candidate can cross a row boundary without take failing.
+ physical = pa.StructArray.from_arrays(
+ [
+ pa.array([[0]] * (rows + 2),
type=pa.list_(pa.int32())),
+ pa.array(["value"] * (rows + 2)),
+ ],
+ names=["__field_mapping",
"__col_{}".format(physical_index)],
+ ).slice(1, rows)
+ with self.assertRaisesRegex(ValueError, "physical column
.* out of range"):
+ assemble_shared_shredding_selected_keys(
+ physical, ["k"], pa.string(),
+ ({0: "k"}, {0: [physical_index]}, set(), 1),
+ )
def setUp(self):
self.tmp = tempfile.mkdtemp()
