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 9afdb70afe [python] Speed up nested VARIANT reads by vectorizing
offsets (#9847)
9afdb70afe is described below
commit 9afdb70afe12af938b2be2359070eac1d8b8bdbb
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):
