From d22f1da28053560c2c9db5a8414f2e5eb0f8139e Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 21 Sep 2026 02:40:13 -0700 Subject: [PATCH] [python] Read VARIANT pages with Parquet OffsetIndex --- paimon-python/README.md | 5 +- .../read/reader/format_pyarrow_reader.py | 3 +- .../pypaimon/tests/parquet_page_index_test.py | 115 ++++++++++++++++++ 3 files changed, 119 insertions(+), 4 deletions(-) diff --git a/paimon-python/README.md b/paimon-python/README.md index cb1836fc57ee..a711f71345eb 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -108,8 +108,9 @@ The command will install the package and core dependencies to your local Python # Parquet page-index reads For row-tracking tables with a Parquet OffsetIndex, PyPaimon can read a -contiguous `_ROW_ID` range without decoding the full row group. This is enabled -by default and can be disabled with the table option: +contiguous `_ROW_ID` range, including standard VARIANT columns, without decoding +the full row group. This is enabled by default and can be disabled with the table +option: ```python table = table.copy({"parquet.filter.columnindex.enabled": "false"}) diff --git a/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py b/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py index df9e3e70020f..ed6dd37215f7 100644 --- a/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py +++ b/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py @@ -542,8 +542,7 @@ def __init__(self, file_io: FileIO, file_format: str, file_path: str, if (self._selected_parquet_row_groups is not None and options is not None and options.parquet_column_index_enabled() - and self._row_group_cache is None - and not self._bounded_variant_read): + and self._row_group_cache is None): from pypaimon.read.reader.parquet_page_index_reader import ( ParquetPageIndexReader, ) diff --git a/paimon-python/pypaimon/tests/parquet_page_index_test.py b/paimon-python/pypaimon/tests/parquet_page_index_test.py index 8a33e0b1dc80..68b14d239e4d 100644 --- a/paimon-python/pypaimon/tests/parquet_page_index_test.py +++ b/paimon-python/pypaimon/tests/parquet_page_index_test.py @@ -42,6 +42,10 @@ FIELDS = [DataField(0, 'id', AtomicType('BIGINT')), DataField(1, 'payload', AtomicType('STRING'))] PAGE_INDEX_OPTIONS = CoreOptions(Options({'parquet.filter.columnindex.enabled': 'true'})) +VARIANT_TYPE = pa.struct([ + pa.field('value', pa.binary(), nullable=False), + pa.field('metadata', pa.binary(), nullable=False), +]) @pytest.fixture @@ -177,6 +181,117 @@ def test_sparse_row_indices_are_normalized(fixture): assert _read(fixture, row_ranges=[(N, N + 10)])[0] is None +@pytest.mark.parametrize('selected_row', [0, 3009, 3011, 4095]) +@pytest.mark.parametrize('page_version', ['1.0', '2.0']) +def test_standard_variant_reads_selected_pages( + tmp_path, selected_row, page_version): + count = 4096 + table = pa.table({ + 'record_value': pa.array([ + None if index % 17 == 0 else { + 'value': hashlib.shake_256(str(index).encode()).digest(512), + 'metadata': hashlib.shake_256( + ('metadata-%d' % index).encode()).digest(512), + } + for index in range(count) + ], type=VARIANT_TYPE), + }) + path = str(tmp_path / 'variant.parquet') + pq.write_table( + table, path, write_page_index=True, data_page_size=4096, + write_batch_size=32, row_group_size=count, + use_dictionary=False, compression='zstd', + data_page_version=page_version) + counter = _CountingLocalFileSystem(skip_instance_cache=True) + file_io = LocalFileIO(str(tmp_path), Options({})) + file_io.filesystem = pafs.PyFileSystem(pafs.FSSpecHandler(counter)) + fields = [DataField(0, 'record_value', AtomicType('VARIANT'))] + + def read(options): + counter.reset_counts() + reader = reader_module.FormatPyArrowReader( + file_io, 'parquet', path, fields, None, + row_ranges=[(selected_row, selected_row)], + batch_size=31, options=options) + try: + if options is not None: + assert reader._page_index_reader is not None + batches = [] + while True: + batch = reader.read_arrow_batch() + if batch is None: + break + batches.append(batch) + return pa.Table.from_batches(batches), sum( + size for _, size in counter.reads) + finally: + reader.close() + + baseline, baseline_bytes = read(None) + with patch.object( + page_module.ParquetPageIndexReader, '_column_payload', + autospec=True, + side_effect=page_module.ParquetPageIndexReader._column_payload, + ) as selected_pages: + optimized, optimized_bytes = read(PAGE_INDEX_OPTIONS) + assert optimized.equals(baseline) + assert optimized.equals(table.slice(selected_row, 1)) + assert [len(call.args[1][4]) for call in selected_pages.call_args_list] == [1, 1] + assert optimized_bytes < baseline_bytes + + +def test_unprofitable_variant_page_selection_uses_ordinary_reader(tmp_path): + count = 4096 + table = pa.table({ + 'record_value': pa.array([ + { + 'value': hashlib.shake_256(str(index).encode()).digest(512), + 'metadata': b'\x01', + } + for index in range(count) + ], type=VARIANT_TYPE), + }) + path = str(tmp_path / 'unprofitable-variant.parquet') + pq.write_table( + table, path, write_page_index=True, data_page_size=4096, + write_batch_size=32, row_group_size=count, + use_dictionary=True, compression='zstd') + reader = reader_module.FormatPyArrowReader( + LocalFileIO(str(tmp_path), Options({})), 'parquet', path, + [DataField(0, 'record_value', AtomicType('VARIANT'))], None, + row_ranges=[(3009, 3009)], options=PAGE_INDEX_OPTIONS) + try: + assert reader._page_index_reader is not None + with patch.object( + page_module.ParquetPageIndexReader, '_column_payload', + side_effect=AssertionError('must use the ordinary reader')): + result = reader.read_arrow_batch() + assert pa.Table.from_batches([result]).equals(table.slice(3009, 1)) + finally: + reader.close() + + +def test_variant_without_offset_index_uses_ordinary_reader(tmp_path): + table = pa.table({ + 'record_value': pa.array([ + {'value': b'one', 'metadata': b'\x01'}, + {'value': b'two', 'metadata': b'\x01'}, + ], type=VARIANT_TYPE), + }) + path = str(tmp_path / 'unindexed-variant.parquet') + pq.write_table(table, path, write_page_index=False) + reader = reader_module.FormatPyArrowReader( + LocalFileIO(str(tmp_path), Options({})), 'parquet', path, + [DataField(0, 'record_value', AtomicType('VARIANT'))], None, + row_ranges=[(1, 1)], options=PAGE_INDEX_OPTIONS) + try: + assert reader._page_index_reader is None + assert pa.Table.from_batches([reader.read_arrow_batch()]).equals( + table.slice(1, 1)) + finally: + reader.close() + + def test_concurrent_readers_and_early_close(fixture): path, table, file_io, _ = fixture