feat(table): plan a PK-vector search from a decoded bucket split

crates/paimon/src/arrow/format/mosaic.rs

@@ -877,6 +877,33 @@ mod tests {
             .await
     }
 
+    /// The byte ranges a read of `data` requests when limited to `row_selection`.
+    async fn read_ranges_with_row_selection(
+        data: Bytes,
+        read_fields: &[DataField],
+        row_selection: Option<Vec<RowRange>>,
+    ) -> crate::Result<Vec<Range<u64>>> {
+        let file_size = data.len() as u64;
+        let calls = Arc::new(Mutex::new(Vec::new()));
+        let _: Vec<RecordBatch> = MosaicFormatReader
+            .read_batch_stream(
+                Box::new(TrackingFileRead {
+                    data,
+                    calls: Arc::clone(&calls),
+                }),
+                file_size,
+                read_fields,
+                None,
+                None,
+                row_selection,
+            )
+            .await?
+            .try_collect()
+            .await?;
+        let ranges = calls.lock().unwrap().clone();
+        Ok(ranges)
+    }
+
     async fn read_ranges_with_predicates(
         data: Bytes,
         read_fields: &[DataField],
@@ -1298,6 +1325,37 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_row_selection_skips_unselected_row_group_reads() {
+        // Three row groups of two rows. A selection inside the last one must not
+        // fetch the column data of the first two: this is what makes a narrow
+        // engine-supplied row range cheaper than reading the file and discarding
+        // rows afterwards, and it is granular to a row group, not to a row.
+        let fields = data_fields();
+        let projected = vec![fields[0].clone()];
+        let data = multi_row_group_mosaic(vec!["id".to_string()]);
+
+        let all = read_ranges_with_row_selection(data.clone(), &projected, None)
+            .await
+            .unwrap();
+        let last_only =
+            read_ranges_with_row_selection(data, &projected, Some(vec![RowRange::new(4, 5)]))
+                .await
+                .unwrap();
+
+        assert!(
+            last_only.len() < all.len(),
+            "a selection in one row group must request fewer ranges than a full read: \
+             {last_only:?} vs {all:?}"
+        );
+        let selected_bytes: u64 = last_only.iter().map(|r| r.end - r.start).sum();
+        let all_bytes: u64 = all.iter().map(|r| r.end - r.start).sum();
+        assert!(
+            selected_bytes < all_bytes,
+            "and fewer bytes: {selected_bytes} vs {all_bytes}"
+        );
+    }
+
     #[tokio::test]
     async fn test_read_predicate_missing_stats_still_filters_rows() {
         let fields = data_fields();

crates/paimon/src/table/data_file_reader.rs

@@ -599,6 +599,36 @@ impl DataFileReader {
         data_fields: Option<Vec<DataField>>,
         dv: Option<Arc<DeletionVector>>,
         local_positions: Vec<i64>,
+    ) -> crate::Result<ArrowRecordBatchStream> {
+        self.read_single_file_stream_local_ranges(
+            split,
+            file_meta,
+            data_fields,
+            dv,
+            coalesce_positions_to_local_ranges(&local_positions),
+        )
+    }
+
+    /// As [`Self::read_single_file_stream_local`], but the selection is already a
+    /// list of file-local inclusive ranges: sorted ascending, non-overlapping, and
+    /// within `[0, file_meta.row_count)`. A caller that already holds ranges — an
+    /// engine-supplied bucket split does — hands them over directly rather than
+    /// expanding them into positions this would only coalesce back.
+    ///
+    /// The emitted rows are always exactly the selected ones, but what that saves is
+    /// the format's business, and it differs: mosaic skips a row group before
+    /// touching its column data, parquet skips pages through the offset index,
+    /// `.row` prunes blocks. Avro is the exception — its reader loads the whole file
+    /// and deserializes every record before applying the selection, so there a
+    /// narrow selection saves only what comes after decoding: Arrow column
+    /// materialization, and whatever the caller does per row.
+    pub(super) fn read_single_file_stream_local_ranges(
+        &self,
+        split: &DataSplit,
+        file_meta: DataFileMeta,
+        data_fields: Option<Vec<DataField>>,
+        dv: Option<Arc<DeletionVector>>,
+        local_ranges: Vec<RowRange>,
     ) -> crate::Result<ArrowRecordBatchStream> {
         // Local-position selection is only sound against a predicate-free reader: a
         // row-filtering predicate drops arbitrary selected rows and desyncs the
@@ -666,12 +696,10 @@ impl DataFileReader {
             }
         };
 
-        // Interpret `local_positions` directly as file-local ranges (no
-        // `to_local_row_ranges`, no `first_row_id`), then fold the DV in.
-        // `merge_row_selection` intersects the selection with the file's
-        // non-deleted ranges, so the reader emits exactly the selected, non-deleted
-        // rows in ascending physical order.
-        let local_ranges = coalesce_positions_to_local_ranges(&local_positions);
+        // Interpret the ranges directly as file-local (no `to_local_row_ranges`, no
+        // `first_row_id`), then fold the DV in. `merge_row_selection` intersects the
+        // selection with the file's non-deleted ranges, so the reader emits exactly
+        // the selected, non-deleted rows in ascending physical order.
         let row_selection =
             merge_row_selection(file_meta.row_count, dv.as_deref(), Some(&local_ranges));
 

crates/paimon/src/table/pk_vector_bucket_split.rs

@@ -126,6 +126,39 @@ impl BucketVectorPayload {
             .as_deref()
             .expect("a decoded payload always carries source metadata")
     }
+
+    /// Consume the payload into the pieces a planner needs, so its decoded metadata
+    /// moves out of the payload rather than being cloned out of it.
+    ///
+    /// Two decoded fields are deliberately left behind. `row_count` is the payload's
+    /// own row count, which the read path derives from the source metadata instead.
+    /// `deletion_vectors_ranges` belongs to deletion-vector index files -- Java
+    /// builds a vector payload through the overload that leaves it null, and a read
+    /// takes its deletion vectors from the bucket's data split -- so a value here
+    /// describes something this payload is not, and is ignored rather than applied.
+    #[cfg_attr(not(test), allow(dead_code))]
+    pub(crate) fn into_parts(self) -> BucketVectorPayloadParts {
+        BucketVectorPayloadParts {
+            index_type: self.index_type,
+            file_name: self.file_name,
+            file_size: self.file_size,
+            external_path: self.external_path,
+            global_index_meta: self.global_index_meta,
+        }
+    }
+}
+
+/// The owned pieces of a [`BucketVectorPayload`], produced by
+/// [`BucketVectorPayload::into_parts`].
+#[cfg_attr(not(test), allow(dead_code))]
+pub(crate) struct BucketVectorPayloadParts {
+    pub(crate) index_type: String,
+    pub(crate) file_name: String,
+    /// As decoded: Java writes a signed length, so a negative value is possible on
+    /// the wire and is rejected where it is converted, not here.
+    pub(crate) file_size: i64,
+    pub(crate) external_path: Option<String>,
+    pub(crate) global_index_meta: GlobalIndexMeta,
 }
 
 impl BucketVectorSearchSplit {
@@ -144,6 +177,19 @@ impl BucketVectorSearchSplit {
         &self.row_ranges_by_file
     }
 
+    /// Consume the split into its three parts, so a planner can take ownership of
+    /// the data split, the payloads and the row ranges without cloning them.
+    #[cfg_attr(not(test), allow(dead_code))]
+    pub(crate) fn into_parts(
+        self,
+    ) -> (
+        DataSplit,
+        Vec<BucketVectorPayload>,
+        IndexMap<String, Vec<RowRange>>,
+    ) {
+        (self.data_split, self.payload_files, self.row_ranges_by_file)
+    }
+
     /// Parse a Java `BucketVectorSearchSplit#serialize` message.
     ///
     /// Integers are big-endian and file names are Java modified UTF-8, following
@@ -454,6 +500,49 @@ fn read_count(cur: &mut &[u8], element: &str) -> crate::Result<usize> {
     Ok(count)
 }
 
+#[cfg(test)]
+impl BucketVectorSearchSplit {
+    /// Assemble a split directly, for tests that need shapes the decoder will not
+    /// produce -- a nested split that wrongly carries row ranges, two splits for one
+    /// bucket, a negative payload size. Production splits always come from
+    /// [`Self::deserialize`].
+    pub(crate) fn new_for_test(
+        data_split: DataSplit,
+        payload_files: Vec<BucketVectorPayload>,
+        row_ranges_by_file: IndexMap<String, Vec<RowRange>>,
+    ) -> Self {
+        Self {
+            data_split,
+            payload_files,
+            row_ranges_by_file,
+        }
+    }
+}
+
+#[cfg(test)]
+impl BucketVectorPayload {
+    #[allow(clippy::too_many_arguments)]
+    pub(crate) fn new_for_test(
+        index_type: &str,
+        file_name: &str,
+        file_size: i64,
+        row_count: i64,
+        deletion_vectors_ranges: Option<IndexMap<String, DeletionVectorMeta>>,
+        external_path: Option<String>,
+        global_index_meta: GlobalIndexMeta,
+    ) -> Self {
+        Self {
+            index_type: index_type.to_string(),
+            file_name: file_name.to_string(),
+            file_size,
+            row_count,
+            deletion_vectors_ranges,
+            external_path,
+            global_index_meta,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;