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 e28c6582c6 [python] Support efficient field selection from wide
VARIANT objects (#10320)
e28c6582c6 is described below
commit e28c6582c6c864b82c22b6bb4768c092aeafdebd
Author: XiaoHongbo <[email protected]>
AuthorDate: Wed Sep 30 09:02:18 2026 +0800
[python] Support efficient field selection from wide VARIANT objects
(#10320)
---
docs/docs/pypaimon/data-types.md | 17 ++
paimon-python/pypaimon/data/__init__.py | 3 +
paimon-python/pypaimon/data/variant_path.py | 256 +++++++++++++++-
paimon-python/pypaimon/tests/variant_path_test.py | 343 ++++++++++++++++++++++
4 files changed, 613 insertions(+), 6 deletions(-)
diff --git a/docs/docs/pypaimon/data-types.md b/docs/docs/pypaimon/data-types.md
index 535bdd6833..a5bcf1e30e 100644
--- a/docs/docs/pypaimon/data-types.md
+++ b/docs/docs/pypaimon/data-types.md
@@ -218,6 +218,23 @@ Supported Paimon type strings for shredded sub-fields:
`BOOLEAN`, `TINYINT`, `SM
## VARIANT Path Updates
+### Decode selected VARIANT fields to Python
+
+`variant_to_pylist` decodes only the requested top-level fields, returning
+one Python dict per row (or `None` for SQL NULL):
+
+```python
+from pypaimon.data import variant_to_pylist
+
+rows = variant_to_pylist(result.column('payload'), ['state.x', 'action.y'])
+```
+
+Field names are literal (`state.x` is not a nested path). Missing fields are
+omitted; VARIANT NULL is `None`. Values use natural Python types, even when
+types differ between rows. This does not reduce Parquet I/O.
+
+### Read and update typed paths
+
Read existing paths as Arrow arrays, use Arrow compute, and replace them
without decoding unrelated fields:
diff --git a/paimon-python/pypaimon/data/__init__.py
b/paimon-python/pypaimon/data/__init__.py
index 692a8ac925..dae7ba93dd 100644
--- a/paimon-python/pypaimon/data/__init__.py
+++ b/paimon-python/pypaimon/data/__init__.py
@@ -24,6 +24,7 @@ if sys.version_info[:2] < (3, 7):
from pypaimon.data.decimal import Decimal
from pypaimon.data.variant_path import (
variant_get,
+ variant_to_pylist,
variant_replace,
variant_set,
)
@@ -32,6 +33,7 @@ __all__ = [
'Timestamp',
'Decimal',
'variant_get',
+ 'variant_to_pylist',
'variant_replace',
'variant_set',
]
@@ -40,6 +42,7 @@ _MODULE_BY_EXPORT = {
'Timestamp': 'pypaimon.data.timestamp',
'Decimal': 'pypaimon.data.decimal',
'variant_get': 'pypaimon.data.variant_path',
+ 'variant_to_pylist': 'pypaimon.data.variant_path',
'variant_replace': 'pypaimon.data.variant_path',
'variant_set': 'pypaimon.data.variant_path',
}
diff --git a/paimon-python/pypaimon/data/variant_path.py
b/paimon-python/pypaimon/data/variant_path.py
index 79e251e9c0..fea3291d5a 100644
--- a/paimon-python/pypaimon/data/variant_path.py
+++ b/paimon-python/pypaimon/data/variant_path.py
@@ -76,6 +76,11 @@ _STRUCTURE_MATCH_INDEX_BUDGET = 8 * 1024 * 1024
_ROOT_INSERT_SPLICE_PAYLOAD_BUDGET = 8 * 1024 * 1024
# Bound per-row Python and NumPy temporaries for tiny payloads.
_ROOT_INSERT_SPLICE_MAX_BATCH_ROWS = 64 * 1024
+_FAST_CHILD_SIZES = np.asarray([
+ _PRIMITIVE_FIXED_SIZES.get(type_info, 0)
+ if type_info not in (_DECIMAL4, _DECIMAL8, _DECIMAL16) else 0
+ for type_info in range(64)
+], dtype=np.uint8)
@functools.lru_cache(maxsize=256)
@@ -1579,6 +1584,141 @@ def _rowwise_replace_chunk(
chunk, values, data, data_start, rebuilt_rows)
+def _flat_object_unsigneds(value, start, count, width):
+ """Read an object header array without creating Python ints per field."""
+ if width in (1, 2, 4):
+ return np.frombuffer(
+ value, dtype=np.dtype('<u%d' % width), count=count,
+ offset=start)
+ raw = np.frombuffer(
+ value, dtype=np.uint8, count=count * 3,
+ offset=start).reshape(count, 3)
+ return (raw[:, 0].astype(np.uint32)
+ | (raw[:, 1].astype(np.uint32) << 8)
+ | (raw[:, 2].astype(np.uint32) << 16))
+
+
+def _flat_object_layout(value):
+ """Validate a root object and read its offsets as a NumPy array."""
+ limit = len(value)
+ _require_range(0, 2, limit)
+ type_info = (value[0] >> 2) & 0x3F
+ size_width = _U32_SIZE if ((type_info >> 4) & 0x1) else 1
+ _require_range(1, size_width, limit)
+ size = _read_unsigned(value, 1, size_width)
+ id_width = ((type_info >> 2) & 0x3) + 1
+ offset_width = (type_info & 0x3) + 1
+ id_start = 1 + size_width
+ offset_start = id_start + size * id_width
+ data_start = offset_start + (size + 1) * offset_width
+ _require_range(0, data_start, limit)
+ offsets = _flat_object_unsigneds(
+ value, offset_start, size + 1, offset_width)
+ end_offset = int(offsets[-1])
+ if size:
+ ordered = np.sort(offsets)
+ if (int(ordered[0]) != 0 or int(ordered[-1]) != end_offset
+ or np.any(np.diff(ordered) == 0)):
+ _malformed("invalid object offsets")
+ else:
+ if end_offset != 0:
+ _malformed("invalid object offsets")
+ ordered = offsets
+ _require_range(data_start, end_offset, limit)
+ return size, id_width, id_start, offset_start, data_start, offsets, ordered
+
+
+def _flat_object_get_chunk(chunk, values, parsed):
+ """Look up wide, top-level object fields together instead of per path."""
+ if (len(parsed) < 2
+ or any(len(path) != 1 or path[0][0] != 'key'
+ for _, path, _ in parsed)):
+ return None
+
+ valid_rows = _valid_row_indices(chunk, values, chunk.field(1))
+ if not len(valid_rows):
+ return [pa.nulls(len(chunk), type=data_type)
+ for _, _, data_type in parsed]
+ first_value = values.view(int(valid_rows[0]))
+ if not first_value or (first_value[0] & 0x3) != _OBJECT:
+ return None
+ first_size = _checked_object_layout(
+ first_value, 0, len(first_value))[0]
+ # Keep the vectorized reader for narrow objects or a few requested keys.
+ if first_size * len(parsed) < 1024:
+ return None
+ # For multi-row float batches, the existing reader vectorizes across rows.
+ # Its per-path cost wins over per-row lookup on modest object widths.
+ if (len(valid_rows) >= _SLOW_PATH_ROWS
+ and all(pa.types.is_float32(data_type)
+ or pa.types.is_float64(data_type)
+ for _, _, data_type in parsed)
+ and (first_size < 256 or first_size * len(parsed) < 4096)):
+ return None
+
+ metadata = _BinaryValues(chunk.field(1))
+ decoded = [[None] * len(chunk) for _ in parsed]
+ # Object field IDs may differ across rows, so the key includes the actual
+ # encoded layout as well as the metadata dictionary and requested IDs.
+ slot_cache = LRUCache(maxsize=64)
+ for row in valid_rows:
+ row = int(row)
+ value = values.view(row)
+ if not value or (value[0] & 0x3) != _OBJECT:
+ return None
+ row_metadata = bytes(metadata.view(row))
+ key_ids = _cached_metadata_key_ids(row_metadata)
+ targets = {}
+ for index, (_, path, _) in enumerate(parsed):
+ key_id = key_ids.get(path[0][1])
+ if key_id is not None:
+ targets.setdefault(key_id, []).append(index)
+
+ (size, id_width, id_start, offset_start, data_start,
+ offsets, ordered_offsets) = _flat_object_layout(value)
+ if data_start + int(offsets[-1]) != len(value):
+ _malformed("trailing bytes after root value")
+ layout = (row_metadata, bytes(value[id_start:offset_start]),
+ id_width, tuple(sorted(targets)))
+ slots = slot_cache.get(layout)
+ if slots is None:
+ field_ids = _flat_object_unsigneds(
+ value, id_start, size, id_width)
+ if len(np.unique(field_ids)) != size:
+ _malformed("duplicate object field id")
+ slots = {int(key_id): slot
+ for slot, key_id in enumerate(field_ids)
+ if int(key_id) in targets}
+ slot_cache[layout] = slots
+ if not targets:
+ continue
+ if not slots:
+ continue
+ selected_offsets = offsets[np.fromiter(
+ slots.values(), dtype=np.int64, count=len(slots))]
+ next_indices = np.searchsorted(
+ ordered_offsets, selected_offsets, side='right')
+ for (key_id, _), offset, next_index in zip(
+ slots.items(), selected_offsets, next_indices):
+ child_pos = data_start + int(offset)
+ child_end = data_start + int(ordered_offsets[next_index])
+ if _checked_value_size(value, child_pos, child_end) != (
+ child_end - child_pos):
+ _malformed("child size does not match container offsets")
+ for index in targets[key_id]:
+ data_type = parsed[index][2]
+ if (pa.types.is_float32(data_type)
+ or pa.types.is_float64(data_type)):
+ decoded[index][row] = _decode_floating(
+ value, child_pos, data_type)
+ else:
+ decoded[index][row] = _decode_exact(
+ value, row_metadata, child_pos, data_type)
+
+ return [pa.array(result, type=data_type)
+ for result, (_, _, data_type) in zip(decoded, parsed)]
+
+
def _variant_get(column, paths: Mapping[str, pa.DataType]):
parsed = []
for path, target_type in paths.items():
@@ -1594,12 +1734,14 @@ def _variant_get(column, paths: Mapping[str,
pa.DataType]):
result_chunks = {path: [] for path in paths}
for chunk in chunks:
values = _BinaryValues(chunk.field(0))
- results = _vectorized_get_chunk(
- chunk,
- values,
- parsed_paths,
- [target_type for _, _, target_type in parsed],
- )
+ results = _flat_object_get_chunk(chunk, values, parsed)
+ if results is None:
+ results = _vectorized_get_chunk(
+ chunk,
+ values,
+ parsed_paths,
+ [target_type for _, _, target_type in parsed],
+ )
if results is not None:
for (path, _, _), result in zip(parsed, results):
result_chunks[path].append(result)
@@ -1646,6 +1788,108 @@ def variant_get(column, path, data_type=None):
return _variant_get(column, {path: data_type})[path]
+@_with_metadata_cache
+def variant_to_pylist(column, fields: Sequence[str]):
+ """Decode selected top-level VARIANT object fields to Python values.
+
+ Values use their natural Python types, even when a field has different
+ types in different rows. A missing field is omitted from its row's dict;
+ a VARIANT null is present with value None, and a SQL null row returns None.
+ Field names are literal names, not VARIANT path expressions.
+ """
+ if isinstance(fields, (str, bytes)):
+ raise TypeError("VARIANT fields must be a sequence of field names")
+ try:
+ names = tuple(dict.fromkeys(fields))
+ except (TypeError, ValueError):
+ raise TypeError(
+ "VARIANT fields must be a sequence of field names") from None
+ if any(not isinstance(name, str) for name in names):
+ raise TypeError("VARIANT field names must be strings")
+
+ chunks, _, _ = _variant_chunks(column)
+ result = []
+ slot_cache = LRUCache(maxsize=64)
+ for chunk in chunks:
+ values = _BinaryValues(chunk.field(0))
+ metadata = _BinaryValues(chunk.field(1))
+ valid = set(map(int, _valid_row_indices(
+ chunk, values, chunk.field(1))))
+ for row in range(len(chunk)):
+ if row not in valid:
+ result.append(None)
+ continue
+ value = bytes(values.view(row))
+ row_metadata = bytes(metadata.view(row))
+ if not value or (value[0] & 0x3) != _OBJECT:
+ raise TypeError("VARIANT root must be an object")
+ variant = GenericVariant(value, row_metadata)
+ key_ids = _cached_metadata_key_ids(row_metadata)
+ targets = {key_ids[name]: name for name in names
+ if name in key_ids}
+ (size, id_width, id_start, offset_start, data_start,
+ offsets, ordered_offsets) = _flat_object_layout(value)
+ if data_start + int(offsets[-1]) != len(value):
+ _malformed("trailing bytes after root value")
+ layout = (row_metadata, value[id_start:offset_start],
+ id_width, tuple(sorted(targets)))
+ selected_slots = slot_cache.get(layout)
+ if selected_slots is None:
+ field_ids = _flat_object_unsigneds(
+ value, id_start, size, id_width)
+ if len(np.unique(field_ids)) != size:
+ _malformed("duplicate object field id")
+ selected_slots = tuple(
+ (slot, targets[int(key_id)])
+ for slot, key_id in enumerate(field_ids)
+ if int(key_id) in targets)
+ slot_cache[layout] = selected_slots
+ if not targets:
+ result.append({})
+ continue
+ slots = np.fromiter(
+ (slot for slot, _ in selected_slots),
+ dtype=np.int64, count=len(selected_slots))
+ starts = offsets[slots]
+ if np.any(starts >= offsets[-1]):
+ _malformed("invalid object offset")
+ next_indices = np.searchsorted(
+ ordered_offsets, starts, side='right')
+ ends = ordered_offsets[next_indices]
+ # Encoded offsets may be uint8/uint16. Adding data_start before
+ # widening can wrap or raise for otherwise valid objects.
+ absolute_starts = starts.astype(np.intp, copy=False) + data_start
+ headers = np.frombuffer(value, dtype=np.uint8)[absolute_starts]
+ basic_types = headers & 0x3
+ type_infos = headers >> 2
+ fixed_sizes = _FAST_CHILD_SIZES[type_infos]
+ is_short_string = basic_types == _SHORT_STR
+ is_fast = (is_short_string
+ | ((basic_types == _PRIMITIVE) & (fixed_sizes != 0)))
+ expected_sizes = np.where(
+ is_short_string, type_infos + 1, fixed_sizes)
+ if np.any(is_fast & (ends - starts != expected_sizes)):
+ _malformed("child size does not match container offsets")
+ selected = {}
+ for index, ((_, name), offset, end) in enumerate(zip(
+ selected_slots, starts, ends)):
+ child_pos = data_start + int(offset)
+ if not is_fast[index]:
+ child_end = data_start + int(end)
+ if basic_types[index] in (_OBJECT, _ARRAY):
+ _validate_value_field_ids(
+ value, child_pos, child_end, len(key_ids))
+ elif _checked_value_size(
+ value, child_pos, child_end) != (
+ child_end - child_pos):
+ _malformed(
+ "child size does not match container offsets")
+ selected[name] = variant._to_python_impl(
+ value, row_metadata, child_pos)
+ result.append(selected)
+ return result
+
+
def _paths_overlap(first: _Path, second: _Path) -> bool:
limit = min(len(first), len(second))
return first[:limit] == second[:limit]
diff --git a/paimon-python/pypaimon/tests/variant_path_test.py
b/paimon-python/pypaimon/tests/variant_path_test.py
index b3a569ea68..e1a070e262 100644
--- a/paimon-python/pypaimon/tests/variant_path_test.py
+++ b/paimon-python/pypaimon/tests/variant_path_test.py
@@ -26,12 +26,15 @@ import pyarrow.compute as pc
from pypaimon.data._variant_binary import _primitive_header
from pypaimon.data.generic_variant import _DOUBLE, GenericVariant
from pypaimon.data.variant_path import (
+ _checked_object_layout,
_compile_paths,
_metadata_cache,
_metadata_key_ids,
_path_positions,
_rebuilt_offsets,
+ _vectorized_get_chunk,
variant_get,
+ variant_to_pylist,
variant_replace,
)
from pypaimon.data.variant_shredding import (
@@ -77,8 +80,161 @@ def _typed_object(fields):
GenericVariant(value, metadata)])
+class TestVariantToPylist(unittest.TestCase):
+
+ def test_preserves_mixed_types_missing_null_and_literal_names(self):
+ column = _variants([
+ {'state.x': 1, 'child': {'count': 2}, 'nullable': None},
+ {'state.x': 1.5, 'child': [3, 4]},
+ {},
+ None,
+ ])
+
+ result = variant_to_pylist(
+ column, ['state.x', 'child', 'nullable', 'absent'])
+
+ self.assertEqual(result, [
+ {'state.x': 1, 'child': {'count': 2}, 'nullable': None},
+ {'state.x': 1.5, 'child': [3, 4]},
+ {},
+ None,
+ ])
+
+ def test_selects_from_wide_objects_and_chunked_arrays(self):
+ first = {'field.%03d' % i: float(i) for i in range(128)}
+ second = {'field.%03d' % i: i for i in range(128)}
+ column = pa.chunked_array([
+ _variants([first]), _variants([second])])
+
+ result = variant_to_pylist(
+ column, ['field.000', 'field.127', 'field.000'])
+
+ self.assertEqual(result, [
+ {'field.000': 0.0, 'field.127': 127.0},
+ {'field.000': 0, 'field.127': 127},
+ ])
+
+ def test_selects_from_same_metadata_with_different_slots(self):
+ fields = {'field.%03d' % i: float(i) for i in range(128)}
+ metadata = GenericVariant.from_python(fields).metadata()
+ key_ids = _metadata_key_ids(metadata)
+ rows = []
+ for names in (list(fields), list(fields)[10:], list(fields)):
+ value = _build_object_value([
+ (key_ids[name], _encode_scalar_to_value_bytes(
+ fields[name], pa.float64()))
+ for name in names
+ ])
+ rows.append(GenericVariant(value, metadata))
+
+ result = variant_to_pylist(
+ GenericVariant.to_arrow_array(rows),
+ ['field.000', 'field.015'])
+
+ self.assertEqual(result, [
+ {'field.000': 0.0, 'field.015': 15.0},
+ {'field.015': 15.0},
+ {'field.000': 0.0, 'field.015': 15.0},
+ ])
+
+ def test_preserves_other_primitive_types(self):
+ timestamp = datetime.datetime(2026, 8, 11, 1, 2, 3, 456000)
+ fields = {
+ 'flag': (True, pa.bool_()),
+ 'text': ('hello', pa.string()),
+ 'binary': (b'abc', pa.binary()),
+ 'decimal': (Decimal('12.30'), pa.decimal128(4, 2)),
+ 'date': (datetime.date(2026, 8, 11), pa.date32()),
+ 'timestamp': (timestamp, pa.timestamp('us')),
+ }
+
+ result = variant_to_pylist(_typed_object(fields), fields)
+
+ self.assertEqual(result, [
+ {name: value for name, (value, _) in fields.items()}])
+
+ def test_rejects_non_object_root_and_invalid_fields(self):
+ with self.assertRaisesRegex(TypeError, "root must be an object"):
+ variant_to_pylist(_variants([[1, 2]]), ['field'])
+ with self.assertRaisesRegex(TypeError, "sequence of field names"):
+ variant_to_pylist(_variants([{}]), 'field')
+ with self.assertRaisesRegex(TypeError, "field names must be strings"):
+ variant_to_pylist(_variants([{}]), [1])
+
+ def test_rejects_truncated_object(self):
+ variant = GenericVariant.from_python({'field': 123})
+ truncated = GenericVariant(
+ variant.value()[:-1], variant.metadata())
+ with self.assertRaisesRegex(ValueError, "MALFORMED_VARIANT"):
+ variant_to_pylist(
+ GenericVariant.to_arrow_array([truncated]), ['field'])
+
+ def test_rejects_truncated_selected_child_before_next_field(self):
+ valid = GenericVariant.from_python({'a': 1.0, 'b': 2.0})
+ key_ids = _metadata_key_ids(valid.metadata())
+ value = _build_object_value([
+ (key_ids['a'], bytes([_primitive_header(_DOUBLE)])),
+ (key_ids['b'], _encode_scalar_to_value_bytes(2.0, pa.float64())),
+ ])
+ column = GenericVariant.to_arrow_array([
+ GenericVariant(value, valid.metadata())])
+
+ with self.assertRaisesRegex(ValueError, "MALFORMED_VARIANT"):
+ variant_to_pylist(column, ['a'])
+ self.assertEqual(variant_to_pylist(column, ['b']), [{'b': 2.0}])
+
+ def test_rejects_truncated_nested_selected_child(self):
+ valid = GenericVariant.from_python(
+ {'a': {'nested': 1.0}, 'b': 2.0})
+ key_ids = _metadata_key_ids(valid.metadata())
+ nested = _build_object_value([
+ (key_ids['nested'], bytes([_primitive_header(_DOUBLE)])),
+ ])
+ value = _build_object_value([
+ (key_ids['a'], nested),
+ (key_ids['b'], _encode_scalar_to_value_bytes(2.0, pa.float64())),
+ ])
+ column = GenericVariant.to_arrow_array([
+ GenericVariant(value, valid.metadata())])
+
+ with self.assertRaisesRegex(ValueError, "MALFORMED_VARIANT"):
+ variant_to_pylist(column, ['a'])
+
+ def test_selected_offsets_cross_encoded_integer_width(self):
+ for count, value in ((100, None), (128, None), (5362, 1.5)):
+ with self.subTest(count=count):
+ name = 'field%04d' % (count - 1)
+ variant = GenericVariant.from_python({
+ 'field%04d' % index: value for index in range(count)
+ })
+ column = GenericVariant.to_arrow_array([variant])
+ self.assertEqual(
+ variant_to_pylist(column, [name]), [{name: value}])
+
+
class TestVariantGet(unittest.TestCase):
+ def test_medium_float_batch_keeps_vectorized_reader(self):
+ column = _variants([
+ {'field%03d' % i: float(i + row) for i in range(128)}
+ for row in range(128)
+ ])
+ paths = {'$.field%03d' % i: pa.float64() for i in range(16)}
+ vectorized_results = []
+
+ def track_vectorized(*args):
+ result = _vectorized_get_chunk(*args)
+ vectorized_results.append(result is not None)
+ return result
+
+ with patch('pypaimon.data.variant_path._vectorized_get_chunk',
+ side_effect=track_vectorized):
+ result = variant_get(column, paths)
+
+ self.assertEqual(vectorized_results, [True])
+ self.assertEqual(result['$.field000'].to_pylist(),
+ [float(row) for row in range(128)])
+
def test_compile_paths_builds_trie_without_prefix_slices(self):
class NoSlicePath(tuple):
def __getitem__(self, item):
@@ -155,6 +311,193 @@ class TestVariantGet(unittest.TestCase):
self.assertEqual(result['$.velocity.x'].to_pylist(), [1.0, 3.0])
self.assertEqual(result['$.velocity.y'].to_pylist(), [-2.0, -4.0])
+ def test_reads_many_flat_fields_without_repeating_vectorized_scan(self):
+ fields = {'field.%03d' % i: float(i) for i in range(128)}
+ paths = {'$["field.%03d"]' % i: pa.float64()
+ for i in range(16)}
+ column = _variants([fields, fields])
+
+ with patch(
+ 'pypaimon.data.variant_path._vectorized_get_chunk',
+ side_effect=AssertionError(
+ "wide flat lookup should be shared")):
+ result = variant_get(column, paths)
+
+ for i, path in enumerate(paths):
+ self.assertEqual(result[path].to_pylist(), [float(i), float(i)])
+
+ def test_reads_many_flat_fields_with_exact_primitive_types(self):
+ timestamp = datetime.datetime(2026, 8, 11, 1, 2, 3, 456000)
+ fields = {
+ 'field.%03d' % i: (float(i), pa.float64())
+ for i in range(128)
+ }
+ selected = {
+ 'flag': (True, pa.bool_()),
+ 'count': (123, pa.int64()),
+ 'text': ('hello', pa.string()),
+ 'binary': (b'abc', pa.binary()),
+ 'decimal': (Decimal('12.30'), pa.decimal128(4, 2)),
+ 'date': (datetime.date(2026, 8, 11), pa.date32()),
+ 'timestamp': (timestamp, pa.timestamp('us')),
+ }
+ fields.update(selected)
+ paths = {'$.%s' % name: data_type
+ for name, (_, data_type) in selected.items()}
+ paths['$["field.000"]'] = pa.float64()
+
+ with patch(
+ 'pypaimon.data.variant_path._vectorized_get_chunk',
+ side_effect=AssertionError(
+ "wide flat lookup should be shared")):
+ result = variant_get(_typed_object(fields), paths)
+
+ for name, (expected, _) in selected.items():
+ self.assertEqual(result['$.%s' % name].to_pylist(), [expected])
+ self.assertEqual(result['$["field.000"]'].to_pylist(), [0.0])
+
+ def test_reads_many_flat_fields_with_complex_types(self):
+ fields = {'field.%03d' % i: float(i) for i in range(128)}
+ fields.update({
+ 'object': {'count': 2, 'flag': True},
+ 'array': [1, 2],
+ 'map': {'left': 1, 'right': 2},
+ })
+ paths = {'$["field.%03d"]' % i: pa.float64()
+ for i in range(5)}
+ paths.update({
+ '$.object': pa.struct([
+ ('count', pa.int64()), ('flag', pa.bool_())]),
+ '$.array': pa.list_(pa.int64()),
+ '$.map': pa.map_(pa.string(), pa.int64()),
+ })
+
+ with patch(
+ 'pypaimon.data.variant_path._vectorized_get_chunk',
+ side_effect=AssertionError(
+ "wide flat lookup should be shared")):
+ result = variant_get(_variants([fields]), paths)
+
+ self.assertEqual(result['$.object'].to_pylist(),
+ [{'count': 2, 'flag': True}])
+ self.assertEqual(result['$.array'].to_pylist(), [[1, 2]])
+ self.assertEqual(result['$.map'].to_pylist(),
+ [[('left', 1), ('right', 2)]])
+
+ def test_many_flat_fields_keep_missing_null_and_mixed_metadata(self):
+ fields = {'field.%03d' % i: float(i) for i in range(128)}
+ paths = {'$["field.%03d"]' % i: pa.float64()
+ for i in range(15)}
+ paths['$["maybe.null"]'] = pa.float64()
+ column = _variants([
+ {**fields, 'maybe.null': None},
+ {**fields, 'field.000': None},
+ fields,
+ None,
+ ])
+
+ result = variant_get(column, paths)
+
+ self.assertEqual(
+ result['$["field.000"]'].to_pylist(),
+ [0.0, None, 0.0, None])
+ self.assertEqual(
+ result['$["maybe.null"]'].to_pylist(),
+ [None, None, None, None])
+ self.assertEqual(
+ result['$["field.014"]'].to_pylist(),
+ [14.0, 14.0, 14.0, None])
+
+ def test_many_flat_fields_allow_nonmonotonic_value_offsets(self):
+ fields = {'field.%03d' % i: float(i) for i in range(128)}
+ variant = GenericVariant.from_python(fields)
+ value = variant.value()
+ size, id_width, id_start, data_start, offsets, _ = (
+ _checked_object_layout(value, 0, len(value)))
+ offset_width = ((value[0] >> 2) & 0x3) + 1
+ offset_start = id_start + size * id_width
+ children = [value[data_start + offsets[i]:data_start + offsets[i + 1]]
+ for i in range(size)]
+ reordered = bytearray(value[:data_start])
+ for i in range(size):
+ new_offset = sum(len(child) for child in children[i + 1:])
+ start = offset_start + i * offset_width
+ reordered[start:start + offset_width] = new_offset.to_bytes(
+ offset_width, 'little')
+ reordered.extend(b''.join(reversed(children)))
+ column = GenericVariant.to_arrow_array([
+ GenericVariant(bytes(reordered), variant.metadata())])
+ paths = {'$["field.%03d"]' % i: pa.float64()
+ for i in range(16)}
+
+ result = variant_get(column, paths)
+
+ for i, path in enumerate(paths):
+ self.assertEqual(result[path].to_pylist(), [float(i)])
+
+ def test_many_flat_fields_with_three_byte_offsets(self):
+ fields = {'field.%03d' % i: 'x' * 1024 for i in range(128)}
+ variant = GenericVariant.from_python(fields)
+ offset_width = ((variant.value()[0] >> 2) & 0x3) + 1
+ self.assertEqual(offset_width, 3)
+ paths = {'$["field.%03d"]' % i: pa.string()
+ for i in range(16)}
+
+ result = variant_get(_variants([fields, fields]), paths)
+
+ for path in paths:
+ self.assertEqual(result[path].to_pylist(), ['x' * 1024] * 2)
+
+ def test_many_flat_fields_same_metadata_different_layouts(self):
+ fields = {'field.%03d' % i: float(i) for i in range(128)}
+ metadata = GenericVariant.from_python(fields).metadata()
+ key_ids = _metadata_key_ids(metadata)
+ rows = []
+ for names in (list(fields), list(fields)[10:], list(fields)):
+ value = _build_object_value([
+ (key_ids[name], _encode_scalar_to_value_bytes(
+ fields[name], pa.float64()))
+ for name in names
+ ])
+ rows.append(GenericVariant(value, metadata))
+ paths = {'$["field.%03d"]' % i: pa.float64()
+ for i in range(16)}
+
+ result = variant_get(GenericVariant.to_arrow_array(rows), paths)
+
+ self.assertEqual(result['$["field.000"]'].to_pylist(),
+ [0.0, None, 0.0])
+ self.assertEqual(result['$["field.015"]'].to_pylist(),
+ [15.0, 15.0, 15.0])
+
+ def test_many_flat_fields_reject_duplicate_id_and_offset(self):
+ fields = {'field.%03d' % i: float(i) for i in range(128)}
+ variant = GenericVariant.from_python(fields)
+ original = variant.value()
+ size, id_width, id_start, _, _, _ = (
+ _checked_object_layout(original, 0, len(original)))
+ offset_width = ((original[0] >> 2) & 0x3) + 1
+ offset_start = id_start + size * id_width
+ paths = {'$["field.%03d"]' % i: pa.float64()
+ for i in range(16)}
+
+ duplicate_id = bytearray(original)
+ duplicate_id[id_start + id_width:id_start + 2 * id_width] = (
+ duplicate_id[id_start:id_start + id_width])
+ with self.assertRaisesRegex(ValueError, 'duplicate object field id'):
+ variant_get(GenericVariant.to_arrow_array([
+ GenericVariant(bytes(duplicate_id), variant.metadata())
+ ]), paths)
+
+ duplicate_offset = bytearray(original)
+ duplicate_offset[
+ offset_start + offset_width:offset_start + 2 * offset_width
+ ] = duplicate_offset[offset_start:offset_start + offset_width]
+ with self.assertRaisesRegex(ValueError, 'invalid object offsets'):
+ variant_get(GenericVariant.to_arrow_array([
+ GenericVariant(bytes(duplicate_offset), variant.metadata())
+ ]), paths)
+
def test_requires_exact_type(self):
cases = (
(_float_variants([1.25]), pa.float64()),
