[core] Rerank primary-key vector candidates with exact distances (#550)

Author: JunRuiLee

crates/paimon/src/table/pk_vector_data_file_reader.rs

@@ -145,7 +145,7 @@ impl DataFilePkVectorReaderFactory {
 /// `None`). The column must be a `FixedSizeList`/`List` of `Float32`; every
 /// non-null row's child slice must have exactly `dimension` elements. Mirrors
 /// the layout handling in `vector_search_builder`.
-fn append_batch_vectors(
+pub(crate) fn append_batch_vectors(
     batch: &arrow_array::RecordBatch,
     field_name: &str,
     dimension: usize,
@@ -206,6 +206,12 @@ fn append_batch_vectors(
         }
         let mut vector = Vec::with_capacity(dimension);
         for i in start..end {
+            if values.is_null(i) {
+                return Err(data_invalid(format!(
+                    "vector row {row} has a null element at index {}",
+                    i - start
+                )));
+            }
             vector.push(values.value(i));
         }
         out.push(Some(vector));
@@ -349,6 +355,35 @@ mod integration_tests {
         }
     }
 
+    /// A present (non-null) vector row whose child slice contains a NULL
+    /// element must fail loud rather than silently defaulting the element to
+    /// `0.0` and corrupting the distance.
+    #[test]
+    fn append_batch_vectors_fails_loud_on_null_element() {
+        let field = vector_field();
+        let read_fields = vec![field.clone()];
+        let arrow_schema = build_target_arrow_schema(&read_fields).unwrap();
+
+        // Row is present, but element index 1 in its child slice is NULL: [1.0, null].
+        let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), 2).with_field(Arc::new(
+            ArrowField::new("element", ArrowDataType::Float32, true),
+        ));
+        builder.values().append_value(1.0);
+        builder.values().append_null();
+        builder.append(true);
+        let vec_array = builder.finish();
+        let batch = RecordBatch::try_new(arrow_schema, vec![Arc::new(vec_array)]).unwrap();
+
+        let mut out: Vec<Option<Vec<f32>>> = Vec::new();
+        let err = append_batch_vectors(&batch, field.name(), 2, &mut out)
+            .expect_err("null child element must fail loud");
+        let msg = err.to_string();
+        assert!(
+            msg.contains("null") && msg.contains("element"),
+            "got: {msg}"
+        );
+    }
+
     /// Write a FixedSizeList<Float32, 2> vector column
     /// (`[1,2]`, NULL, `[3,4]`) as a parquet data file, build the factory over
     /// its split, preload via `create`, and assert the whole-file sequential

crates/paimon/src/table/pk_vector_indexed_split_read.rs

@@ -61,7 +61,7 @@ pub(crate) struct PkVectorIndexedSplit {
 /// positions, validating bounds and ordering. Ranges must be non-empty, each
 /// within `[0, row_count)`, strictly ascending and non-overlapping (touching
 /// ranges allowed). Expansion is inclusive `from..=to`.
-fn expand_ranges(ranges: &[RowRange], row_count: i64) -> crate::Result<Vec<i64>> {
+pub(crate) fn expand_ranges(ranges: &[RowRange], row_count: i64) -> crate::Result<Vec<i64>> {
     if ranges.is_empty() {
         return Err(data_invalid("indexed split must select at least one row"));
     }

crates/paimon/src/table/pk_vector_orchestrator.rs

@@ -112,6 +112,7 @@ pub(crate) struct PkVectorSearchSplit {
 /// the cross-bucket merge dimensions a lone `PkVectorSearchResult` lacks;
 /// `split_index` is the re-association handle back to
 /// `splits[split_index].data_split`.
+#[derive(Clone)]
 pub(crate) struct PkVectorCandidate {
     pub split_index: usize,
     pub partition: BinaryRow,
@@ -121,6 +122,14 @@ pub(crate) struct PkVectorCandidate {
     pub distance: f32,
 }
 
+/// Best-first indexed (approximate) and exact-fallback candidate lists, each
+/// already globally bounded. The indexed list may be over-fetched for a later
+/// exact rerank; the exact list is bounded to the caller's final limit.
+pub(crate) struct OrchestratorSearchResult {
+    pub(crate) indexed: Vec<PkVectorCandidate>,
+    pub(crate) exact: Vec<PkVectorCandidate>,
+}
+
 /// 5-level BEST_FIRST (smallest = best) key. Level 1 orders distance with
 /// `java_float_compare` so a NaN distance (e.g. from a non-finite stored vector
 /// under inner product) sorts last rather than winning Top-1. Level 2 uses the
@@ -145,6 +154,20 @@ fn global_top_k(mut candidates: Vec<PkVectorCandidate>, limit: usize) -> Vec<PkV
     candidates
 }
 
+/// Merge already-bounded indexed and exact candidate lists into a single
+/// best-first list truncated to `limit`. The final global Top-K over the union
+/// is what guarantees the merged result is independent of how the two inputs
+/// were individually bounded.
+pub(crate) fn merge_candidates(
+    indexed: Vec<PkVectorCandidate>,
+    exact: Vec<PkVectorCandidate>,
+    limit: usize,
+) -> Vec<PkVectorCandidate> {
+    let mut all = indexed;
+    all.extend(exact);
+    global_top_k(all, limit)
+}
+
 /// Group Top-K survivors by `(partition, bucket, data_file_name)`, re-associate
 /// each group's file to its real `DataFileMeta` + aligned deletion file in the
 /// source bucket split, and build one `PkVectorIndexedSplit` per file. Groups are
@@ -197,7 +220,15 @@ pub(crate) fn build_indexed_splits(
         }
 
         // Re-associate the file to its DataFileMeta + aligned deletion file.
-        let source = &splits[split_index].data_split;
+        let source = &splits
+            .get(split_index)
+            .ok_or_else(|| {
+                data_invalid(format!(
+                    "vector search hit references split index {split_index} out of range (splits: {})",
+                    splits.len()
+                ))
+            })?
+            .data_split;
         let file_idx = source
             .data_files()
             .iter()
@@ -301,10 +332,13 @@ impl PkVectorOrchestrator {
     }
 
     /// Run the eager per-bucket search + cross-bucket global Top-K and return the
-    /// best-first survivors (through the full 5-level tie-break, raw distance
-    /// preserved). The exact-reader factory is split-scoped: it receives the
-    /// current split index and split so a caller can build a reader keyed to the
-    /// specific split/file. `skip_exact_fallback` forwards to `bucket_search`.
+    /// indexed (approximate) and exact-fallback survivors as two separate
+    /// best-first lists (each through the full 5-level tie-break, raw distance
+    /// preserved). The indexed list is bounded to `indexed_limit` (over-fetched
+    /// for a later exact rerank); the exact list is bounded to `limit`. The
+    /// exact-reader factory is split-scoped: it receives the current split index
+    /// and split so a caller can build a reader keyed to the specific split/file.
+    /// `skip_exact_fallback` forwards to `bucket_search`.
     ///
     /// `residual_by_split`, when present, carries one per-file allow-list of
     /// physical row positions per split (indexed parallel to `splits`): only
@@ -320,6 +354,7 @@ impl PkVectorOrchestrator {
         query: &[f32],
         metric: VectorSearchMetric,
         limit: usize,
+        indexed_limit: usize,
         ann_searcher: Option<&dyn PkVectorAnnSearcher>,
         exact_reader_factory: &mut (dyn for<'s, 'f> FnMut(
             usize,
@@ -330,11 +365,14 @@ impl PkVectorOrchestrator {
         search_options: &HashMap<String, String>,
         skip_exact_fallback: bool,
         residual_by_split: Option<&[HashMap<String, RoaringTreemap>]>,
-    ) -> crate::Result<Vec<PkVectorCandidate>> {
+    ) -> crate::Result<OrchestratorSearchResult> {
         // Eager input-shape validation (Java checkArgument parity).
         if limit == 0 {
             return Err(data_invalid("vector search limit must be positive"));
         }
+        if indexed_limit == 0 {
+            return Err(data_invalid("vector indexed search limit must be positive"));
+        }
         if query.is_empty() {
             return Err(data_invalid("vector search query must not be empty"));
         }
@@ -346,8 +384,9 @@ impl PkVectorOrchestrator {
             }
         }
 
-        // Eager per-bucket search -> tagged candidates.
-        let mut candidates: Vec<PkVectorCandidate> = Vec::new();
+        // Eager per-bucket search -> tagged candidates, kept split by path.
+        let mut indexed_candidates: Vec<PkVectorCandidate> = Vec::new();
+        let mut exact_candidates: Vec<PkVectorCandidate> = Vec::new();
         for (split_index, split) in splits.iter().enumerate() {
             let dvs = build_bucket_dv_map(&self.reader, split).await?;
             // Wrap the split-scoped factory into bucket_search's per-file signature.
@@ -357,38 +396,41 @@ impl PkVectorOrchestrator {
                 exact_reader_factory(split_index, split, file)
             });
             let residual_ranges = residual_by_split.map(|per_split| &per_split[split_index]);
-            let results = bucket_search(
+            let result = bucket_search(
                 ann_searcher,
                 &split.ann_segments,
                 &split.active_files,
                 &dvs,
                 &mut bucket_factory,
                 query,
                 metric,
+                indexed_limit,
                 limit,
                 search_options,
                 skip_exact_fallback,
                 residual_ranges,
             )
             .await?;
-            for PkVectorSearchResult {
+            let tag = |PkVectorSearchResult {
+                           data_file_name,
+                           row_position,
+                           distance,
+                       }: PkVectorSearchResult| PkVectorCandidate {
+                split_index,
+                partition: split.data_split.partition().clone(),
+                bucket: split.data_split.bucket(),
                 data_file_name,
                 row_position,
                 distance,
-            } in results
-            {
-                candidates.push(PkVectorCandidate {
-                    split_index,
-                    partition: split.data_split.partition().clone(),
-                    bucket: split.data_split.bucket(),
-                    data_file_name,
-                    row_position,
-                    distance,
-                });
-            }
+            };
+            indexed_candidates.extend(result.indexed.into_iter().map(&tag));
+            exact_candidates.extend(result.exact.into_iter().map(&tag));
         }
 
-        Ok(global_top_k(candidates, limit))
+        Ok(OrchestratorSearchResult {
+            indexed: global_top_k(indexed_candidates, indexed_limit),
+            exact: global_top_k(exact_candidates, limit),
+        })
     }
 }
 
@@ -547,6 +589,18 @@ mod tests {
         assert!(survivors.is_empty());
     }
 
+    #[test]
+    fn merge_candidates_takes_global_top_k_over_union() {
+        // indexed and exact each already bounded; the union's global Top-K must be
+        // independent of how each side was bounded.
+        let indexed = vec![cand(0, 0, "f", 0, 0.1), cand(0, 0, "f", 1, 0.4)];
+        let exact = vec![cand(0, 0, "g", 0, 0.2)];
+        let out = merge_candidates(indexed, exact, 2);
+        assert_eq!(out.len(), 2);
+        assert_eq!(out[0].distance, 0.1);
+        assert_eq!(out[1].distance, 0.2); // 0.2 (exact) beats 0.4 (indexed)
+    }
+
     #[test]
     fn builds_two_splits_with_ascending_position_ordered_scores() {
         // One bucket, two files. file-a hits at global order [pos=10, pos=2];
@@ -628,6 +682,18 @@ mod tests {
         assert!(format!("{err:?}").contains("duplicate"), "got: {err:?}");
     }
 
+    #[test]
+    fn build_indexed_splits_fails_loud_on_out_of_range_split_index() {
+        // A malformed candidate whose split_index is beyond the splits slice must
+        // fail loud, not panic on a bare slice index.
+        let splits = vec![search_split(0, vec![data_file("f", 10)])];
+        let survivors = vec![cand(5, 0, "f", 0, 1.0)];
+        let err = build_indexed_splits(survivors, &splits, VectorSearchMetric::L2)
+            .map(|_| ())
+            .expect_err("out-of-range split index must error");
+        assert!(format!("{err:?}").contains("out of range"), "got: {err:?}");
+    }
+
     #[test]
     fn rejects_same_group_key_from_different_splits() {
         // Two buckets share (partition, bucket, file_name) but sit at different
@@ -949,19 +1015,23 @@ mod e2e_tests {
         // expects; the split index/split are unused here.
         let mut wrapped =
             as_split_factory(|_: usize, _: &PkVectorSearchSplit, f: &BucketActiveFile| factory(f));
-        let survivors = orch
+        let result = orch
             .search_candidates(
                 splits,
                 query,
                 metric,
                 limit,
+                limit,
                 ann,
                 &mut wrapped,
                 opts,
                 false,
                 None,
             )
             .await?;
+        // Merge the two bounded lists into the best-first survivors the
+        // materialization path expects.
+        let survivors = merge_candidates(result.indexed, result.exact, limit);
         let indexed_splits = build_indexed_splits(survivors, splits, metric)?;
         let mut out = Vec::new();
         for indexed in indexed_splits {
@@ -991,6 +1061,7 @@ mod e2e_tests {
                 &[0.0, 0.0],
                 VectorSearchMetric::L2,
                 0,
+                0,
                 None,
                 &mut factory,
                 &opts,
@@ -1020,6 +1091,7 @@ mod e2e_tests {
                 &[],
                 VectorSearchMetric::L2,
                 5,
+                5,
                 None,
                 &mut factory,
                 &opts,
@@ -1360,12 +1432,13 @@ mod e2e_tests {
             },
         );
         let opts = HashMap::new();
-        let cands = PkVectorOrchestrator::new(make_reader(file_io, table_path))
+        let result = PkVectorOrchestrator::new(make_reader(file_io, table_path))
             .search_candidates(
                 &[split],
                 &[0.0, 0.0],
                 VectorSearchMetric::L2,
                 2,
+                2,
                 None,
                 &mut factory,
                 &opts,
@@ -1374,6 +1447,8 @@ mod e2e_tests {
             )
             .await
             .unwrap();
+        // Merge the two bounded lists into the best-first survivors.
+        let cands = merge_candidates(result.indexed, result.exact, 2);
         // Best-first: pos1 (d=1), pos2 (d=4).
         assert_eq!(
             cands
@@ -1432,12 +1507,13 @@ mod e2e_tests {
         allowed.insert(2);
         let residual_by_split = vec![HashMap::from([("r.mosaic".to_string(), allowed)])];
         let opts = HashMap::new();
-        let cands = PkVectorOrchestrator::new(make_reader(file_io, table_path))
+        let result = PkVectorOrchestrator::new(make_reader(file_io, table_path))
             .search_candidates(
                 &[split],
                 &[0.0, 0.0],
                 VectorSearchMetric::L2,
                 3,
+                3,
                 None,
                 &mut factory,
                 &opts,
@@ -1446,6 +1522,8 @@ mod e2e_tests {
             )
             .await
             .unwrap();
+        // Merge the two bounded lists into the best-first survivors.
+        let cands = merge_candidates(result.indexed, result.exact, 3);
         // Best-first among allowed positions: pos2 (d=4) then pos0 (d=9).
         assert_eq!(
             cands
@@ -1493,6 +1571,7 @@ mod e2e_tests {
                 &[0.0, 0.0],
                 VectorSearchMetric::L2,
                 3,
+                3,
                 None,
                 &mut factory,
                 &opts,
@@ -1533,12 +1612,13 @@ mod e2e_tests {
             },
         );
         let opts = HashMap::new();
-        let cands = PkVectorOrchestrator::new(make_reader(file_io, table_path))
+        let result = PkVectorOrchestrator::new(make_reader(file_io, table_path))
             .search_candidates(
                 &[split],
                 &[0.0, 0.0],
                 VectorSearchMetric::L2,
                 2,
+                2,
                 None,
                 &mut factory,
                 &opts,
@@ -1547,7 +1627,8 @@ mod e2e_tests {
             )
             .await
             .unwrap();
-        assert!(cands.is_empty());
+        assert!(result.indexed.is_empty());
+        assert!(result.exact.is_empty());
     }
 
     #[test]