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 25a3b5a96b [python] Support batch size for blob read (#8502)
25a3b5a96b is described below
commit 25a3b5a96b238ed938c6c82142d472b54f2ecdbc
Author: umi <[email protected]>
AuthorDate: Fri Jul 10 08:36:35 2026 +0800
[python] Support batch size for blob read (#8502)
---
.../pypaimon/read/reader/concat_batch_reader.py | 229 ++++++++++++++++-----
paimon-python/pypaimon/read/split_read.py | 1 +
paimon-python/pypaimon/tests/blob_test.py | 153 ++++++++++++++
.../tests/data_evolution_deletion_vector_test.py | 137 +++++++++++-
4 files changed, 459 insertions(+), 61 deletions(-)
diff --git a/paimon-python/pypaimon/read/reader/concat_batch_reader.py
b/paimon-python/pypaimon/read/reader/concat_batch_reader.py
index 51b6f573a6..f17092279c 100644
--- a/paimon-python/pypaimon/read/reader/concat_batch_reader.py
+++ b/paimon-python/pypaimon/read/reader/concat_batch_reader.py
@@ -31,6 +31,24 @@ from pypaimon.utils.range import Range
_MIN_BATCH_SIZE_TO_REFILL = 1024
+class _BlobFileState:
+
+ def __init__(
+ self,
+ file: DataFileMeta,
+ supplier: Callable,
+ selected_ranges: List[Range],
+ ):
+ self.file = file
+ self.supplier = supplier
+ self.selected_ranges = selected_ranges
+ self.selected_count = sum(row_range.count() for row_range in
selected_ranges)
+ self.reader = None
+ self.reader_initialized = False
+ self.selected_range_index = 0
+ self.selected_position_base = 0
+
+
class ConcatBatchReader(RecordBatchReader):
def __init__(self, reader_suppliers: List[Callable], file_io=None,
@@ -241,7 +259,7 @@ class BlobFallbackBatchReader(RecordBatchReader):
def __init__(self, file_reader_suppliers: List[Tuple[DataFileMeta,
Callable]],
field_name: str, output_type, row_ranges:
Optional[List[Range]] = None,
- blob_as_descriptor: bool = False, deletion_vector=None):
+ blob_as_descriptor: bool = False, deletion_vector=None,
batch_size: int = 1024):
self._file_reader_suppliers = file_reader_suppliers
self._field_name = field_name
self._output_type = output_type
@@ -253,28 +271,37 @@ class BlobFallbackBatchReader(RecordBatchReader):
else:
self._deletion_vector_range, self._deletion_vector =
deletion_vector
self._returned = False
- self._readers: List[RecordBatchReader] = []
+ self._batch_size = max(1, batch_size)
+ self._target_ranges = self._compute_target_ranges()
+ self._target_range_index = 0
+ self._next_row_id = (
+ self._target_ranges[0].from_
+ if self._target_ranges
+ else None
+ )
+ self._file_states = []
+ for file, supplier in self._file_reader_suppliers:
+ selected_ranges = self._selected_ranges(file)
+ if selected_ranges:
+ self._file_states.append(_BlobFileState(file, supplier,
selected_ranges))
def read_arrow_batch(self) -> Optional[RecordBatch]:
- if self._returned:
+ batch_row_ids = self._next_batch_row_ids()
+ if not batch_row_ids:
return None
- self._returned = True
groups: Dict[int, Dict[int, Tuple[object, bool]]] = {}
- target_row_ids = self._target_row_ids()
- for file, supplier in self._file_reader_suppliers:
- row_ids = self._selected_row_ids(file)
- blob_values = self._read_blob_values(file, supplier)
- if len(blob_values) != len(row_ids):
- raise ValueError(
- "Blob fallback reader returned an unexpected row count "
- f"for {file.file_name}: expect {len(row_ids)}, got
{len(blob_values)}."
- )
- if not row_ids:
+ batch_first = batch_row_ids[0]
+ batch_last = batch_row_ids[-1]
+ for state in self._file_states:
+ if not self._state_overlaps_batch(state, batch_first, batch_last):
+ continue
+ blob_values = self._read_blob_values(state, batch_row_ids)
+ if not blob_values:
continue
- group = groups.setdefault(file.max_sequence_number, {})
- for row_id, blob in zip(row_ids, blob_values):
+ group = groups.setdefault(state.file.max_sequence_number, {})
+ for row_id, blob in blob_values.items():
if row_id in group:
raise ValueError(
"Blob files within the same max sequence should not
overlap."
@@ -293,7 +320,7 @@ class BlobFallbackBatchReader(RecordBatchReader):
return None
result = []
- for row_id in target_row_ids:
+ for row_id in batch_row_ids:
found = False
for max_sequence_number in sorted(groups.keys(), reverse=True):
candidate = groups[max_sequence_number].get(row_id)
@@ -312,38 +339,49 @@ class BlobFallbackBatchReader(RecordBatchReader):
names=[self._field_name],
)
- def _target_row_ids(self) -> List[int]:
- file_ranges = [
+ def _compute_target_ranges(self) -> List[Range]:
+ ranges = Range.sort_and_merge_overlap([
file.row_id_range()
for file, _ in self._file_reader_suppliers
- ]
- ranges = [
- Range(
- min(row_range.from_ for row_range in file_ranges),
- max(row_range.to for row_range in file_ranges),
- )
- ]
+ ])
if self._row_ranges is not None:
ranges = Range.and_(ranges, self._row_ranges)
- return [
- row_id
- for row_id in self._expand_ranges(ranges)
- if not self._is_deleted(row_id)
- ]
+ return ranges
- def _selected_row_ids(self, file: DataFileMeta) -> List[int]:
+ def _selected_ranges(self, file: DataFileMeta) -> List[Range]:
ranges = [file.row_id_range()]
if self._row_ranges is not None:
ranges = Range.and_(ranges, self._row_ranges)
- return self._expand_ranges(ranges)
+ return ranges
+
+ def _next_batch_row_ids(self) -> List[int]:
+ batch_row_ids = []
+ while (
+ len(batch_row_ids) < self._batch_size
+ and self._target_range_index < len(self._target_ranges)
+ ):
+ target_range = self._target_ranges[self._target_range_index]
+ if self._next_row_id is None or self._next_row_id <
target_range.from_:
+ self._next_row_id = target_range.from_
+
+ while (
+ self._next_row_id <= target_range.to
+ and len(batch_row_ids) < self._batch_size
+ ):
+ row_id = self._next_row_id
+ self._next_row_id += 1
+ if not self._is_deleted(row_id):
+ batch_row_ids.append(row_id)
+
+ if self._next_row_id > target_range.to:
+ self._target_range_index += 1
+ self._next_row_id = (
+ self._target_ranges[self._target_range_index].from_
+ if self._target_range_index < len(self._target_ranges)
+ else None
+ )
- @staticmethod
- def _expand_ranges(ranges: List[Range]) -> List[int]:
- return [
- row_id
- for row_range in ranges
- for row_id in range(row_range.from_, row_range.to + 1)
- ]
+ return batch_row_ids
def _is_deleted(self, row_id: int) -> bool:
if self._deletion_vector is None:
@@ -357,25 +395,116 @@ class BlobFallbackBatchReader(RecordBatchReader):
row_id - self._deletion_vector_range.from_
)
- def _read_blob_values(self, file: DataFileMeta, supplier: Callable) ->
List[object]:
- reader = supplier()
+ @staticmethod
+ def _state_overlaps_batch(
+ state: _BlobFileState, batch_first: int, batch_last: int
+ ) -> bool:
+ selected_ranges = state.selected_ranges
+ while (
+ state.selected_range_index < len(selected_ranges)
+ and selected_ranges[state.selected_range_index].to < batch_first
+ ):
+ state.selected_position_base += (
+ selected_ranges[state.selected_range_index].count()
+ )
+ state.selected_range_index += 1
+ return (
+ state.selected_range_index < len(selected_ranges)
+ and selected_ranges[state.selected_range_index].from_ <= batch_last
+ )
+
+ def _read_blob_values(
+ self, state: _BlobFileState, batch_row_ids: List[int]
+ ) -> Dict[int, object]:
+ positions_and_row_ids = self._selected_positions_and_row_ids(
+ state, batch_row_ids
+ )
+ if not positions_and_row_ids:
+ return {}
+
+ reader = self._reader_for_state(state)
if reader is None:
- return []
- self._readers.append(reader)
+ return {}
+
try:
+ blob_lengths = [reader.blob_lengths[pos] for pos, _ in
positions_and_row_ids]
+ blob_offsets = [reader.blob_offsets[pos] for pos, _ in
positions_and_row_ids]
iterator = BlobRecordIterator(
reader._file_io,
reader.file_path,
- reader.blob_lengths,
- reader.blob_offsets,
+ blob_lengths,
+ blob_offsets,
self._field_name,
reader._input_stream,
)
- return [row.values[0] for row in iterator]
+
+ blobs = []
+ for row in iterator:
+ blobs.append(row.values[0])
+ return {
+ row_id: blob
+ for (_, row_id), blob in zip(positions_and_row_ids, blobs)
+ }
except AttributeError as e:
raise TypeError("Blob fallback reader expects FormatBlobReader
suppliers.") from e
- def close(self) -> None:
- for reader in self._readers:
+ @staticmethod
+ def _selected_positions_and_row_ids(
+ state: _BlobFileState, batch_row_ids: List[int]
+ ) -> List[Tuple[int, int]]:
+ selected_ranges = state.selected_ranges
+ positions_and_row_ids = []
+ for row_id in batch_row_ids:
+ while (
+ state.selected_range_index < len(selected_ranges)
+ and selected_ranges[state.selected_range_index].to < row_id
+ ):
+ state.selected_position_base += (
+ selected_ranges[state.selected_range_index].count()
+ )
+ state.selected_range_index += 1
+ if state.selected_range_index >= len(selected_ranges):
+ break
+ row_range = selected_ranges[state.selected_range_index]
+ if row_range.from_ <= row_id <= row_range.to:
+ positions_and_row_ids.append(
+ (
+ state.selected_position_base + row_id -
row_range.from_,
+ row_id,
+ )
+ )
+ return positions_and_row_ids
+
+ def _reader_for_state(self, state: _BlobFileState):
+ if state.reader_initialized:
+ return state.reader
+
+ reader = state.supplier()
+ if reader is None:
+ state.reader_initialized = True
+ return None
+ actual_rows = len(reader.blob_lengths)
+ expected_rows = state.selected_count
+ if actual_rows != expected_rows:
+ reader.close()
+ raise ValueError(
+ "Blob fallback reader returned an unexpected row count "
+ f"for {state.file.file_name}: expect {expected_rows}, got "
+ f"{actual_rows}."
+ )
+
+ state.reader = reader
+ state.reader_initialized = True
+ return reader
+
+ @staticmethod
+ def _close_state_reader(state: _BlobFileState) -> None:
+ reader = state.reader
+ state.reader = None
+ state.reader_initialized = False
+ if reader is not None:
reader.close()
- self._readers = []
+
+ def close(self) -> None:
+ for state in self._file_states:
+ self._close_state_reader(state)
diff --git a/paimon-python/pypaimon/read/split_read.py
b/paimon-python/pypaimon/read/split_read.py
index d230975d5c..0b01ae4507 100644
--- a/paimon-python/pypaimon/read/split_read.py
+++ b/paimon-python/pypaimon/read/split_read.py
@@ -1283,6 +1283,7 @@ class DataEvolutionSplitRead(SplitRead):
self.row_ranges,
CoreOptions.blob_as_descriptor(self.table.options),
deletion_vector=deletion_vector,
+ batch_size=batch_size,
)
else:
# Create concatenated reader for multiple files
diff --git a/paimon-python/pypaimon/tests/blob_test.py
b/paimon-python/pypaimon/tests/blob_test.py
index 14508d1e10..cd79d5d432 100644
--- a/paimon-python/pypaimon/tests/blob_test.py
+++ b/paimon-python/pypaimon/tests/blob_test.py
@@ -28,7 +28,9 @@ import pyarrow as pa
from pypaimon import CatalogFactory, Schema
from pypaimon.common.file_io import FileIO
from pypaimon.filesystem.local_file_io import LocalFileIO
+from pypaimon.manifest.schema.data_file_meta import DataFileMeta
from pypaimon.common.options import Options
+from pypaimon.read.reader.concat_batch_reader import BlobFallbackBatchReader
from pypaimon.read.reader.format_blob_reader import BlobRecordIterator,
FormatBlobReader
from pypaimon.schema.data_types import AtomicType, DataField
from pypaimon.table.row.blob import Blob, BlobData, BlobRef, BlobDescriptor,
BlobViewStruct, BlobView
@@ -246,6 +248,157 @@ class BlobTest(unittest.TestCase):
self.assertFalse(blob.is_resolved())
self.assertEqual(blob.view_struct, view_struct)
+ def test_blob_fallback_batch_reader_respects_batch_size(self):
+ created_readers = []
+
+ class FakeBlobReader:
+ def __init__(self):
+ self._file_io = None
+ self.file_path = "fake.blob"
+ self.blob_lengths = [20, 20, 20, 20, 20]
+ self.blob_offsets = [0, 100, 200, 300, 400]
+ self._input_stream = None
+ self.closed = False
+
+ def close(self):
+ self.closed = True
+
+ def supplier():
+ reader = FakeBlobReader()
+ created_readers.append(reader)
+ return reader
+
+ data_file = DataFileMeta(
+ file_name="fake.blob",
+ file_size=0,
+ row_count=5,
+ min_key=None,
+ max_key=None,
+ key_stats=None,
+ value_stats=None,
+ min_sequence_number=0,
+ max_sequence_number=0,
+ schema_id=0,
+ level=0,
+ extra_files=[],
+ first_row_id=10,
+ file_path="fake.blob",
+ )
+ reader = BlobFallbackBatchReader(
+ [(data_file, supplier)],
+ "picture",
+ pa.large_binary(),
+ blob_as_descriptor=True,
+ batch_size=2,
+ )
+
+ first = reader.read_arrow_batch()
+ second = reader.read_arrow_batch()
+ third = reader.read_arrow_batch()
+ self.assertIsNone(reader.read_arrow_batch())
+
+ self.assertEqual(first.num_rows, 2)
+ self.assertEqual(second.num_rows, 2)
+ self.assertEqual(third.num_rows, 1)
+ offsets = []
+ for batch in (first, second, third):
+ offsets.extend(
+ BlobDescriptor.deserialize(value.as_py()).offset
+ for value in batch.column("picture")
+ )
+ self.assertEqual(offsets, [4, 104, 204, 304, 404])
+ self.assertEqual(1, len(created_readers))
+ self.assertFalse(created_readers[0].closed)
+
+ reader.close()
+ self.assertTrue(created_readers[0].closed)
+
+ def test_blob_fallback_batch_reader_reuses_version_readers(self):
+ created_by_file = {}
+
+ class FakeBlobReader:
+ def __init__(self, file_path, blob_lengths, blob_offsets):
+ self._file_io = None
+ self.file_path = file_path
+ self.blob_lengths = blob_lengths
+ self.blob_offsets = blob_offsets
+ self._input_stream = None
+ self.closed = False
+
+ def close(self):
+ self.closed = True
+
+ def data_file(name, max_sequence_number):
+ return DataFileMeta(
+ file_name=name,
+ file_size=0,
+ row_count=5,
+ min_key=None,
+ max_key=None,
+ key_stats=None,
+ value_stats=None,
+ min_sequence_number=max_sequence_number,
+ max_sequence_number=max_sequence_number,
+ schema_id=0,
+ level=0,
+ extra_files=[],
+ first_row_id=0,
+ file_path=name,
+ )
+
+ def supplier(file_path, blob_lengths, blob_offsets):
+ def create_reader():
+ reader = FakeBlobReader(file_path, blob_lengths, blob_offsets)
+ created_by_file.setdefault(file_path, []).append(reader)
+ return reader
+ return create_reader
+
+ old_file = data_file("old.blob", 1)
+ new_file = data_file("new.blob", 2)
+ reader = BlobFallbackBatchReader(
+ [
+ (
+ old_file,
+ supplier(
+ "old.blob",
+ [20, 20, 20, 20, 20],
+ [0, 100, 200, 300, 400],
+ ),
+ ),
+ (
+ new_file,
+ supplier(
+ "new.blob",
+ [-2, 20, -2, 20, -2],
+ [-1, 1000, -1, 3000, -1],
+ ),
+ ),
+ ],
+ "picture",
+ pa.large_binary(),
+ blob_as_descriptor=True,
+ batch_size=2,
+ )
+
+ offsets = []
+ batch = reader.read_arrow_batch()
+ while batch is not None:
+ offsets.extend(
+ BlobDescriptor.deserialize(value.as_py()).offset
+ for value in batch.column("picture")
+ )
+ batch = reader.read_arrow_batch()
+
+ self.assertEqual([4, 1004, 204, 3004, 404], offsets)
+ self.assertEqual(1, len(created_by_file["old.blob"]))
+ self.assertEqual(1, len(created_by_file["new.blob"]))
+ self.assertFalse(created_by_file["old.blob"][0].closed)
+ self.assertFalse(created_by_file["new.blob"][0].closed)
+
+ reader.close()
+ self.assertTrue(created_by_file["old.blob"][0].closed)
+ self.assertTrue(created_by_file["new.blob"][0].closed)
+
def test_blob_data_interface_compliance(self):
"""Test that BlobData properly implements Blob interface."""
test_data = b"interface test data"
diff --git
a/paimon-python/pypaimon/tests/data_evolution_deletion_vector_test.py
b/paimon-python/pypaimon/tests/data_evolution_deletion_vector_test.py
index 650adb95f2..448a8ab1b0 100644
--- a/paimon-python/pypaimon/tests/data_evolution_deletion_vector_test.py
+++ b/paimon-python/pypaimon/tests/data_evolution_deletion_vector_test.py
@@ -57,7 +57,8 @@ class _OneBatchReader(RecordBatchReader):
class _BlobFallbackBatchReaderForTest(BlobFallbackBatchReader):
def __init__(
- self, files, values_by_file_name, row_ranges=None, deletion_vector=None
+ self, files, values_by_file_name, row_ranges=None,
deletion_vector=None,
+ batch_size=1024
):
super().__init__(
[(file, lambda: None) for file in files],
@@ -66,15 +67,18 @@ class
_BlobFallbackBatchReaderForTest(BlobFallbackBatchReader):
row_ranges=row_ranges,
blob_as_descriptor=False,
deletion_vector=deletion_vector,
+ batch_size=batch_size,
)
self._values_by_file_name = values_by_file_name
- def _read_blob_values(self, file, supplier):
- values = self._values_by_file_name[file.file_name]
- return [
- values[row_id - file.first_row_id]
- for row_id in self._selected_row_ids(file)
- ]
+ def _read_blob_values(self, state, batch_row_ids):
+ values = self._values_by_file_name[state.file.file_name]
+ return {
+ row_id: values[pos]
+ for pos, row_id in self._selected_positions_and_row_ids(
+ state, batch_row_ids
+ )
+ }
def _file(name, first_row_id, row_count, max_sequence_number):
@@ -213,6 +217,121 @@ class DataEvolutionDeletionVectorTest(unittest.TestCase):
)
self.assertIsNone(reader.read_arrow_batch())
+ def test_blob_fallback_batch_reader_does_not_eof_on_row_id_gap(self):
+ reader = _BlobFallbackBatchReaderForTest(
+ [
+ _file("blob-left.blob", 0, 1, 1),
+ _file("blob-right.blob", 2, 1, 1),
+ ],
+ {
+ "blob-left.blob": [BlobData(b"left")],
+ "blob-right.blob": [BlobData(b"right")],
+ },
+ row_ranges=[Range(0, 2)],
+ batch_size=1,
+ )
+
+ values = []
+ for batch in iter(reader.read_arrow_batch, None):
+ values.extend(batch.column(0).to_pylist())
+
+ self.assertEqual([b"left", b"right"], values)
+
+ def test_blob_fallback_batch_reader_skips_states_outside_batch(self):
+ class
_CountingBlobFallbackBatchReader(_BlobFallbackBatchReaderForTest):
+ def __init__(self, files, values_by_file_name, batch_size):
+ super().__init__(files, values_by_file_name,
batch_size=batch_size)
+ self.selected_calls = []
+
+ def _selected_positions_and_row_ids(self, state, batch_row_ids):
+ self.selected_calls.append(
+ (state.file.file_name, list(batch_row_ids))
+ )
+ return BlobFallbackBatchReader._selected_positions_and_row_ids(
+ state, batch_row_ids
+ )
+
+ reader = _CountingBlobFallbackBatchReader(
+ [
+ _file("blob-0.blob", 0, 1, 1),
+ _file("blob-10.blob", 10, 1, 1),
+ _file("blob-20.blob", 20, 1, 1),
+ ],
+ {
+ "blob-0.blob": [BlobData(b"0")],
+ "blob-10.blob": [BlobData(b"10")],
+ "blob-20.blob": [BlobData(b"20")],
+ },
+ batch_size=1,
+ )
+ try:
+ self.assertEqual([b"0"],
reader.read_arrow_batch().column(0).to_pylist())
+ self.assertEqual([("blob-0.blob", [0])], reader.selected_calls)
+
+ self.assertEqual([b"10"],
reader.read_arrow_batch().column(0).to_pylist())
+ self.assertEqual(
+ [("blob-0.blob", [0]), ("blob-10.blob", [10])],
+ reader.selected_calls,
+ )
+ finally:
+ reader.close()
+
+ def
test_blob_fallback_batch_reader_avoids_full_row_id_materialization(self):
+ class _LazyBlobFallbackBatchReaderForTest(BlobFallbackBatchReader):
+ def _read_blob_values(self, state, batch_row_ids):
+ return {
+ row_id: BlobData(str(row_id).encode("utf-8"))
+ for row_id in batch_row_ids
+ }
+
+ reader = _LazyBlobFallbackBatchReaderForTest(
+ [(_file("large.blob", 0, 1_000_000, 1), lambda: None)],
+ "blob_col",
+ pa.binary(),
+ batch_size=2,
+ )
+ try:
+ self.assertNotIn("_target_row_ids", reader.__dict__)
+ for state in reader._file_states:
+ self.assertFalse(hasattr(state, "selected_row_ids"))
+ self.assertFalse(hasattr(state, "row_id_to_pos"))
+ self.assertFalse(hasattr(state,
"reader_uses_selected_positions"))
+
+ batch = reader.read_arrow_batch()
+ self.assertEqual([b"0", b"1"], batch.column(0).to_pylist())
+ finally:
+ reader.close()
+
+ def test_blob_fallback_batch_reader_advances_sparse_range_cursor(self):
+ reader = _BlobFallbackBatchReaderForTest(
+ [_file("sparse.blob", 0, 100, 1)],
+ {
+ "sparse.blob": [
+ BlobData(b"0"),
+ BlobData(b"10"),
+ BlobData(b"20"),
+ BlobData(b"30"),
+ ],
+ },
+ row_ranges=[Range(0, 0), Range(10, 10), Range(20, 20), Range(30,
30)],
+ batch_size=2,
+ )
+ try:
+ state = reader._file_states[0]
+
+ first = reader.read_arrow_batch()
+ self.assertEqual([b"0", b"10"], first.column(0).to_pylist())
+ self.assertEqual(1, state.selected_range_index)
+ self.assertEqual(1, state.selected_position_base)
+
+ second = reader.read_arrow_batch()
+ self.assertEqual([b"20", b"30"], second.column(0).to_pylist())
+ self.assertEqual(3, state.selected_range_index)
+ self.assertEqual(3, state.selected_position_base)
+ self.assertIsNone(reader.read_arrow_batch())
+ finally:
+ reader.close()
+
def
test_data_evolution_merge_reader_aligns_blob_with_row_ranges_and_dv(self):
row_ranges = [Range(1, 4)]
deletion_vector = BitmapDeletionVector()
@@ -234,20 +353,16 @@ class DataEvolutionDeletionVectorTest(unittest.TestCase):
],
{
"blob-old.blob": [
- BlobData(b"old-0"),
BlobData(b"old-1"),
BlobData(b"old-2"),
BlobData(b"old-3"),
BlobData(b"old-4"),
- BlobData(b"old-5"),
],
"blob-new.blob": [
- Blob.PLACE_HOLDER,
Blob.PLACE_HOLDER,
BlobData(b"new-2"),
BlobData(b"new-3"),
BlobData(b"new-4"),
- Blob.PLACE_HOLDER,
],
},
row_ranges=row_ranges,
