[python] Plan chunk-shuffled native reads with IndexedSplit

paimon-python/pypaimon/globalindex/indexed_split.py

@@ -15,8 +15,11 @@
 # specific language governing permissions and limitations
 # under the License.
 
-"""
-IndexedSplit wraps a Split with row ranges and optional scores.
+"""IndexedSplit wraps a Split with row ranges and optional scores.
+
+Ranges use the coordinate system of the table read path: stable row IDs for
+row-tracked tables and split-local physical positions for tables without row
+tracking.
 """
 
 from typing import List, Optional
@@ -43,19 +46,17 @@ def __init__(
         data_split: 'Split',
         row_ranges: List['Range'],
         scores: Optional[List[float]] = None,
-        exact_merged_row_count: Optional[int] = None,
     ):
         self._data_split = data_split
         self._row_ranges = row_ranges
         self._scores = scores
-        self._exact_merged_row_count = exact_merged_row_count
 
     def data_split(self) -> 'Split':
         """Return the underlying data split."""
         return self._data_split
 
     def row_ranges(self) -> List['Range']:
-        """Return the row ranges from global index."""
+        """Return ranges in the coordinate system of the read path."""
         return self._row_ranges
 
     def scores(self) -> Optional[List[float]]:
@@ -90,8 +91,6 @@ def row_count(self) -> int:
         return sum(r.count() for r in self._row_ranges)
 
     def merged_row_count(self):
-        if self._exact_merged_row_count is not None:
-            return self._exact_merged_row_count
         return self.row_count
 
     # Delegate other properties to data_split
@@ -157,7 +156,6 @@ def __eq__(self, other):
             self._data_split == other._data_split
             and self._row_ranges == other._row_ranges
             and self._scores == other._scores
-            and self._exact_merged_row_count == other._exact_merged_row_count
         )
 
     def __hash__(self):
@@ -166,10 +164,8 @@ def __hash__(self):
             id(self._data_split),
             tuple(self._row_ranges),
             scores_hash,
-            self._exact_merged_row_count,
         ))
 
     def __repr__(self):
         return (f"IndexedSplit(data_split={self._data_split}, "
-                f"row_ranges={self._row_ranges}, scores={self._scores}, "
-                f"exact_merged_row_count={self._exact_merged_row_count})")
+                f"row_ranges={self._row_ranges}, scores={self._scores})")

paimon-python/pypaimon/multimodal/temporal.py

@@ -1228,11 +1228,7 @@ def fetch(self, row_ids):
             allowed = Range.and_(wanted, self._split_ranges[split_index])
             if not allowed:
                 continue
-            indexed = IndexedSplit(
-                split,
-                allowed,
-                exact_merged_row_count=sum(r.count() for r in allowed),
-            )
+            indexed = IndexedSplit(split, allowed)
             if auth_result is not None:
                 indexed = QueryAuthSplit(indexed, auth_result)
             selected_splits.append(indexed)

paimon-python/pypaimon/read/datasource/torch_dataset.py

@@ -218,11 +218,7 @@ def select_indexed_splits(
         if not allowed:
             continue
 
-        indexed = IndexedSplit(
-            split,
-            allowed,
-            exact_merged_row_count=sum(r.count() for r in allowed),
-        )
+        indexed = IndexedSplit(split, allowed)
         selected.append(
             QueryAuthSplit(indexed, auth_result)
             if auth_result is not None else indexed

paimon-python/pypaimon/read/native_plan.py

@@ -371,7 +371,9 @@ def native_plan(
         row_ranges: Optional[List[Tuple[int, int]]] = None,
         incremental_range: Optional[Tuple[int, int]] = None,
         row_position_slice: Optional[Tuple[int, int]] = None,
-        row_position_shard: Optional[Tuple[int, int]] = None) -> Plan:
+        row_position_shard: Optional[Tuple[int, int]] = None,
+        chunk_shuffle: Optional[Tuple[int, int]] = None,
+        shard: Optional[Tuple[int, int]] = None) -> Plan:
     """Plan with pypaimon_rust, preserving snapshot metadata.
 
     Native conversion or planning failures are handled by TableScan, which
@@ -390,12 +392,20 @@ def native_plan(
         scan = scan.with_row_position_slice(*row_position_slice)
     if row_position_shard is not None:
         scan = scan.with_row_position_shard(*row_position_shard)
+    if chunk_shuffle is not None:
+        seed, chunk_size = chunk_shuffle
+        scan = scan.with_chunk_shuffle(str(seed), chunk_size)
+    if shard is not None:
+        scan = scan.with_shard(*shard)
     rust_plan = scan.plan()
     rust_splits = rust_plan.splits()
     pfields = _partition_fields(table)
     # Trimmed primary keys decode per-file min/max keys (PK merge-on-read).
     kfields = table.trimmed_primary_keys_fields
-    splits = [deserialize_split_v1(split.serialize(), pfields, kfields) for split in rust_splits]
+    splits = [
+        deserialize_split_v1(split.serialize(), pfields, kfields)
+        for split in rust_splits
+    ]
     if table.options.native_read_enabled():
         # Retain the opaque Rust split next to the Python metadata view. The
         # normal planner/reader contract remains a Python Split list, while

paimon-python/pypaimon/read/scanner/chunk_shuffle_split_generator.py

@@ -26,7 +26,6 @@
 from pypaimon.manifest.schema.data_file_meta import DataFileMeta
 from pypaimon.manifest.schema.manifest_entry import ManifestEntry
 from pypaimon.read.scanner.split_generator import AbstractSplitGenerator
-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
@@ -47,17 +46,16 @@ def _null_safe_partition_key(partition_values) -> tuple:
 
 @dataclass
 class _PhysicalRowSlice:
-    """A half-open physical row slice containing visible rows."""
+    """File-local inclusive ranges containing only visible rows."""
 
-    start_inclusive: int
-    end_exclusive: int
+    row_ranges: List[Range]
     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,
-        )
+    def to_closed_row_id_ranges(self, first_row_id: int) -> List[Range]:
+        return [
+            Range(first_row_id + row_range.from_, first_row_id + row_range.to)
+            for row_range in self.row_ranges
+        ]
 
 
 class _LiveRowRangeSlicer:
@@ -91,35 +89,49 @@ def take(self, expected_live_rows: int) -> Optional[_PhysicalRowSlice]:
         if self._physical_position >= self._physical_row_count:
             return None
 
-        start = self._physical_position
         live_rows = 0
+        row_ranges = []
 
         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
+                if take > 0:
+                    row_ranges.append(Range(
+                        self._physical_position,
+                        self._physical_position + take - 1,
+                    ))
+                    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
+                    if needed > 0:
+                        row_ranges.append(Range(
+                            self._physical_position,
+                            self._physical_position + needed - 1,
+                        ))
+                        self._physical_position += needed
+                        live_rows += needed
                 else:
-                    self._physical_position += live_run
-                    live_rows += live_run
+                    if live_run > 0:
+                        row_ranges.append(Range(
+                            self._physical_position,
+                            self._physical_position + live_run - 1,
+                        ))
+                        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 _PhysicalRowSlice(
-                    start,
-                    self._physical_position,
+                    row_ranges,
                     live_rows,
                 )
 
@@ -129,7 +141,7 @@ def take(self, expected_live_rows: int) -> Optional[_PhysicalRowSlice]:
 
         if live_rows == 0:
             return None
-        return _PhysicalRowSlice(start, self._physical_position, live_rows)
+        return _PhysicalRowSlice(row_ranges, live_rows)
 
     def _skip_deleted_positions_at_cursor(self) -> None:
         while (
@@ -187,7 +199,7 @@ class ChunkShuffleSplitGeneratorBase(AbstractSplitGenerator):
       7. Map each chunk through :meth:`_chunk_to_split`.
 
     Subclasses implement the three abstract hooks. Chunks ride on existing
-    reader wrappers (``SlicedSplit`` / ``IndexedSplit``).
+    reader wrapper (``IndexedSplit``).
     """
 
     def __init__(
@@ -323,16 +335,9 @@ def _live_row_slicer(
 
 @dataclass
 class _FileSegment:
-    """A contiguous slice of a data file inside one chunk.
-
-    start/end are half-open row offsets within the file when the chunk
-    boundary falls inside the file; both are None when the chunk owns
-    the full file (so SlicedSplit's shard_file_idx_map can skip it and
-    treat the file as full — see sliced_split.py:73-78).
-    """
+    """Visible file-local row ranges from a data file inside one chunk."""
     file: DataFileMeta
-    start: Optional[int]
-    end: Optional[int]
+    row_ranges: List[Range]
     live_row_count: int
 
 
@@ -381,27 +386,13 @@ def _slice_group_into_chunks(
                 if physical_slice is None:
                     break
 
-                if (
-                    physical_slice.start_inclusive == 0
-                    and physical_slice.end_exclusive == file.row_count
-                ):
-                    current.append(
-                        _FileSegment(
-                            file,
-                            None,
-                            None,
-                            physical_slice.live_row_count,
-                        )
-                    )
-                else:
-                    current.append(
-                        _FileSegment(
-                            file,
-                            physical_slice.start_inclusive,
-                            physical_slice.end_exclusive,
-                            physical_slice.live_row_count,
-                        )
+                current.append(
+                    _FileSegment(
+                        file,
+                        physical_slice.row_ranges,
+                        physical_slice.live_row_count,
                     )
+                )
 
                 current_rows += physical_slice.live_row_count
 
@@ -412,11 +403,26 @@ def _slice_group_into_chunks(
 
     def _chunk_to_split(self, chunk: _Chunk) -> Split:
         files: List[DataFileMeta] = []
-        shard_file_idx_map = {}
+        row_ranges = []
+        ranges_use_row_ids = self.table.options.row_tracking_enabled()
+        split_offset = 0
         for seg in chunk.segments:
             files.append(seg.file)
-            if seg.start is not None and seg.end is not None:
-                shard_file_idx_map[seg.file.file_name] = (seg.start, seg.end)
+            if ranges_use_row_ids:
+                if seg.file.first_row_id is None:
+                    raise ValueError(
+                        "Row-tracked file '%s' is missing first_row_id"
+                        % seg.file.file_name
+                    )
+                range_base = seg.file.first_row_id
+            else:
+                range_base = split_offset
+            row_ranges.extend(
+                Range(range_base + row_range.from_,
+                      range_base + row_range.to)
+                for row_range in seg.row_ranges
+            )
+            split_offset += seg.file.row_count
 
         # set_file_path is already done once per unique file in
         # ChunkShuffleSplitGeneratorBase.create_splits.
@@ -435,19 +441,11 @@ def _chunk_to_split(self, chunk: _Chunk) -> Split:
             snapshot_id=self.snapshot_id,
         )
 
-        exact_merged_row_count = sum(
-            seg.live_row_count for seg in chunk.segments
+        return IndexedSplit(
+            data_split,
+            Range.sort_and_merge_overlap(row_ranges, True),
+            scores=None,
         )
-        if (
-            shard_file_idx_map
-            or data_split.merged_row_count() != exact_merged_row_count
-        ):
-            return SlicedSplit(
-                data_split,
-                shard_file_idx_map,
-                exact_merged_row_count=exact_merged_row_count,
-            )
-        return data_split
 
 
 # ---------------------------------------------------------------------------
@@ -461,11 +459,11 @@ class _AlignedGroupSegment:
 
     ``files`` is the entire group (may include blob/vector siblings),
     so the reader sees every column file even when only a slice of the
-    group's row_id range lands in this chunk. ``row_range`` is the
-    inclusive global row_id range this segment owns.
+    group's row_id range lands in this chunk. ``row_ranges`` are the
+    inclusive global row-id ranges of visible rows this segment owns.
     """
     files: List[DataFileMeta]
-    row_range: Range
+    row_ranges: List[Range]
     live_row_count: int
 
 
@@ -554,11 +552,10 @@ def _slice_group_into_chunks(
                 physical_slice = slicer.take(avail)
                 if physical_slice is None:
                     break
-                seg_range = physical_slice.to_closed_row_id_range(first_row_id)
                 current.append(
                     _AlignedGroupSegment(
                         group_files,
-                        seg_range,
+                        physical_slice.to_closed_row_id_ranges(first_row_id),
                         physical_slice.live_row_count,
                     )
                 )
@@ -573,14 +570,14 @@ def _chunk_to_split(self, chunk: _Chunk) -> Split:
         segments = chunk.segments
         if len(segments) == 1:
             all_files = segments[0].files
-            row_ranges = [segments[0].row_range]
+            row_ranges = list(segments[0].row_ranges)
         else:
             all_files = []
             row_ranges = []
             for seg in segments:
                 all_files.extend(seg.files)
-                row_ranges.append(seg.row_range)
-            row_ranges.sort(key=lambda r: r.from_)
+                row_ranges.extend(seg.row_ranges)
+        row_ranges = Range.sort_and_merge_overlap(row_ranges, True)
 
         data_deletion_files = self._get_deletion_files_for_split(
             all_files,
@@ -599,9 +596,6 @@ def _chunk_to_split(self, chunk: _Chunk) -> Split:
             data_split,
             row_ranges,
             scores=None,
-            exact_merged_row_count=sum(
-                seg.live_row_count for seg in segments
-            ),
         )
 
     @staticmethod