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 a4bfec1bd2 [python] Prune Parquet row groups for row-id range reads
(#8893)
a4bfec1bd2 is described below
commit a4bfec1bd2242847510c09ff7faf017a9ee2efbb
Author: XiaoHongbo <[email protected]>
AuthorDate: Wed Jul 29 13:55:19 2026 +0800
[python] Prune Parquet row groups for row-id range reads (#8893)
---
.../pypaimon/read/reader/data_file_batch_reader.py | 49 ++-
.../pypaimon/read/reader/format_pyarrow_reader.py | 192 +++++++++-
paimon-python/pypaimon/read/split_read.py | 39 +-
.../pypaimon/tests/parquet_row_range_test.py | 393 +++++++++++++++++++++
4 files changed, 658 insertions(+), 15 deletions(-)
diff --git a/paimon-python/pypaimon/read/reader/data_file_batch_reader.py
b/paimon-python/pypaimon/read/reader/data_file_batch_reader.py
index 1de9f7a405..2bdf38689e 100644
--- a/paimon-python/pypaimon/read/reader/data_file_batch_reader.py
+++ b/paimon-python/pypaimon/read/reader/data_file_batch_reader.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
-from typing import List, Optional
+from typing import List, Optional, Tuple
import pyarrow as pa
import pyarrow.compute as pc
@@ -106,6 +106,35 @@ def _constructed_to_string_array(array, file_type):
return pa.array(out, type=pa.string())
+class _RowIdRangeCursor:
+ """Generate selected row offsets batch by batch from compact ranges."""
+
+ def __init__(self, ranges: List[Tuple[int, int]]):
+ self._ranges = ranges
+ self._range_index = 0
+ self._next_offset = ranges[0][0] if ranges else None
+
+ def take(self, count: int) -> List[int]:
+ offsets = []
+ while len(offsets) < count:
+ if self._range_index >= len(self._ranges):
+ raise ValueError(
+ "Row range offsets were exhausted before the reader")
+ lower, upper = self._ranges[self._range_index]
+ current = max(lower, self._next_offset)
+ take_count = min(count - len(offsets), upper - current + 1)
+ offsets.extend(range(current, current + take_count))
+ current += take_count
+ if current > upper:
+ self._range_index += 1
+ self._next_offset = (
+ self._ranges[self._range_index][0]
+ if self._range_index < len(self._ranges) else None)
+ else:
+ self._next_offset = current
+ return offsets
+
+
class DataFileBatchReader(RecordBatchReader):
"""
Reads record batch from files of different formats
@@ -120,7 +149,13 @@ class DataFileBatchReader(RecordBatchReader):
file_io: Optional[FileIO] = None,
row_id_offsets: Optional[List[int]] = None,
file_data_fields: Optional[List[DataField]] = None,
- target_data_fields: Optional[List[DataField]] = None):
+ target_data_fields: Optional[List[DataField]] = None,
+ row_id_offset_ranges: Optional[List[Tuple[int, int]]] = None):
+ if (row_id_offsets is not None
+ and row_id_offset_ranges is not None):
+ raise ValueError(
+ "row_id_offsets and row_id_offset_ranges cannot both "
+ "be provided")
self.format_reader = format_reader
self.index_mapping = index_mapping
self.partition_info = partition_info
@@ -130,6 +165,9 @@ class DataFileBatchReader(RecordBatchReader):
self.first_row_id = first_row_id
self.row_id_offsets = row_id_offsets
self._row_id_cursor = 0
+ self._row_id_range_cursor = (
+ _RowIdRangeCursor(row_id_offset_ranges)
+ if row_id_offset_ranges is not None else None)
self.max_sequence_number = max_sequence_number
self.system_fields = system_fields
self.file_io = file_io
@@ -351,6 +389,13 @@ class DataFileBatchReader(RecordBatchReader):
row_ids = [self.first_row_id + o for o in
self.row_id_offsets[self._row_id_cursor:end]]
arrays[idx] = pa.array(row_ids, type=pa.int64())
self._row_id_cursor = end
+ elif self._row_id_range_cursor is not None:
+ row_ids = [
+ self.first_row_id + offset
+ for offset in self._row_id_range_cursor.take(
+ record_batch.num_rows)
+ ]
+ arrays[idx] = pa.array(row_ids, type=pa.int64())
else:
row_id_range = range(self.first_row_id, self.first_row_id +
record_batch.num_rows)
arrays[idx] = pa.array(row_id_range, type=pa.int64())
diff --git a/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py
b/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py
index 956ca39f46..565ec1f5b1 100644
--- a/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py
+++ b/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py
@@ -15,7 +15,8 @@
# specific language governing permissions and limitations
# under the License.
-from typing import Any, Dict, List, Optional, Set
+from collections import deque
+from typing import Any, Deque, Dict, Iterator, List, Optional, Set, Tuple
import pyarrow as pa
import pyarrow.dataset as ds
@@ -57,10 +58,61 @@ class FormatPyArrowReader(RecordBatchReader):
push_down_predicate: Any, batch_size: int = 1024,
options: CoreOptions = None,
nested_name_paths: Optional[List[List[str]]] = None,
- predicate_field_names: Optional[Set[str]] = None):
+ predicate_field_names: Optional[Set[str]] = None,
+ row_indices: Optional[List[int]] = None,
+ row_ranges: Optional[List[Tuple[int, int]]] = None):
self._predicate_field_names = predicate_field_names or set()
file_path_for_pyarrow = file_io.to_filesystem_path(file_path)
self.dataset = ds.dataset(file_path_for_pyarrow, format=file_format,
filesystem=file_io.filesystem)
+ self._range_slicer = None
+ self._selected_parquet_row_groups = None
+ self._exhausted = False
+ if row_indices is not None and row_ranges is not None:
+ raise ValueError(
+ "row_indices and row_ranges cannot both be provided")
+ row_selection_supplied = (
+ row_indices is not None or row_ranges is not None)
+ if row_selection_supplied and file_format == 'parquet':
+ if push_down_predicate is not None:
+ raise ValueError(
+ "row selections cannot be combined with a scanner-level "
+ "push-down predicate because filtering shifts row "
+ "positions")
+ runs = (
+ _normalize_runs(row_ranges)
+ if row_ranges is not None
+ else _to_runs(row_indices)
+ )
+ if not runs:
+ self._exhausted = True
+ else:
+ fragment = next(iter(self.dataset.get_fragments()), None)
+ row_group_fragments = (
+ list(fragment.split_by_row_group())
+ if fragment is not None else [])
+ selected_infos = []
+ selected_ids = []
+ offset = 0
+ run_index = 0
+ for row_group_fragment in row_group_fragments:
+ row_group = row_group_fragment.row_groups[0]
+ row_count = row_group.num_rows
+ lower, upper = offset, offset + row_count - 1
+ while (
+ run_index < len(runs)
+ and runs[run_index][1] < lower):
+ run_index += 1
+ if (run_index < len(runs)
+ and runs[run_index][0] <= upper):
+ selected_infos.append((offset, row_count))
+ selected_ids.append(row_group.id)
+ offset += row_count
+ if not selected_ids:
+ self._exhausted = True
+ else:
+ self._selected_parquet_row_groups = selected_ids
+ self._range_slicer = _RowRunSlicer(
+ selected_infos, runs)
self._file_format = file_format
self.read_fields = read_fields
self._read_field_names = [f.name for f in read_fields]
@@ -115,12 +167,15 @@ class FormatPyArrowReader(RecordBatchReader):
# Read projected VARIANT columns in bounded batches.
self._parquet_file = None
- if self._bounded_variant_read:
+ if (self._bounded_variant_read
+ or self._selected_parquet_row_groups is not None):
import pyarrow.parquet as pq
# ParquetFile(filesystem=...) is unavailable in PyArrow 6.
self._parquet_file = pq.ParquetFile(
file_io.filesystem.open_input_file(file_path_for_pyarrow))
- if self._parquet_file is not None:
+ if self._exhausted:
+ self._raw_batches = iter(())
+ elif self._parquet_file is not None:
self._raw_batches = self._iter_row_group_batches()
else:
reader = self.dataset.scanner(
@@ -192,6 +247,8 @@ class FormatPyArrowReader(RecordBatchReader):
return columns
def _select_existing_fields(self, batch):
+ if not self.existing_fields:
+ return _zero_column_batch(batch.num_rows)
columns = []
fields = []
for name in self.existing_fields:
@@ -218,11 +275,15 @@ class FormatPyArrowReader(RecordBatchReader):
column = column.flatten()[index]
columns.append(column)
names.append(field.name)
+ if not columns:
+ return _zero_column_batch(batch.num_rows)
return pa.RecordBatch.from_arrays(columns, names=names)
def _surviving_row_group_ids(self):
total = self._parquet_file.num_row_groups
if self._scan_filter is None:
+ if self._selected_parquet_row_groups is not None:
+ return self._selected_parquet_row_groups
return range(total)
try:
ids = set()
@@ -236,7 +297,10 @@ class FormatPyArrowReader(RecordBatchReader):
return range(total)
def read_arrow_batch(self) -> Optional[RecordBatch]:
- batch = next(self._raw_batches, None)
+ if self._range_slicer is not None:
+ batch = self._range_slicer.next_batch(self._raw_batches)
+ else:
+ batch = next(self._raw_batches, None)
if batch is None:
return None
return self._post_process_batch(batch)
@@ -349,6 +413,124 @@ def _path_exists_in_arrow_schema(schema: pa.Schema, path:
List[str]) -> bool:
return True
+def _zero_column_batch(num_rows: int) -> RecordBatch:
+ """Build a zero-column batch without losing its logical row count."""
+ empty_struct = pa.Array.from_buffers(
+ pa.struct([]), num_rows, [None], children=[])
+ return pa.RecordBatch.from_struct_array(empty_struct)
+
+
+def _to_runs(row_indices: List[int]) -> List[Tuple[int, int]]:
+ """Collapse row indices into sorted, distinct, inclusive runs."""
+ if not row_indices:
+ return []
+ sorted_indices = sorted(set(row_indices))
+ runs = []
+ start = previous = sorted_indices[0]
+ for index in sorted_indices[1:]:
+ if index == previous + 1:
+ previous = index
+ continue
+ runs.append((start, previous))
+ start = previous = index
+ runs.append((start, previous))
+ return runs
+
+
+def _normalize_runs(
+ row_ranges: List[Tuple[int, int]]) -> List[Tuple[int, int]]:
+ """Sort and merge inclusive row ranges without expanding their rows."""
+ if not row_ranges:
+ return []
+ ranges = sorted(row_ranges)
+ merged = []
+ for lower, upper in ranges:
+ if lower > upper:
+ raise ValueError(
+ "Invalid row range: {} > {}".format(lower, upper))
+ if merged and lower <= merged[-1][1] + 1:
+ merged[-1] = (merged[-1][0], max(merged[-1][1], upper))
+ else:
+ merged.append((lower, upper))
+ return merged
+
+
+class _RowRunSlicer:
+ """Slice selected Parquet row groups down to requested file-local rows."""
+
+ def __init__(
+ self,
+ selected_infos: List[Tuple[int, int]],
+ runs: List[Tuple[int, int]]):
+ self._segments = []
+ concatenated_offset = 0
+ for file_offset, row_count in selected_infos:
+ self._segments.append((
+ concatenated_offset,
+ concatenated_offset + row_count,
+ file_offset,
+ ))
+ concatenated_offset += row_count
+ self._runs = [(lower, upper + 1) for lower, upper in runs]
+ self._stream_offset = 0
+ self._segment_index = 0
+ self._run_index = 0
+ self._pending: Deque[RecordBatch] = deque()
+
+ def next_batch(
+ self, batches: Iterator[RecordBatch]) -> Optional[RecordBatch]:
+ while not self._pending:
+ batch = next(batches, None)
+ if batch is None:
+ return None
+ self._slice_batch(batch)
+ return self._pending.popleft()
+
+ def _slice_batch(self, batch: RecordBatch) -> None:
+ batch_start = self._stream_offset
+ batch_end = batch_start + batch.num_rows
+ self._stream_offset = batch_end
+ position = batch_start
+
+ while position < batch_end:
+ while (
+ self._segment_index < len(self._segments)
+ and position >= self._segments[
+ self._segment_index][1]):
+ self._segment_index += 1
+ if self._segment_index >= len(self._segments):
+ return
+
+ segment_start, segment_end, file_start = self._segments[
+ self._segment_index]
+ part_end = min(batch_end, segment_end)
+ local_start = file_start + position - segment_start
+ local_end = file_start + part_end - segment_start
+
+ while (
+ self._run_index < len(self._runs)
+ and self._runs[self._run_index][1] <= local_start):
+ self._run_index += 1
+ run_index = self._run_index
+ while (
+ run_index < len(self._runs)
+ and self._runs[run_index][0] < local_end):
+ run_start, run_end = self._runs[run_index]
+ lower = max(local_start, run_start)
+ upper = min(local_end, run_end)
+ if lower < upper:
+ offset = (
+ position - batch_start + lower - local_start)
+ self._pending.append(
+ batch.slice(offset, upper - lower))
+ if run_end <= local_end:
+ run_index += 1
+ else:
+ break
+ self._run_index = run_index
+ position = part_end
+
+
def _contains_variant(data_type) -> bool:
if isinstance(data_type, AtomicType):
return data_type.type.upper() == 'VARIANT'
diff --git a/paimon-python/pypaimon/read/split_read.py
b/paimon-python/pypaimon/read/split_read.py
index e3d7f3f16d..d34c1800e0 100644
--- a/paimon-python/pypaimon/read/split_read.py
+++ b/paimon-python/pypaimon/read/split_read.py
@@ -225,19 +225,37 @@ class SplitRead(ABC):
file_path = self._aligned_extra_file_path(file, row_sidecar_file)
file_format = ROW_SIDECAR_FORMAT
- # Convert global row_ranges (IndexedSplit) to local row_indices for
native pushdown.
+ # Prepare file-local native row selection. Existing native formats
+ # consume row indices; Parquet keeps compact ranges to avoid expanding
+ # large selections into millions of Python integers.
row_indices = None
+ parquet_row_ranges = None
if effective_row_ranges is not None:
row_index_formats = (CoreOptions.FILE_FORMAT_BLOB,
CoreOptions.FILE_FORMAT_VORTEX,
CoreOptions.FILE_FORMAT_LANCE,
CoreOptions.FILE_FORMAT_ROW)
if file_format in row_index_formats:
- row_indices = []
- for r in effective_row_ranges:
- start = r.from_ - file.first_row_id
- end = r.to - file.first_row_id
- row_indices.extend(range(start, end + 1))
+ row_indices = [
+ row_id - file.first_row_id
+ for row_range in effective_row_ranges
+ for row_id in range(row_range.from_, row_range.to + 1)
+ ]
+ elif (file_format == CoreOptions.FILE_FORMAT_PARQUET
+ and read_arrow_predicate is None):
+ parquet_row_ranges = []
+ merged_ranges = Range.sort_and_merge_overlap(
+ effective_row_ranges, True)
+ for r in merged_ranges:
+ start = max(0, r.from_ - file.first_row_id)
+ end = min(
+ file.row_count - 1,
+ r.to - file.first_row_id,
+ )
+ if end >= start:
+ parquet_row_ranges.append((start, end))
+ if not parquet_row_ranges:
+ return EmptyRecordBatchReader()
# Map nested paths into the order the format reader will see.
nested_path_by_name = self._nested_path_by_name()
@@ -338,7 +356,8 @@ class SplitRead(ABC):
ordered_read_fields, read_arrow_predicate,
batch_size=batch_size,
options=self.table.options,
nested_name_paths=ordered_nested_paths,
- predicate_field_names=predicate_fields)
+ predicate_field_names=predicate_fields,
+ row_ranges=parquet_row_ranges)
elif file_format == CoreOptions.FILE_FORMAT_ROW:
if has_nested:
raise NotImplementedError(
@@ -394,6 +413,7 @@ class SplitRead(ABC):
system_fields,
file_io=self.table.file_io,
row_id_offsets=row_indices,
+ row_id_offset_ranges=parquet_row_ranges,
file_data_fields=file_read_fields,
target_data_fields=target_fields)
else:
@@ -409,11 +429,14 @@ class SplitRead(ABC):
system_fields,
file_io=self.table.file_io,
row_id_offsets=row_indices,
+ row_id_offset_ranges=parquet_row_ranges,
file_data_fields=file_read_fields,
target_data_fields=target_fields)
# For non-Vortex formats, wrap with RowIdFilterRecordBatchReader
- if row_ranges is not None and row_indices is None:
+ if (row_ranges is not None
+ and row_indices is None
+ and parquet_row_ranges is None):
reader = RowIdFilterRecordBatchReader(reader, file.first_row_id,
effective_row_ranges)
# For formats without native shard support, wrap with ShardBatchReader
diff --git a/paimon-python/pypaimon/tests/parquet_row_range_test.py
b/paimon-python/pypaimon/tests/parquet_row_range_test.py
new file mode 100644
index 0000000000..6bba358924
--- /dev/null
+++ b/paimon-python/pypaimon/tests/parquet_row_range_test.py
@@ -0,0 +1,393 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+################################################################################
+
+import os
+import shutil
+import tempfile
+import unittest
+from unittest import mock
+
+import pyarrow as pa
+import pyarrow.fs as pafs
+import pyarrow.parquet as pq
+
+from pypaimon import CatalogFactory, Schema
+from pypaimon.read.reader.format_pyarrow_reader import FormatPyArrowReader
+from pypaimon.schema.data_types import AtomicType, DataField
+
+
+N = 300
+ROW_GROUP_SIZE = 64
+TABLE_OPTIONS = {
+ "row-tracking.enabled": "true",
+ "data-evolution.enabled": "true",
+ "read.batch-size": "50",
+}
+VARIANT_TYPE = pa.struct(
+ [
+ pa.field("value", pa.binary(), nullable=False),
+ pa.field("metadata", pa.binary(), nullable=False),
+ ]
+)
+
+
+class _LocalFileIO:
+ filesystem = pafs.LocalFileSystem()
+
+ def to_filesystem_path(self, path):
+ return path
+
+
+def _commit_with_row_groups(table, data):
+ original = table.file_io.write_parquet
+
+ def patched(path, arrow_table, **kwargs):
+ kwargs.setdefault("row_group_size", ROW_GROUP_SIZE)
+ return original(path, arrow_table, **kwargs)
+
+ table.file_io.write_parquet = patched
+ try:
+ write_builder = table.new_batch_write_builder()
+ writer = write_builder.new_write()
+ commit = write_builder.new_commit()
+ writer.write_arrow(data)
+ commit.commit(writer.prepare_commit())
+ writer.close()
+ commit.close()
+ finally:
+ table.file_io.write_parquet = original
+
+
+def _data_file_paths(table):
+ warehouse_root = table.table_path.replace("file://", "")
+ paths = []
+ for root, _, files in os.walk(warehouse_root):
+ parts = set(os.path.relpath(root, warehouse_root).split(os.sep))
+ if parts & {"manifest", "schema", "snapshot", "index"}:
+ continue
+ for file_name in files:
+ if file_name.endswith(".parquet"):
+ paths.append(os.path.join(root, file_name))
+ return sorted(paths)
+
+
+def _rows(schema, count):
+ return pa.table(
+ {
+ "group_id": ["g{}".format(index // 100)
+ for index in range(count)],
+ "payload": [bytes([index % 256]) * 64
+ for index in range(count)],
+ "value": list(range(count)),
+ },
+ schema=schema,
+ )
+
+
+class ToRunsTest(unittest.TestCase):
+
+ def test_runs_are_sorted_deduplicated_and_empty_safe(self):
+ from pypaimon.read.reader.format_pyarrow_reader import (
+ _normalize_runs,
+ _to_runs,
+ )
+
+ self.assertEqual(_to_runs([]), [])
+ self.assertEqual(_to_runs([1, 2, 3]), [(1, 3)])
+ self.assertEqual(
+ _to_runs([9, 4, 1, 3, 4]),
+ [(1, 1), (3, 4), (9, 9)],
+ )
+ self.assertEqual(_normalize_runs([]), [])
+ self.assertEqual(
+ _normalize_runs([
+ (900_000_000, 999_999_999),
+ (0, 900_000_001),
+ ]),
+ [(0, 999_999_999)],
+ )
+
+
+class RowIdRangeCursorTest(unittest.TestCase):
+
+ def test_generates_offsets_across_batches_and_ranges(self):
+ from pypaimon.read.reader.data_file_batch_reader import (
+ _RowIdRangeCursor,
+ )
+
+ cursor = _RowIdRangeCursor([(10, 12), (100, 102)])
+ self.assertEqual(cursor.take(2), [10, 11])
+ self.assertEqual(cursor.take(3), [12, 100, 101])
+ self.assertEqual(cursor.take(1), [102])
+
+
+class ParquetRowRangeTest(unittest.TestCase):
+
+ @classmethod
+ def setUpClass(cls):
+ cls.tempdir = tempfile.mkdtemp()
+ cls.catalog = CatalogFactory.create(
+ {"warehouse": os.path.join(cls.tempdir, "warehouse")})
+ cls.catalog.create_database("default", False)
+ cls.arrow_schema = pa.schema(
+ [
+ ("group_id", pa.string()),
+ ("payload", pa.large_binary()),
+ ("value", pa.int32()),
+ ]
+ )
+ identifier = "default.parquet_row_range"
+ cls.catalog.create_table(
+ identifier,
+ Schema.from_pyarrow_schema(
+ cls.arrow_schema,
+ options=TABLE_OPTIONS,
+ ),
+ False,
+ )
+ table = cls.catalog.get_table(identifier)
+ _commit_with_row_groups(table, _rows(cls.arrow_schema, N))
+ cls.table = cls.catalog.get_table(identifier)
+
+ data_files = _data_file_paths(cls.table)
+ assert len(data_files) == 1, data_files
+ cls.data_file = data_files[0]
+ assert pq.ParquetFile(cls.data_file).num_row_groups == 5
+
+ @classmethod
+ def tearDownClass(cls):
+ shutil.rmtree(cls.tempdir, ignore_errors=True)
+
+ def _read(
+ self,
+ predicate_factory=None,
+ columns=("group_id", "value", "_ROW_ID")):
+ read_builder = self.table.new_read_builder().with_projection(
+ list(columns))
+ if predicate_factory is not None:
+ predicate_builder = read_builder.new_predicate_builder()
+ read_builder = read_builder.with_filter(
+ predicate_factory(predicate_builder))
+ splits = read_builder.new_scan().plan().splits()
+ return read_builder.new_read().to_arrow(splits)
+
+ def test_reads_only_intersecting_row_groups(self):
+ result = self._read(
+ lambda builder: builder.between("_ROW_ID", 100, 163))
+ self.assertEqual(
+ result.column("_ROW_ID").to_pylist(),
+ list(range(100, 164)),
+ )
+
+ reader = FormatPyArrowReader(
+ _LocalFileIO(),
+ "parquet",
+ self.data_file,
+ [DataField(0, "value", AtomicType("INT"))],
+ None,
+ batch_size=50,
+ row_ranges=[(100, 163)],
+ )
+ self.assertIsNotNone(reader._parquet_file)
+ self.assertEqual(
+ reader._selected_parquet_row_groups,
+ [1, 2],
+ )
+
+ def test_direct_row_indices_are_sorted_and_deduplicated(self):
+ reader = FormatPyArrowReader(
+ _LocalFileIO(),
+ "parquet",
+ self.data_file,
+ [DataField(0, "value", AtomicType("INT"))],
+ None,
+ batch_size=2,
+ row_indices=[130, 11, 10, 10],
+ )
+ values = []
+ while True:
+ batch = reader.read_arrow_batch()
+ if batch is None:
+ break
+ values.extend(batch.column(0).to_pylist())
+ self.assertEqual(values, [10, 11, 130])
+
+ def test_system_only_projection_preserves_rows(self):
+ result = self._read(
+ lambda builder: builder.between("_ROW_ID", 2, 7),
+ columns=("_ROW_ID",),
+ )
+ self.assertEqual(
+ result.column("_ROW_ID").to_pylist(),
+ list(range(2, 8)),
+ )
+
+ def test_missing_only_projection_preserves_rows(self):
+ reader = FormatPyArrowReader(
+ _LocalFileIO(),
+ "parquet",
+ self.data_file,
+ [DataField(99, "missing", AtomicType("INT"))],
+ None,
+ batch_size=2,
+ row_ranges=[(2, 7)],
+ )
+ values = []
+ while True:
+ batch = reader.read_arrow_batch()
+ if batch is None:
+ break
+ values.extend(batch.column(0).to_pylist())
+ self.assertEqual(values, [None] * 6)
+
+ def test_skips_per_row_python_range_filter(self):
+ from pypaimon.read.reader.row_range_filter_record_reader import (
+ RowIdFilterRecordBatchReader,
+ )
+
+ with mock.patch.object(
+ RowIdFilterRecordBatchReader,
+ "_is_row_in_range",
+ wraps=RowIdFilterRecordBatchReader._is_row_in_range,
+ ) as range_spy:
+ result = self._read(
+ lambda builder: builder.between("_ROW_ID", 100, 163))
+ self.assertEqual(result.num_rows, 64)
+ range_spy.assert_not_called()
+
+ def test_disjoint_ranges_preserve_order_and_payload(self):
+ result = self._read(
+ lambda builder: builder.or_predicates(
+ [
+ builder.between("_ROW_ID", 10, 20),
+ builder.between("_ROW_ID", 130, 140),
+ builder.between("_ROW_ID", 290, 299),
+ ]
+ ),
+ columns=("payload", "value", "_ROW_ID"),
+ )
+ expected = (
+ list(range(10, 21))
+ + list(range(130, 141))
+ + list(range(290, 300))
+ )
+ self.assertEqual(
+ result.column("_ROW_ID").to_pylist(), expected)
+ self.assertEqual(result.column("value").to_pylist(), expected)
+ self.assertEqual(
+ result.column(
+ result.schema.get_field_index("payload")).to_pylist(),
+ [bytes([index % 256]) * 64 for index in expected],
+ )
+
+ def test_column_predicate_falls_back_without_position_shift(self):
+ result = self._read(
+ lambda builder: builder.and_predicates(
+ [
+ builder.between("_ROW_ID", 64, 250),
+ builder.greater_or_equal("value", 200),
+ ]
+ )
+ )
+ self.assertEqual(
+ result.column("_ROW_ID").to_pylist(),
+ list(range(200, 251)),
+ )
+ self.assertEqual(
+ result.column("value").to_pylist(),
+ list(range(200, 251)),
+ )
+
+ result = self._read(
+ lambda builder: builder.and_predicates(
+ [
+ builder.between("_ROW_ID", 90, 210),
+ builder.equal("group_id", "g1"),
+ ]
+ )
+ )
+ self.assertEqual(
+ result.column("_ROW_ID").to_pylist(),
+ list(range(100, 200)),
+ )
+
+ def test_full_scan_is_unchanged(self):
+ real_parquet_file = pq.ParquetFile
+ with mock.patch.object(
+ pq,
+ "ParquetFile",
+ wraps=real_parquet_file,
+ ) as parquet_file_spy:
+ result = self._read(columns=("value",))
+ self.assertEqual(
+ sorted(result.column("value").to_pylist()),
+ list(range(N)),
+ )
+ parquet_file_spy.assert_not_called()
+
+ def test_projected_variant_uses_same_exact_row_slicing(self):
+ path = os.path.join(self.tempdir, "variant-row-ranges.parquet")
+ payload = pa.array(
+ [
+ {"value": "v{}".format(index).encode(), "metadata": b"m"}
+ for index in range(N)
+ ],
+ type=VARIANT_TYPE,
+ )
+ pq.write_table(
+ pa.table({"payload": payload}),
+ path,
+ row_group_size=ROW_GROUP_SIZE,
+ )
+ requested = list(range(60, 71)) + list(range(130, 134))
+ reader = FormatPyArrowReader(
+ _LocalFileIO(),
+ "parquet",
+ path,
+ [DataField(0, "payload", AtomicType("VARIANT"))],
+ None,
+ batch_size=50,
+ row_ranges=[(60, 70), (130, 133)],
+ )
+ self.assertIsNotNone(reader._parquet_file)
+ self.assertEqual(
+ reader._selected_parquet_row_groups,
+ [0, 1, 2],
+ )
+ values = []
+ while True:
+ batch = reader.read_arrow_batch()
+ if batch is None:
+ break
+ values.extend(
+ batch.column(
+ batch.schema.get_field_index("payload")).to_pylist())
+ self.assertEqual(
+ values,
+ [
+ {
+ "value": "v{}".format(index).encode(),
+ "metadata": b"m",
+ }
+ for index in requested
+ ],
+ )
+
+
+if __name__ == "__main__":
+ unittest.main()
