feat(vindex): search multiple primary-key vector queries in one pass

Add a batch entry to the bucket-local vector search so N query vectors share
one pass over each ANN segment and each uncovered data file. The ANN searcher
gains search_batch: it builds the live-row mask once (query-independent), opens
one reader per segment, and drives visit_batch_vector_search, mapping results
per query. The bucket search gains bucket_search_batch with per-query bounded
heaps; a single query short-circuits to the existing single-query path so its
result stays identical. The exact-fallback per-file search already scored all
queries in one stream pass. Single-query wrappers are retained throughout.

Author: JunRuiLee

crates/paimon/src/table/pk_vector_data_file_reader.rs

@@ -519,6 +519,60 @@ mod integration_tests {
         assert_eq!(streamed[0][1].row_position, 3);
     }
 
+    /// A multi-query `search_file` returns independent per-query Top-K lists: the
+    /// slot for a query in a batch is identical to that query searched alone (no
+    /// cross-query bleed), and a shared `is_excluded` applies to every query.
+    #[tokio::test]
+    async fn search_file_multi_query_returns_independent_per_query_top_k() {
+        let rows = vec![
+            Some(vec![0.0, 0.0]),
+            Some(vec![1.0, 0.0]),
+            Some(vec![2.0, 0.0]),
+            Some(vec![3.0, 0.0]),
+        ];
+        let (factory, file_name) =
+            build_factory(&rows, rows.len() as i64, "memory:/pkvdfr_multiquery").await;
+        let active = BucketActiveFile {
+            file_name,
+            row_count: rows.len() as i64,
+        };
+        // Exclude physical position 1 for all queries (shared predicate).
+        let is_excluded = |pos: i64| pos == 1;
+        let q0 = [0.0f32, 0.0]; // nearest is pos 0
+        let q1 = [3.0f32, 0.0]; // nearest is pos 3
+
+        let batch = factory
+            .search_file(
+                &active,
+                &[&q0, &q1],
+                VectorSearchMetric::L2,
+                2,
+                &is_excluded,
+            )
+            .await
+            .unwrap();
+        assert_eq!(batch.len(), 2, "one result list per query");
+
+        // Each query searched alone must equal its slot in the batch.
+        let only_q0 = factory
+            .search_file(&active, &[&q0], VectorSearchMetric::L2, 2, &is_excluded)
+            .await
+            .unwrap();
+        let only_q1 = factory
+            .search_file(&active, &[&q1], VectorSearchMetric::L2, 2, &is_excluded)
+            .await
+            .unwrap();
+        assert_eq!(batch[0], only_q0[0]);
+        assert_eq!(batch[1], only_q1[0]);
+
+        // Sanity: distinct nearest neighbours, and the excluded position is absent.
+        assert_eq!(batch[0][0].row_position, 0);
+        assert_eq!(batch[1][0].row_position, 3);
+        assert!(batch
+            .iter()
+            .all(|list| list.iter().all(|r| r.row_position != 1)));
+    }
+
     /// A `DataFileMeta.row_count` larger than the file's real row count means the
     /// stream ends early; the search must fail loud rather than return a short
     /// result.

crates/paimon/src/table/pk_vector_orchestrator.rs

@@ -1100,18 +1100,18 @@ mod e2e_tests {
         hits: Vec<PkVectorSearchResult>,
     }
     impl PkVectorAnnSearcher for FakeAnn {
-        fn search(
+        fn search_batch(
             &self,
             _segment: &BucketAnnSegment,
-            _query: &[f32],
+            queries: &[&[f32]],
             _metric: VectorSearchMetric,
             _limit: usize,
             _active_source_files: &HashSet<String>,
             _dvs: &HashMap<String, Arc<DeletionVector>>,
             _opts: &HashMap<String, String>,
             _residual_ranges: Option<&HashMap<String, roaring::RoaringTreemap>>,
-        ) -> crate::Result<Vec<PkVectorSearchResult>> {
-            Ok(self.hits.clone())
+        ) -> crate::Result<Vec<Vec<PkVectorSearchResult>>> {
+            Ok(queries.iter().map(|_| self.hits.clone()).collect())
         }
     }
 

crates/paimon/src/table/vector_search_builder.rs

@@ -493,7 +493,9 @@ impl<'a> VectorSearchBuilder<'a> {
         );
 
         // Real ANN scorer: preload each segment's bytes (keyed by resolved,
-        // globally unique path) and drive the vindex reader from memory.
+        // globally unique path) and drive the vindex reader from memory. The reader
+        // is opened once per segment and every query in the batch is searched
+        // against it, mirroring the shared-reader batch search.
         let segment_bytes = preload_segment_bytes(self.table.file_io(), &plan.splits).await?;
         // Fail loud on a config/segment metric mismatch before scoring, mirroring
         // Java `PkVectorAnnSegmentSearcher.search`.
@@ -505,8 +507,8 @@ impl<'a> VectorSearchBuilder<'a> {
         };
         let search_options = options.clone();
         let field_name = pk_col.to_string();
-        let scorer: crate::vindex::pkvector::ann::Scorer =
-            Box::new(move |segment: &BucketAnnSegment, search: &VectorSearch| {
+        let scorer: crate::vindex::pkvector::ann::BatchScorer = Box::new(
+            move |segment: &BucketAnnSegment, searches: &[VectorSearch]| {
                 let data = segment_bytes
                     .get(&segment.path)
                     .ok_or_else(|| crate::Error::DataInvalid {
@@ -523,15 +525,16 @@ impl<'a> VectorSearchBuilder<'a> {
                     VectorIndexBackend::Lumina => {
                         let mut reader =
                             LuminaVectorGlobalIndexReader::new(io_meta, options.clone());
-                        reader.visit_vector_search(search, |_| Ok(Cursor::new(data)))
+                        reader.visit_batch_vector_search(searches, |_| Ok(Cursor::new(data)))
                     }
                     VectorIndexBackend::Vindex => {
                         let mut reader =
                             VindexVectorGlobalIndexReader::new(io_meta, options.clone());
-                        reader.visit_vector_search(search, |_| Ok(Cursor::new(data)))
+                        reader.visit_batch_vector_search(searches, |_| Ok(Cursor::new(data)))
                     }
                 }
-            });
+            },
+        );
         let ann_searcher = VindexAnnSearcher::new(field_name, scorer);
 
         // Residual (post-recall) filtering: for each candidate file, re-read its