From 7fde333adddcd01aa71e6bfd53013591893cc6d4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=96=86=E5=AE=87?= Date: Sat, 25 Jul 2026 15:37:30 +0800 Subject: [PATCH 1/2] [python] chunk_shuffle should skip deleted records during counting --- .../scanner/chunk_shuffle_split_generator.py | 277 ++++++++++++++++-- .../pypaimon/read/scanner/file_scanner.py | 2 - paimon-python/pypaimon/read/split_read.py | 5 + .../data_evolution_deletion_vector_test.py | 36 +++ .../chunk_shuffle_split_generator_test.py | 244 ++++++++++++++- 5 files changed, 519 insertions(+), 45 deletions(-) diff --git a/paimon-python/pypaimon/read/scanner/chunk_shuffle_split_generator.py b/paimon-python/pypaimon/read/scanner/chunk_shuffle_split_generator.py index 6493b2e6fb3f..368f1edee5eb 100644 --- a/paimon-python/pypaimon/read/scanner/chunk_shuffle_split_generator.py +++ b/paimon-python/pypaimon/read/scanner/chunk_shuffle_split_generator.py @@ -19,8 +19,9 @@ from abc import abstractmethod from collections import defaultdict from dataclasses import dataclass -from typing import Any, List, Optional, Tuple +from typing import Any, Iterator, List, Optional, Tuple +from pypaimon.deletionvectors.deletion_vector import DeletionVector from pypaimon.globalindex.indexed_split import IndexedSplit from pypaimon.manifest.schema.data_file_meta import DataFileMeta from pypaimon.manifest.schema.manifest_entry import ManifestEntry @@ -28,6 +29,8 @@ from pypaimon.read.sliced_split import SlicedSplit from pypaimon.read.split import DataSplit, Split from pypaimon.table.row.generic_row import GenericRow +from pypaimon.table.source.deletion_file import DeletionFile +from pypaimon.utils.data_evolution_utils import retrieve_anchor_file from pypaimon.utils.range import Range from pypaimon.utils.range_helper import RangeHelper @@ -42,6 +45,115 @@ def _null_safe_partition_key(partition_values) -> tuple: return tuple((v is None, v) for v in partition_values) +@dataclass +class _LiveRowRange: + """A physical half-open range containing ``live_rows`` visible rows.""" + + start: int + end: int + live_rows: int + + +class _LiveRowRangeSlicer: + """Map requested live-row counts to contiguous physical row ranges. + + Deleted positions must be sorted and unique. Each position is consumed + once, so slicing costs O(number of output ranges + DV cardinality) + instead of O(physical row count). + """ + + def __init__( + self, + physical_row_count: int, + deleted_positions: Iterator[int], + ): + if physical_row_count < 0: + raise ValueError( + f"physical_row_count must be non-negative, got {physical_row_count}" + ) + self._physical_row_count = physical_row_count + self._deleted_positions = iter(deleted_positions) + self._physical_position = 0 + self._last_deleted_position = None + self._next_deleted_position = self._advance_deleted_position() + + def take(self, expected_live_rows: int) -> Optional[_LiveRowRange]: + if expected_live_rows <= 0: + raise ValueError( + f"expected_live_rows must be positive, got {expected_live_rows}" + ) + if self._physical_position >= self._physical_row_count: + return None + + start = self._physical_position + live_rows = 0 + + while self._physical_position < self._physical_row_count: + if self._next_deleted_position is None: + take = min( + expected_live_rows - live_rows, + self._physical_row_count - self._physical_position, + ) + self._physical_position += take + live_rows += take + else: + live_run = self._next_deleted_position - self._physical_position + needed = expected_live_rows - live_rows + if needed <= live_run: + self._physical_position += needed + live_rows += needed + else: + self._physical_position += live_run + live_rows += live_run + + if live_rows == expected_live_rows: + # Deleted rows have zero live-row weight. Attach a deletion run + # immediately after the boundary to this range so the next + # range starts at a live row (or EOF). + self._skip_deleted_positions_at_cursor() + return _LiveRowRange( + start, + self._physical_position, + live_rows, + ) + + # The current live run was insufficient, so the cursor must be at + # the next deleted position. Consume it and continue. + self._skip_deleted_positions_at_cursor() + + if live_rows == 0: + return None + return _LiveRowRange(start, self._physical_position, live_rows) + + def _skip_deleted_positions_at_cursor(self) -> None: + while ( + self._next_deleted_position is not None + and self._next_deleted_position == self._physical_position + ): + self._physical_position += 1 + self._next_deleted_position = self._advance_deleted_position() + + def _advance_deleted_position(self) -> Optional[int]: + position = next(self._deleted_positions, None) + if position is None: + return None + if position < 0 or position >= self._physical_row_count: + raise ValueError( + f"Deletion vector position {position} is outside physical row " + f"range [0, {self._physical_row_count})." + ) + if ( + self._last_deleted_position is not None + and position <= self._last_deleted_position + ): + raise ValueError( + "Deletion vector positions must be strictly increasing, but found " + f"{position} after {self._last_deleted_position}." + ) + self._last_deleted_position = position + return position + + @dataclass class _Chunk: """A unit of work for one DataLoader read. ``segments`` carries @@ -68,10 +180,8 @@ class ChunkShuffleSplitGeneratorBase(AbstractSplitGenerator): 6. If sharded, take this worker's slice via balanced ``_compute_shard_range``. 7. Map each chunk through :meth:`_chunk_to_split`. - Subclasses implement the three abstract hooks. Reader paths - (``RawFileSplitRead`` for append, ``DataEvolutionSplitRead`` for DE) - are unchanged because chunks ride on existing wrappers - (``SlicedSplit`` / ``IndexedSplit``). + Subclasses implement the three abstract hooks. Chunks ride on existing + reader wrappers (``SlicedSplit`` / ``IndexedSplit``). """ def __init__( @@ -86,6 +196,9 @@ def __init__( super().__init__(table, target_split_size, open_file_cost, deletion_files_map) self.seed = seed self.chunk_size = chunk_size + # Planning-only cache. Readers continue to load DVs through their + # existing split-local factories. + self._deletion_vector_cache = {} def create_splits(self, file_entries: List[ManifestEntry]) -> List[Split]: if not file_entries: @@ -146,9 +259,64 @@ def _slice_group_into_chunks(self, entries: List[ManifestEntry]) -> List[List[An def _chunk_to_split(self, chunk: _Chunk) -> Split: """Wrap a chunk into a Split that the existing readers consume.""" + def _deletion_file( + self, + partition: GenericRow, + bucket: int, + file_name: str, + ) -> Optional[DeletionFile]: + partition_key = (tuple(partition.values), bucket) + return self.deletion_files_map.get(partition_key, {}).get(file_name) + + def _live_row_slicer( + self, + physical_row_count: int, + deletion_file: Optional[DeletionFile], + ) -> Optional[_LiveRowRangeSlicer]: + if deletion_file is None or deletion_file.cardinality == 0: + return _LiveRowRangeSlicer(physical_row_count, iter(())) + + cardinality = deletion_file.cardinality + if cardinality is not None: + if cardinality < 0 or cardinality > physical_row_count: + raise ValueError( + f"Deletion vector cardinality {cardinality} is outside valid " + f"range [0, {physical_row_count}]." + ) + if cardinality == physical_row_count: + return None + + cache_key = ( + deletion_file.dv_index_path, + deletion_file.offset, + deletion_file.length, + ) + deletion_vector = self._deletion_vector_cache.get(cache_key) + if deletion_vector is None: + deletion_vector = DeletionVector.read(self.table.file_io, deletion_file) + self._deletion_vector_cache[cache_key] = deletion_vector + + actual_cardinality = deletion_vector.get_cardinality() + if cardinality is not None and cardinality != actual_cardinality: + raise ValueError( + f"Deletion vector cardinality mismatch, metadata is {cardinality} " + f"but bitmap contains {actual_cardinality} positions." + ) + if actual_cardinality > physical_row_count: + raise ValueError( + f"Deletion vector cardinality {actual_cardinality} exceeds physical " + f"row count {physical_row_count}." + ) + if actual_cardinality == physical_row_count: + return None + return _LiveRowRangeSlicer( + physical_row_count, + iter(deletion_vector.bit_map()), + ) + # --------------------------------------------------------------------------- -# Append (non-DE, non-DV) implementation +# Append implementation # --------------------------------------------------------------------------- @@ -167,7 +335,7 @@ class _FileSegment: class AppendChunkShuffleSplitGenerator(ChunkShuffleSplitGeneratorBase): - """Chunk-shuffled splits for plain append tables (non-PK, non-DV, non-DE).""" + """Chunk-shuffled splits for plain append tables (non-PK, non-DE).""" def _sort_key(self, entry: ManifestEntry): return ( @@ -180,8 +348,9 @@ def _slice_group_into_chunks( self, entries: List[ManifestEntry] ) -> List[List[_FileSegment]]: """Cut a (partition, bucket) group into chunks of at most - ``self.chunk_size`` rows. ``chunk_size`` is a hard upper bound: - the last chunk may be smaller, but no chunk exceeds it. + ``self.chunk_size`` live rows. ``chunk_size`` is a hard upper bound: + the last chunk may be smaller, but no chunk exceeds it after DV + filtering. """ chunks: List[List[_FileSegment]] = [] current: List[_FileSegment] = [] @@ -189,9 +358,16 @@ def _slice_group_into_chunks( for entry in entries: file = entry.file - offset = 0 - remaining = file.row_count - while remaining > 0: + deletion_file = self._deletion_file( + entry.partition, + entry.bucket, + file.file_name, + ) + slicer = self._live_row_slicer(file.row_count, deletion_file) + if slicer is None: + continue + + while True: avail = self.chunk_size - current_rows if avail <= 0: chunks.append(current) @@ -199,16 +375,18 @@ def _slice_group_into_chunks( current_rows = 0 avail = self.chunk_size - take = min(remaining, avail) + live_range = slicer.take(avail) + if live_range is None: + break - if take == file.row_count and offset == 0: + if live_range.start == 0 and live_range.end == file.row_count: current.append(_FileSegment(file, None, None)) else: - current.append(_FileSegment(file, offset, offset + take)) + current.append( + _FileSegment(file, live_range.start, live_range.end) + ) - current_rows += take - offset += take - remaining -= take + current_rows += live_range.live_rows if current: chunks.append(current) @@ -225,13 +403,18 @@ def _chunk_to_split(self, chunk: _Chunk) -> Split: # set_file_path is already done once per unique file in # ChunkShuffleSplitGeneratorBase.create_splits. + data_deletion_files = self._get_deletion_files_for_split( + files, + chunk.partition, + chunk.bucket, + ) data_split = DataSplit( files=files, partition=chunk.partition, bucket=chunk.bucket, raw_convertible=True, - data_deletion_files=None, + data_deletion_files=data_deletion_files, ) if shard_file_idx_map: @@ -295,11 +478,43 @@ def _slice_group_into_chunks( chunks: List[List[_AlignedGroupSegment]] = [] current: List[_AlignedGroupSegment] = [] current_rows = 0 + partition = entries[0].partition + bucket = entries[0].bucket for group_range, group_files in aligned_groups: - offset = 0 - group_rows = group_range.count() - while offset < group_rows: + anchor = None + deletion_file = None + if self.deletion_files_map: + anchor = retrieve_anchor_file(group_files) + deletion_file = self._deletion_file( + partition, + bucket, + anchor.file_name, + ) + + physical_row_count = group_range.count() + first_row_id = group_range.from_ + if deletion_file is not None: + anchor_range = anchor.row_id_range() + if ( + anchor_range.from_ > group_range.from_ + or anchor_range.to < group_range.to + ): + raise ValueError( + f"Data evolution anchor range {anchor_range} does not contain " + f"aligned group range {group_range}." + ) + physical_row_count = anchor.row_count + first_row_id = anchor_range.from_ + + slicer = self._live_row_slicer( + physical_row_count, + deletion_file, + ) + if slicer is None: + continue + + while True: avail = self.chunk_size - current_rows if avail <= 0: chunks.append(current) @@ -307,14 +522,15 @@ def _slice_group_into_chunks( current_rows = 0 avail = self.chunk_size - take = min(group_rows - offset, avail) + live_range = slicer.take(avail) + if live_range is None: + break seg_range = Range( - group_range.from_ + offset, - group_range.from_ + offset + take - 1, + first_row_id + live_range.start, + first_row_id + live_range.end - 1, ) current.append(_AlignedGroupSegment(group_files, seg_range)) - current_rows += take - offset += take + current_rows += live_range.live_rows if current: chunks.append(current) @@ -334,12 +550,17 @@ def _chunk_to_split(self, chunk: _Chunk) -> Split: row_ranges.append(seg.row_range) row_ranges.sort(key=lambda r: r.from_) + data_deletion_files = self._get_deletion_files_for_split( + all_files, + chunk.partition, + chunk.bucket, + ) data_split = DataSplit( files=all_files, partition=chunk.partition, bucket=chunk.bucket, raw_convertible=False, - data_deletion_files=None, + data_deletion_files=data_deletion_files, ) return IndexedSplit(data_split, row_ranges, scores=None) diff --git a/paimon-python/pypaimon/read/scanner/file_scanner.py b/paimon-python/pypaimon/read/scanner/file_scanner.py index e952a9f3a4c8..3442fedb3a82 100755 --- a/paimon-python/pypaimon/read/scanner/file_scanner.py +++ b/paimon-python/pypaimon/read/scanner/file_scanner.py @@ -612,8 +612,6 @@ def with_chunk_shuffle(self, seed: int, chunk_size: int) -> 'FileScanner': def _validate_chunk_shuffle_compat(self) -> None: if self.table.is_primary_key_table: raise ValueError("chunk_shuffle only supports append tables") - if self.deletion_vectors_enabled: - raise ValueError("chunk_shuffle not supported with deletion vectors") if self.start_pos_of_this_subtask is not None: raise ValueError("chunk_shuffle cannot combine with with_slice") if self.limit is not None: diff --git a/paimon-python/pypaimon/read/split_read.py b/paimon-python/pypaimon/read/split_read.py index 6ced11f0c127..34947db61b4b 100644 --- a/paimon-python/pypaimon/read/split_read.py +++ b/paimon-python/pypaimon/read/split_read.py @@ -764,6 +764,11 @@ def raw_reader_supplier(self, file: DataFileMeta, dv_factory: Optional[Callable] row_tracking_enabled=True) dv = dv_factory() if dv_factory else None if dv: + if file.file_name in shard_file_idx_map: + dv = PositionMappedDeletionVector( + dv, + file_offset=start_pos, + ) return ApplyDeletionVectorReader(RowPositionReader(file_batch_reader), dv) else: return file_batch_reader 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 448a8ab1b0c1..4ef3748f78c1 100644 --- a/paimon-python/pypaimon/tests/data_evolution_deletion_vector_test.py +++ b/paimon-python/pypaimon/tests/data_evolution_deletion_vector_test.py @@ -16,6 +16,7 @@ # under the License. import unittest +from unittest.mock import Mock import pyarrow as pa @@ -32,7 +33,9 @@ MergeAllBatchReader, ) from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader +from pypaimon.read.sliced_split import SlicedSplit from pypaimon.read.split import DataSplit +from pypaimon.read.split_read import RawFileSplitRead from pypaimon.table.row.blob import Blob, BlobData from pypaimon.table.row.generic_row import GenericRow from pypaimon.table.source.deletion_file import DeletionFile @@ -160,6 +163,39 @@ def test_apply_deletion_vector_reader_uses_mapped_deletion_vector(self): self.assertTrue(reader.deletion_vector().is_deleted(1)) self.assertFalse(reader.deletion_vector().is_deleted(2)) + def test_append_sliced_reader_maps_positions_to_original_file_offsets(self): + file = _file("slice.parquet", 0, 10, 1) + data_split = DataSplit( + files=[file], + partition=GenericRow([], []), + bucket=0, + raw_convertible=True, + data_deletion_files=None, + ) + sliced_split = SlicedSplit( + data_split, + {"slice.parquet": (5, 10)}, + ) + split_read = RawFileSplitRead.__new__(RawFileSplitRead) + split_read.split = sliced_split + split_read._get_final_read_data_fields = Mock(return_value=[]) + split_read.file_reader_supplier = Mock( + return_value=_OneBatchReader([5, 6, 7, 8, 9]) + ) + deletion_vector = BitmapDeletionVector() + deletion_vector.delete(7) + + reader = split_read.raw_reader_supplier( + file, + dv_factory=lambda: deletion_vector, + ) + + self.assertEqual( + [5, 6, 8, 9], + reader.read_arrow_batch().column(0).to_pylist(), + ) + self.assertIsNone(reader.read_arrow_batch()) + def test_data_evolution_merge_reader_handles_fully_deleted_file(self): deletion_vector = BitmapDeletionVector() deletion_vector.delete(0) diff --git a/paimon-python/pypaimon/tests/scanner/chunk_shuffle_split_generator_test.py b/paimon-python/pypaimon/tests/scanner/chunk_shuffle_split_generator_test.py index ce69096d7cd6..243f2520df77 100644 --- a/paimon-python/pypaimon/tests/scanner/chunk_shuffle_split_generator_test.py +++ b/paimon-python/pypaimon/tests/scanner/chunk_shuffle_split_generator_test.py @@ -26,19 +26,22 @@ import shutil import tempfile import unittest -from unittest.mock import Mock +from unittest.mock import Mock, patch import pyarrow as pa from pypaimon import CatalogFactory, Schema +from pypaimon.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector from pypaimon.globalindex.indexed_split import IndexedSplit from pypaimon.manifest.schema.data_file_meta import DataFileMeta from pypaimon.read.scanner.chunk_shuffle_split_generator import ( AppendChunkShuffleSplitGenerator, DataEvolutionChunkShuffleSplitGenerator, + _LiveRowRangeSlicer, ) from pypaimon.read.sliced_split import SlicedSplit from pypaimon.read.split import DataSplit +from pypaimon.table.source.deletion_file import DeletionFile from pypaimon.utils.range import Range @@ -63,33 +66,41 @@ def _mock_entry(partition_values, bucket, file_name, row_count, file_size=1024): return entry -def _make_generator(seed, chunk_size, table=None): +def _make_generator(seed, chunk_size, table=None, deletion_files_map=None): if table is None: table = _mock_table() return AppendChunkShuffleSplitGenerator( table, target_split_size=128 * 1024 * 1024, open_file_cost=4 * 1024 * 1024, - deletion_files_map=None, + deletion_files_map=deletion_files_map, seed=seed, chunk_size=chunk_size, ) -def _make_de_generator(seed, chunk_size, table=None): +def _make_de_generator(seed, chunk_size, table=None, deletion_files_map=None): if table is None: table = _mock_table() return DataEvolutionChunkShuffleSplitGenerator( table, target_split_size=128 * 1024 * 1024, open_file_cost=4 * 1024 * 1024, - deletion_files_map=None, + deletion_files_map=deletion_files_map, seed=seed, chunk_size=chunk_size, ) -def _mock_de_entry(partition_values, bucket, file_name, first_row_id, row_count, file_size=1024): +def _mock_de_entry( + partition_values, + bucket, + file_name, + first_row_id, + row_count, + file_size=1024, + max_sequence_number=0, +): """A DE-flavoured mock entry: file carries first_row_id and a real Range so :meth:`row_id_range` and ``Range.overlaps`` work.""" entry = Mock() @@ -101,6 +112,7 @@ def _mock_de_entry(partition_values, bucket, file_name, first_row_id, row_count, file.file_size = file_size file.row_count = row_count file.first_row_id = first_row_id + file.max_sequence_number = max_sequence_number file.row_id_range = lambda f=first_row_id, c=row_count: Range(f, f + c - 1) file.set_file_path = Mock() entry.file = file @@ -139,6 +151,43 @@ def _split_rows(split): return split.row_count +def _bitmap_deletion_vector(*positions): + deletion_vector = BitmapDeletionVector() + for position in positions: + deletion_vector.delete(position) + return deletion_vector + + +class LiveRowRangeSlicerTest(unittest.TestCase): + + def test_slices_by_live_rows_and_absorbs_deletion_runs(self): + slicer = _LiveRowRangeSlicer(10, iter([0, 3, 4, 9])) + + first = slicer.take(3) + second = slicer.take(3) + + self.assertEqual((first.start, first.end, first.live_rows), (0, 6, 3)) + self.assertEqual((second.start, second.end, second.live_rows), (6, 10, 3)) + self.assertIsNone(slicer.take(3)) + + def test_returns_smaller_tail_and_skips_fully_deleted_source(self): + tail_slicer = _LiveRowRangeSlicer(6, iter([1, 4])) + tail = tail_slicer.take(10) + self.assertEqual((tail.start, tail.end, tail.live_rows), (0, 6, 4)) + self.assertIsNone(tail_slicer.take(1)) + + deleted_slicer = _LiveRowRangeSlicer(3, iter([0, 1, 2])) + self.assertIsNone(deleted_slicer.take(1)) + + def test_rejects_invalid_deletion_positions(self): + with self.assertRaisesRegex(ValueError, "outside physical row range"): + _LiveRowRangeSlicer(3, iter([-1])) + + slicer = _LiveRowRangeSlicer(3, iter([0, 0])) + with self.assertRaisesRegex(ValueError, "strictly increasing"): + slicer.take(1) + + class ChunkShuffleSplitGeneratorAlgoTest(unittest.TestCase): def test_no_entries_returns_empty(self): @@ -200,6 +249,76 @@ def test_chunk_size_larger_than_total(self): self.assertIsInstance(splits[0], DataSplit) self.assertEqual(_split_rows(splits[0]), 60) + def test_deletion_vector_slices_by_live_rows_and_is_attached(self): + entry = _mock_entry([], 0, 'f1', 10) + deletion_file = DeletionFile('dv.index', 10, 20, cardinality=4) + deletion_files_map = {((), 0): {'f1': deletion_file}} + deletion_vector = _bitmap_deletion_vector(0, 3, 4, 9) + gen = _make_generator( + seed=1, + chunk_size=3, + deletion_files_map=deletion_files_map, + ) + + with patch( + 'pypaimon.read.scanner.chunk_shuffle_split_generator.DeletionVector.read', + return_value=deletion_vector, + ) as read: + splits = gen.create_splits([entry]) + + self.assertEqual( + sorted(s.shard_file_idx_map()['f1'] for s in splits), + [(0, 6), (6, 10)], + ) + for split in splits: + self.assertEqual( + split.data_split().data_deletion_files, + [deletion_file], + ) + read.assert_called_once_with(gen.table.file_io, deletion_file) + + def test_planning_cache_reuses_same_deletion_vector_descriptor(self): + entries = [ + _mock_entry([], 0, 'f1', 5), + _mock_entry([], 0, 'f2', 5), + ] + deletion_file = DeletionFile('dv.index', 10, 20, cardinality=1) + deletion_files_map = { + ((), 0): { + 'f1': deletion_file, + 'f2': deletion_file, + } + } + gen = _make_generator( + seed=1, + chunk_size=2, + deletion_files_map=deletion_files_map, + ) + + with patch( + 'pypaimon.read.scanner.chunk_shuffle_split_generator.DeletionVector.read', + return_value=_bitmap_deletion_vector(1), + ) as read: + gen.create_splits(entries) + + self.assertEqual(read.call_count, 1) + + def test_fully_deleted_file_does_not_create_a_segment(self): + entry = _mock_entry([], 0, 'f1', 5) + deletion_file = DeletionFile('dv.index', 10, 20, cardinality=5) + gen = _make_generator( + seed=1, + chunk_size=2, + deletion_files_map={((), 0): {'f1': deletion_file}}, + ) + + with patch( + 'pypaimon.read.scanner.chunk_shuffle_split_generator.DeletionVector.read', + ) as read: + self.assertEqual(gen.create_splits([entry]), []) + + read.assert_not_called() + def test_deterministic_same_seed_same_order(self): entries = [_mock_entry([], 0, f'f{i}', 100) for i in range(20)] gen1 = _make_generator(seed=42, chunk_size=50) @@ -453,12 +572,11 @@ def test_pk_table_rejected(self): with self.assertRaisesRegex(ValueError, "only supports append tables"): scan.plan() - def test_dv_table_rejected(self): + def test_dv_table_supported(self): table = self._append_table('cs_dv', options={'deletion-vectors.enabled': 'true'}) scan = table.new_read_builder().new_scan() scan.with_chunk_shuffle(seed=1, chunk_size=100) - with self.assertRaisesRegex(ValueError, "deletion vectors"): - scan.plan() + self.assertEqual(scan.plan().splits(), []) def test_with_slice_then_chunk_shuffle_rejected(self): table = self._append_table('cs_slice') @@ -616,6 +734,56 @@ def test_blob_propagates_when_group_split(self): files = sorted(f.file_name for f in s.files) self.assertEqual(files, ['g0.blob', 'g0.parquet']) + def test_deletion_vector_uses_anchor_and_slices_by_live_rows(self): + anchor = _mock_de_entry( + [], + 0, + 'anchor.parquet', + 100, + 10, + max_sequence_number=1, + ) + blob = _mock_de_entry( + [], + 0, + 'field.blob', + 100, + 10, + max_sequence_number=2, + ) + deletion_file = DeletionFile('dv.index', 10, 20, cardinality=4) + deletion_files_map = {((), 0): {'anchor.parquet': deletion_file}} + gen = _make_de_generator( + seed=1, + chunk_size=3, + deletion_files_map=deletion_files_map, + ) + + with patch( + 'pypaimon.read.scanner.chunk_shuffle_split_generator.DeletionVector.read', + return_value=_bitmap_deletion_vector(1, 2, 7, 9), + ) as read: + splits = gen.create_splits([blob, anchor]) + + ranges = sorted( + (row_range.from_, row_range.to) + for split in splits + for row_range in split.row_ranges() + ) + self.assertEqual(ranges, [(100, 104), (105, 109)]) + for split in splits: + self.assertEqual( + sorted(f.file_name for f in split.files), + ['anchor.parquet', 'field.blob'], + ) + deletion_by_name = dict( + (file.file_name, dv) + for file, dv in zip(split.files, split.data_deletion_files) + ) + self.assertEqual(deletion_by_name['anchor.parquet'], deletion_file) + self.assertIsNone(deletion_by_name['field.blob']) + read.assert_called_once_with(gen.table.file_io, deletion_file) + def test_deterministic_same_seed(self): entries = [_mock_de_entry([], 0, f'g{i:02d}.parquet', i * 100, 100) for i in range(20)] gen1 = _make_de_generator(seed=42, chunk_size=100) @@ -705,19 +873,22 @@ def setUpClass(cls): def tearDownClass(cls): shutil.rmtree(cls.tempdir, ignore_errors=True) - def _create_de_table(self, name): + def _create_de_table(self, name, deletion_vectors_enabled=False): pa_schema = pa.schema([ ('id', pa.int32()), ('value', pa.string()), ('payload', pa.large_binary()), ]) + options = { + 'row-tracking.enabled': 'true', + 'data-evolution.enabled': 'true', + 'blob.target-file-size': '1 b', + } + if deletion_vectors_enabled: + options['deletion-vectors.enabled'] = 'true' schema = Schema.from_pyarrow_schema( pa_schema, - options={ - 'row-tracking.enabled': 'true', - 'data-evolution.enabled': 'true', - 'blob.target-file-size': '1 b', - }, + options=options, ) identifier = f'default.{name}' self.catalog.create_table(identifier, schema, False) @@ -744,6 +915,14 @@ def _commit_full_rows(self, table, pa_schema, ids): tc.close() return commit_messages + @staticmethod + def _delete_by_row_id(table, row_ids): + write_builder = table.new_batch_write_builder() + commit_messages = write_builder.new_update().delete_by_row_id(row_ids) + table_commit = write_builder.new_commit() + table_commit.commit(commit_messages) + table_commit.close() + def _assert_commit_has_main_and_multiple_blob_files(self, commit_messages): all_files = [f for msg in commit_messages for f in msg.new_files] main_files = [f for f in all_files if not DataFileMeta.is_blob_file(f.file_name)] @@ -816,6 +995,41 @@ def plan_sigs(worker): for worker in range(4): self.assertEqual(plan_sigs(worker), plan_sigs(worker)) + def test_python_delete_then_chunk_shuffle_read(self): + table, pa_schema = self._create_de_table( + 'cs_de_python_delete', + deletion_vectors_enabled=True, + ) + commit_messages = self._commit_full_rows( + table, + pa_schema, + list(range(12)), + ) + self._assert_commit_has_main_and_multiple_blob_files(commit_messages) + self._delete_by_row_id(table, [0, 2, 3, 7, 11]) + + read_builder = table.new_read_builder() + splits = read_builder.new_scan() \ + .with_chunk_shuffle(seed=42, chunk_size=3) \ + .plan() \ + .splits() + self._assert_splits_include_blob_files(splits) + + table_read = read_builder.new_read() + rows_per_split = [ + table_read.to_arrow([split]).num_rows + for split in splits + ] + self.assertEqual(sorted(rows_per_split), [1, 3, 3]) + + actual = table_read.to_arrow(splits).sort_by('id') + expected_ids = [1, 4, 5, 6, 8, 9, 10] + self.assertEqual(actual.column('id').to_pylist(), expected_ids) + self.assertEqual( + actual.column('payload').to_pylist(), + self._payloads(expected_ids), + ) + if __name__ == '__main__': unittest.main() From c16e43b1926c71ceb92892916a697e70ce081728 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=96=86=E5=AE=87?= Date: Mon, 27 Jul 2026 11:16:31 +0800 Subject: [PATCH 2/2] add tests --- .../scanner/chunk_shuffle_split_generator.py | 52 +++--- .../chunk_shuffle_split_generator_test.py | 149 +++++++++++++++++- 2 files changed, 178 insertions(+), 23 deletions(-) diff --git a/paimon-python/pypaimon/read/scanner/chunk_shuffle_split_generator.py b/paimon-python/pypaimon/read/scanner/chunk_shuffle_split_generator.py index 368f1edee5eb..4419d8d99adb 100644 --- a/paimon-python/pypaimon/read/scanner/chunk_shuffle_split_generator.py +++ b/paimon-python/pypaimon/read/scanner/chunk_shuffle_split_generator.py @@ -46,12 +46,18 @@ def _null_safe_partition_key(partition_values) -> tuple: @dataclass -class _LiveRowRange: - """A physical half-open range containing ``live_rows`` visible rows.""" +class _PhysicalRowSlice: + """A half-open physical row slice containing visible rows.""" - start: int - end: int - live_rows: int + start_inclusive: int + end_exclusive: int + live_row_count: int + + def to_closed_row_id_range(self, first_row_id: int) -> Range: + return Range( + first_row_id + self.start_inclusive, + first_row_id + self.end_exclusive - 1, + ) class _LiveRowRangeSlicer: @@ -77,7 +83,7 @@ def __init__( self._last_deleted_position = None self._next_deleted_position = self._advance_deleted_position() - def take(self, expected_live_rows: int) -> Optional[_LiveRowRange]: + def take(self, expected_live_rows: int) -> Optional[_PhysicalRowSlice]: if expected_live_rows <= 0: raise ValueError( f"expected_live_rows must be positive, got {expected_live_rows}" @@ -111,7 +117,7 @@ def take(self, expected_live_rows: int) -> Optional[_LiveRowRange]: # immediately after the boundary to this range so the next # range starts at a live row (or EOF). self._skip_deleted_positions_at_cursor() - return _LiveRowRange( + return _PhysicalRowSlice( start, self._physical_position, live_rows, @@ -123,7 +129,7 @@ def take(self, expected_live_rows: int) -> Optional[_LiveRowRange]: if live_rows == 0: return None - return _LiveRowRange(start, self._physical_position, live_rows) + return _PhysicalRowSlice(start, self._physical_position, live_rows) def _skip_deleted_positions_at_cursor(self) -> None: while ( @@ -201,6 +207,8 @@ def __init__( self._deletion_vector_cache = {} def create_splits(self, file_entries: List[ManifestEntry]) -> List[Split]: + """TODO: Lazily initialize DataSplits to avoid creating too many objects.""" + if not file_entries: return [] @@ -375,18 +383,25 @@ def _slice_group_into_chunks( current_rows = 0 avail = self.chunk_size - live_range = slicer.take(avail) - if live_range is None: + physical_slice = slicer.take(avail) + if physical_slice is None: break - if live_range.start == 0 and live_range.end == file.row_count: + if ( + physical_slice.start_inclusive == 0 + and physical_slice.end_exclusive == file.row_count + ): current.append(_FileSegment(file, None, None)) else: current.append( - _FileSegment(file, live_range.start, live_range.end) + _FileSegment( + file, + physical_slice.start_inclusive, + physical_slice.end_exclusive, + ) ) - current_rows += live_range.live_rows + current_rows += physical_slice.live_row_count if current: chunks.append(current) @@ -522,15 +537,12 @@ def _slice_group_into_chunks( current_rows = 0 avail = self.chunk_size - live_range = slicer.take(avail) - if live_range is None: + physical_slice = slicer.take(avail) + if physical_slice is None: break - seg_range = Range( - first_row_id + live_range.start, - first_row_id + live_range.end - 1, - ) + seg_range = physical_slice.to_closed_row_id_range(first_row_id) current.append(_AlignedGroupSegment(group_files, seg_range)) - current_rows += live_range.live_rows + current_rows += physical_slice.live_row_count if current: chunks.append(current) diff --git a/paimon-python/pypaimon/tests/scanner/chunk_shuffle_split_generator_test.py b/paimon-python/pypaimon/tests/scanner/chunk_shuffle_split_generator_test.py index 243f2520df77..f1e4d5e6575f 100644 --- a/paimon-python/pypaimon/tests/scanner/chunk_shuffle_split_generator_test.py +++ b/paimon-python/pypaimon/tests/scanner/chunk_shuffle_split_generator_test.py @@ -166,14 +166,36 @@ def test_slices_by_live_rows_and_absorbs_deletion_runs(self): first = slicer.take(3) second = slicer.take(3) - self.assertEqual((first.start, first.end, first.live_rows), (0, 6, 3)) - self.assertEqual((second.start, second.end, second.live_rows), (6, 10, 3)) + self.assertEqual( + ( + first.start_inclusive, + first.end_exclusive, + first.live_row_count, + ), + (0, 6, 3), + ) + self.assertEqual( + ( + second.start_inclusive, + second.end_exclusive, + second.live_row_count, + ), + (6, 10, 3), + ) + self.assertEqual(first.to_closed_row_id_range(100), Range(100, 105)) self.assertIsNone(slicer.take(3)) def test_returns_smaller_tail_and_skips_fully_deleted_source(self): tail_slicer = _LiveRowRangeSlicer(6, iter([1, 4])) tail = tail_slicer.take(10) - self.assertEqual((tail.start, tail.end, tail.live_rows), (0, 6, 4)) + self.assertEqual( + ( + tail.start_inclusive, + tail.end_exclusive, + tail.live_row_count, + ), + (0, 6, 4), + ) self.assertIsNone(tail_slicer.take(1)) deleted_slicer = _LiveRowRangeSlicer(3, iter([0, 1, 2])) @@ -277,6 +299,63 @@ def test_deletion_vector_slices_by_live_rows_and_is_attached(self): ) read.assert_called_once_with(gen.table.file_io, deletion_file) + def test_two_deletion_vector_files_share_one_live_row_chunk(self): + # Each file contributes 5 - 2 = 3 live rows, so both should fit + # exactly in one six-row chunk. + entries = [ + _mock_entry([], 0, 'f1', 5), + _mock_entry([], 0, 'f2', 5), + ] + first_deletion_file = DeletionFile( + 'dv.index', + 10, + 20, + cardinality=2, + ) + second_deletion_file = DeletionFile( + 'dv.index', + 30, + 20, + cardinality=2, + ) + deletion_files_map = { + ((), 0): { + 'f1': first_deletion_file, + 'f2': second_deletion_file, + } + } + deletion_vectors = { + first_deletion_file: _bitmap_deletion_vector(1, 3), + second_deletion_file: _bitmap_deletion_vector(0, 4), + } + gen = _make_generator( + seed=1, + chunk_size=6, + deletion_files_map=deletion_files_map, + ) + + with patch( + 'pypaimon.read.scanner.chunk_shuffle_split_generator.DeletionVector.read', + side_effect=lambda _, deletion_file: deletion_vectors[deletion_file], + ) as read: + splits = gen.create_splits(entries) + + self.assertEqual(len(splits), 1) + self.assertIsInstance(splits[0], DataSplit) + self.assertEqual( + [file.file_name for file in splits[0].files], + ['f1', 'f2'], + ) + self.assertEqual( + splits[0].data_deletion_files, + [first_deletion_file, second_deletion_file], + ) + self.assertEqual(read.call_count, 2) + self.assertEqual( + {mock_call.args[1] for mock_call in read.call_args_list}, + {first_deletion_file, second_deletion_file}, + ) + def test_planning_cache_reuses_same_deletion_vector_descriptor(self): entries = [ _mock_entry([], 0, 'f1', 5), @@ -784,6 +863,70 @@ def test_deletion_vector_uses_anchor_and_slices_by_live_rows(self): self.assertIsNone(deletion_by_name['field.blob']) read.assert_called_once_with(gen.table.file_io, deletion_file) + def test_multiple_deletion_vector_groups_share_one_live_row_chunk(self): + # The groups contribute 3 live rows each despite spanning 5 and 6 + # physical rows, so both should become segments of the same chunk. + entries = [ + _mock_de_entry([], 0, 'g0.parquet', 100, 5), + _mock_de_entry([], 0, 'g1.parquet', 200, 6), + ] + first_deletion_file = DeletionFile( + 'dv.index', + 10, + 20, + cardinality=2, + ) + second_deletion_file = DeletionFile( + 'dv.index', + 30, + 20, + cardinality=3, + ) + deletion_files_map = { + ((), 0): { + 'g0.parquet': first_deletion_file, + 'g1.parquet': second_deletion_file, + } + } + deletion_vectors = { + first_deletion_file: _bitmap_deletion_vector(1, 3), + second_deletion_file: _bitmap_deletion_vector(0, 2, 5), + } + gen = _make_de_generator( + seed=1, + chunk_size=6, + deletion_files_map=deletion_files_map, + ) + + with patch( + 'pypaimon.read.scanner.chunk_shuffle_split_generator.DeletionVector.read', + side_effect=lambda _, deletion_file: deletion_vectors[deletion_file], + ) as read: + splits = gen.create_splits(entries) + + self.assertEqual(len(splits), 1) + self.assertIsInstance(splits[0], IndexedSplit) + self.assertEqual( + [ + (row_range.from_, row_range.to) + for row_range in splits[0].row_ranges() + ], + [(100, 104), (200, 205)], + ) + self.assertEqual( + [file.file_name for file in splits[0].files], + ['g0.parquet', 'g1.parquet'], + ) + self.assertEqual( + splits[0].data_deletion_files, + [first_deletion_file, second_deletion_file], + ) + self.assertEqual(read.call_count, 2) + self.assertEqual( + {mock_call.args[1] for mock_call in read.call_args_list}, + {first_deletion_file, second_deletion_file}, + ) + def test_deterministic_same_seed(self): entries = [_mock_de_entry([], 0, f'g{i:02d}.parquet', i * 100, 100) for i in range(20)] gen1 = _make_de_generator(seed=42, chunk_size=100)