This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit c1f92f1a415850f7724256631177cbab506c4f80 Author: XiaoHongbo <[email protected]> AuthorDate: Wed Sep 16 10:29:49 2026 +0800 [python] Speed up nested VARIANT reads by vectorizing offsets (#9847) --- .../pypaimon/read/reader/format_pyarrow_reader.py | 16 +++++------ .../tests/format_pyarrow_variant_row_group_test.py | 32 +++++++++++++++++++++- 2 files changed, 39 insertions(+), 9 deletions(-) diff --git a/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py b/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py index f3403e260b..7ce7ec1d7a 100644 --- a/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py +++ b/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py @@ -25,6 +25,7 @@ from concurrent.futures import Future from typing import Any, Callable, Deque, Dict, Iterator, List, Optional, Set, Tuple import pyarrow as pa +import pyarrow.compute as pc import pyarrow.dataset as ds from pyarrow import RecordBatch @@ -1125,11 +1126,10 @@ def _normalized_offsets(column): [None, column.buffers()[1]], offset=column.offset, ) - raw_offsets = offsets_array.to_pylist() - start = raw_offsets[0] - end = raw_offsets[-1] - offsets = [value - start for value in raw_offsets] - for index, is_null in enumerate(column.is_null().to_pylist()): - if is_null: - offsets[index] = None - return pa.array(offsets, type=offsets_array.type), start, end + start = offsets_array[0].as_py() + end = offsets_array[-1].as_py() + offsets = pc.subtract(offsets_array, pa.scalar(start, type=offsets_array.type)) + if column.null_count: + mask = pa.concat_arrays([column.is_null(), pa.array([False])]) + offsets = pc.if_else(mask, pa.scalar(None, type=offsets_array.type), offsets) + return offsets, start, end diff --git a/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py b/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py index bb0f06bc6e..ef58bcbe51 100644 --- a/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py +++ b/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py @@ -33,7 +33,10 @@ from pypaimon.data.variant_shredding import ( shredding_schema_to_arrow_type, shred_variant_column, ) -from pypaimon.read.reader.format_pyarrow_reader import FormatPyArrowReader +from pypaimon.read.reader.format_pyarrow_reader import ( + FormatPyArrowReader, + _normalized_offsets, +) from pypaimon.schema.data_types import ( ArrayType, AtomicType, @@ -71,6 +74,33 @@ def _drain(reader): return rows, columns, content_keys +class NormalizedOffsetsTest(unittest.TestCase): + + def test_nested_offsets(self): + for arrow_type, values in [ + (pa.list_(pa.int64()), [[1, 2], None, [], [3], None]), + (pa.large_list(pa.int64()), [[1, 2], None, [], [3], None]), + (pa.map_(pa.string(), pa.int64()), + [[('a', 1), ('b', 2)], None, [], [('c', 3)], None]), + ]: + array = pa.array(values, type=arrow_type) + for column in [array, array.slice(1, 3), array.slice(4), + array.slice(2, 0), pa.array([], type=arrow_type)]: + with self.subTest(type=arrow_type, length=len(column), offset=column.offset): + offset_type = pa.int64() if pa.types.is_large_list(arrow_type) else pa.int32() + raw = pa.Array.from_buffers( + offset_type, len(column) + 1, + [None, column.buffers()[1]], offset=column.offset).to_pylist() + expected = [value - raw[0] for value in raw] + for i, null in enumerate(column.is_null().to_pylist()): + if null: + expected[i] = None + offsets, start, end = _normalized_offsets(column) + self.assertEqual((start, end), (raw[0], raw[-1])) + self.assertEqual(offsets.type, offset_type) + self.assertEqual(offsets.to_pylist(), expected) + + class VariantRowGroupReaderTest(unittest.TestCase): def setUp(self):
