diff --git a/paimon-python/pypaimon/read/reader/concat_batch_reader.py b/paimon-python/pypaimon/read/reader/concat_batch_reader.py index 51b6f573a63a..f17092279c26 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 @@ _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 @@ def __init__(self, file_reader_suppliers: List[Tuple[DataFileMeta, Callable]], 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 @@ def read_arrow_batch(self) -> Optional[RecordBatch]: 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 @@ def read_arrow_batch(self) -> Optional[RecordBatch]: 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 @@ def _is_deleted(self, row_id: int) -> bool: 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 9167b7824abd..bf5fcfc9f9ce 100644 --- a/paimon-python/pypaimon/read/split_read.py +++ b/paimon-python/pypaimon/read/split_read.py @@ -1287,6 +1287,7 @@ def _create_union_reader(self, need_merge_files: List[DataFileMeta], deletion_ve 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 14508d1e10b9..cd79d5d432dd 100644 --- a/paimon-python/pypaimon/tests/blob_test.py +++ b/paimon-python/pypaimon/tests/blob_test.py @@ -28,7 +28,9 @@ 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 @@ def test_blob_view_struct_roundtrip(self): 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 650adb95f243..448a8ab1b0c1 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 @@ def close(self): 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 @@ def __init__( 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 @@ def test_blob_fallback_batch_reader_applies_deletion_vector(self): ) 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 @@ def test_data_evolution_merge_reader_aligns_blob_with_row_ranges_and_dv(self): ], { "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,