fix(index): tighten batch IVF gate, move ranking to CPU runtime, record partitions_searched

Three follow-ups on the shared-scan batch IVF path:

- Gate eligibility on whether the *selected* index_segments cover every
  requested fragment, not whether the whole logical index does. A subset
  selected via with_index_segments could otherwise let the batch node search
  only the selected segments and silently drop a fragment covered solely by an
  unselected segment. Extracted the coverage check into
  fragments_missing_from_index_segments, shared by the gate and knn_combined so
  eligibility and fallback stay in lockstep.

- Run the per-query centroid ranking on the dedicated CPU runtime
  (find_partitions_batch_on_cpu) instead of the async worker: the ranking is
  pure CPU and batch width multiplies it, so a wide batch over a large centroid
  set could monopolize a Tokio worker.

- Record partitions_searched on ANNIvfBatchExec. It built the metric but never
  incremented it, so EXPLAIN ANALYZE reported 0 for every batch query. Report
  the distinct partitions read -- the shared I/O this node exists to save --
  mirroring the single-query ANNIvfSubIndexExec.

Tests: partial-segment fallback, CPU-runtime ranking, and a
partitions_searched=2 assertion (distinct union, not the per-query sum of 4).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: sezruby

rust/lance/src/dataset/scanner.rs

@@ -5545,8 +5545,8 @@ impl Scanner {
     /// - no refine step (the batch path does not yet rerank);
     /// - fixed nprobes (`minimum_nprobes == maximum_nprobes`) — see below;
     /// - every segment an IVF index with a flat-style sub-index (i.e. not HNSW);
-    /// - all target fragments indexed (or `fast_search`, which ignores
-    ///   unindexed fragments).
+    /// - every target fragment covered by the *selected* `index_segments` (or
+    ///   `fast_search`, which searches only the selected segments anyway).
     ///
     /// The fixed-nprobes requirement is a *correctness* gate, not just an
     /// optimization. The shared-scan path searches exactly `minimum_nprobes`
@@ -5625,11 +5625,47 @@ impl Scanner {
         if self.fast_search {
             return Ok(true);
         }
-        // The batch node only searches indexed partitions, so any unindexed
-        // target fragment would silently drop rows; fall back in that case.
-        let unindexed_fragments =
-            self.retain_target_fragments(self.dataset.unindexed_fragments(index_name).await?);
-        Ok(unindexed_fragments.is_empty())
+        // The batch node only searches the selected `index_segments`, so any
+        // target fragment those segments do not cover would silently drop rows
+        // (the single-query path re-scores such fragments on a flat fallback in
+        // `knn_combined`). Measure coverage against the selected segments -- not
+        // the whole logical index -- so a subset selected via
+        // `with_index_segments` cannot hide a fragment that an unselected
+        // segment happens to cover; fall back whenever any remain.
+        let uncovered_fragments = self
+            .fragments_missing_from_index_segments(index_name, index_segments)
+            .await?;
+        Ok(uncovered_fragments.is_empty())
+    }
+
+    /// Target fragments the given `index_segments` do not cover.
+    ///
+    /// The ANN scan reads only the selected segments' partitions, so these are
+    /// exactly the fragments the single-query path re-scores on a flat fallback
+    /// in [`Self::knn_combined`]. Coverage is measured against the *selected*
+    /// segments rather than every segment of the logical index (which
+    /// `Dataset::unindexed_fragments` would do): a caller may select a subset
+    /// via [`with_index_segments`](Self::with_index_segments) while another,
+    /// unselected segment covers one of the requested fragments.
+    async fn fragments_missing_from_index_segments(
+        &self,
+        index_name: &str,
+        index_segments: &[IndexMetadata],
+    ) -> Result<Vec<Fragment>> {
+        if let Some(target_fragments) = &self.fragments {
+            let indexed_fragments = self.get_indexed_frags(index_segments);
+            Ok(target_fragments
+                .iter()
+                .filter(|fragment| !indexed_fragments.contains(fragment.id as u32))
+                .cloned()
+                .collect())
+        } else if self.index_segments.is_some() {
+            // An explicit segment selection with no fragment restriction searches
+            // exactly those segments; there is nothing to fall back for.
+            Ok(Vec::new())
+        } else {
+            self.dataset.unindexed_fragments(index_name).await
+        }
     }
 
     async fn batch_indexed_vector_search(
@@ -5748,18 +5784,9 @@ impl Scanner {
         mut knn_node: Arc<dyn ExecutionPlan>,
         filter_plan: &ExprFilterPlan,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        let fallback_fragments = if let Some(target_fragments) = &self.fragments {
-            let indexed_fragments = self.get_indexed_frags(indexed_segments);
-            target_fragments
-                .iter()
-                .filter(|fragment| !indexed_fragments.contains(fragment.id as u32))
-                .cloned()
-                .collect::<Vec<_>>()
-        } else if self.index_segments.is_some() {
-            Vec::new()
-        } else {
-            self.dataset.unindexed_fragments(index_name).await?
-        };
+        let fallback_fragments = self
+            .fragments_missing_from_index_segments(index_name, indexed_segments)
+            .await?;
 
         let has_fallback = !fallback_fragments.is_empty();
         let has_stale = !stale_rows.is_empty();
@@ -9773,6 +9800,23 @@ mod test {
             plan
         );
 
+        // The batch node loads each probed partition once and scores every query
+        // that probes it, so it must report the *distinct* partitions read: with
+        // 2 partitions and nprobes(2), both queries probe both partitions, so the
+        // union is 2 -- not the per-query sum (2 queries x 2 = 4), and never 0
+        // (which is what a dropped metric would show). This guards the observed
+        // `partitions_searched` against silently regressing to either.
+        let analyzed = scan.analyze_plan().await.unwrap();
+        let batch_line = analyzed
+            .lines()
+            .find(|line| line.contains("ANNIvfBatch"))
+            .expect("analyzed plan should contain the ANNIvfBatch node");
+        assert!(
+            batch_line.contains("partitions_searched=2"),
+            "batch node must report the distinct partitions searched, got:\n{}",
+            batch_line
+        );
+
         let batch = scan.try_into_batch().await.unwrap();
         assert_query_index_field(&batch);
         assert_eq!(
@@ -10056,6 +10100,80 @@ mod test {
         }
     }
 
+    /// The shared-scan fast path is only equivalent to repeated single-query
+    /// search when the selected `index_segments` cover every requested fragment.
+    /// With one segment per fragment, requesting both fragments but selecting
+    /// only the first segment leaves fragment 1 covered solely by the
+    /// *unselected* segment: the batch node would search just the selected
+    /// segment and silently drop it. Eligibility must be computed from the
+    /// selected segments' coverage (as `knn_combined` does), not the whole
+    /// logical index, so the scanner falls back to the per-query loop, which
+    /// re-scores the uncovered fragment on the flat path and returns every row.
+    #[tokio::test]
+    async fn test_batch_knn_indexed_partial_segment_selection_falls_back() {
+        let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false)
+            .await
+            .unwrap();
+        // One segment per fragment: segment_ids[0] covers fragment 0 (i=0..200),
+        // segment_ids[1] covers fragment 1 (i=200..400).
+        let segment_ids = test_ds.make_segmented_vector_index().await.unwrap();
+        let dataset = &test_ds.dataset;
+        let fragments = dataset.fragments();
+        assert_eq!(fragments.len(), 2, "base dataset should have two fragments");
+
+        let (queries, _query_values) = batch_knn_two_queries();
+        // k covers every row in both requested fragments (200 each), so a
+        // complete search returns 400 rows per query.
+        let k = 400;
+
+        let mut scan = dataset.scan();
+        scan.nearest("vec", &queries, k).unwrap();
+        scan.nprobes(2);
+        // Request both indexed fragments but select only the segment covering
+        // fragment 0; fragment 1 is covered only by the unselected segment.
+        scan.with_fragments(vec![fragments[0].clone(), fragments[1].clone()]);
+        scan.with_index_segments(vec![segment_ids[0]]).unwrap();
+        scan.project(&["i"]).unwrap();
+
+        let plan = scan.explain_plan(false).await.unwrap();
+        assert!(
+            !plan.contains("ANNIvfBatch"),
+            "a requested fragment outside the selected segments must force a fallback, \
+             not a shared scan that drops it, got:\n{plan}"
+        );
+
+        let batch = scan.try_into_batch().await.unwrap();
+        assert_query_index_field(&batch);
+        assert_eq!(
+            batch.num_rows(),
+            2 * k,
+            "each query must return all 400 rows across both requested fragments"
+        );
+        let query_indices = batch[QUERY_INDEX_COL].as_primitive::<Int32Type>();
+        for query_index in 0..2 {
+            let rows_for_query = query_indices
+                .iter()
+                .filter(|value| *value == Some(query_index))
+                .count();
+            assert_eq!(
+                rows_for_query, k,
+                "query_index {query_index} must cover both fragments (got {rows_for_query})"
+            );
+        }
+        // Fragment 0 (i in 0..200) comes from the selected segment; fragment 1
+        // (i in 200..400) must appear via the flat fallback.
+        let i_array = batch["i"].as_primitive::<Int32Type>();
+        assert!(
+            i_array
+                .iter()
+                .any(|v| v.is_some_and(|val| (0..200).contains(&val)))
+                && i_array
+                    .iter()
+                    .any(|v| v.is_some_and(|val| (200..400).contains(&val))),
+            "results must include rows from both the selected segment and the flat-fallback fragment"
+        );
+    }
+
     /// A wide batch probes more distinct partitions than one streaming chunk holds
     /// (`STREAMING_SEARCH_BATCH_SIZE` = 16), so `search_partitions_batch` scores
     /// them in several `spawn_cpu` dispatches, threading the per-query top-k heaps

rust/lance/src/io/exec/knn.rs

@@ -126,6 +126,58 @@ async fn find_partitions_on_cpu(
         .map_err(|e| DataFusionError::Execution(format!("Failed to find partitions: {}", e)))
 }
 
+/// Per-query IVF partition rankings: for each query, its probed-partition ids
+/// and the corresponding centroid distances, in query order.
+type BatchPartitionRankings = (Vec<Arc<UInt32Array>>, Vec<Arc<Float32Array>>);
+
+/// Rank every query vector in a batch against the IVF centroids on the CPU
+/// runtime.
+///
+/// [`VectorIndex::find_partitions`] is pure CPU work, and a wide batch over a
+/// large centroid set multiplies it enough to monopolize a Tokio worker and
+/// stall unrelated async progress. Dispatch the whole ranking loop as a single
+/// `spawn_cpu` job -- mirroring the single-query [`find_partitions_on_cpu`] --
+/// so it stays off the async executor threads.
+///
+/// `query.key` holds all `query_count` vectors concatenated and `dim` is the
+/// per-vector width. Returns each query's probed-partition list and the
+/// corresponding centroid distances, in query order.
+async fn find_partitions_batch_on_cpu(
+    index: Arc<dyn VectorIndex>,
+    query: Query,
+    query_count: usize,
+    dim: usize,
+) -> DataFusionResult<BatchPartitionRankings> {
+    spawn_cpu(move || -> Result<BatchPartitionRankings> {
+        let mut partitions_per_query = Vec::with_capacity(query_count);
+        let mut dists_per_query = Vec::with_capacity(query_count);
+        for query_index in 0..query_count {
+            let mut single_query = query.clone();
+            single_query.key = query.key.slice(query_index * dim, dim);
+            // Probe a fixed number of partitions per query. The scanner only
+            // routes here when `minimum_nprobes == maximum_nprobes` and
+            // `minimum_nprobes > 0` (see `Scanner::batch_index_search_supported`),
+            // so this is exactly what the single-query path would search -- no
+            // adaptive `early_pruning` floor or late-search expansion applies,
+            // making the batch result identical to repeated single-query search.
+            // No clamp is needed: the gate rejects `nprobes(0)` (which the
+            // single-query path treats as "probe nothing") rather than silently
+            // searching one partition here.
+            debug_assert!(
+                single_query.minimum_nprobes > 0,
+                "batch node reached with nprobes(0); the scanner gate should have fallen back"
+            );
+            single_query.maximum_nprobes = Some(single_query.minimum_nprobes);
+            let (partitions, q_c_dists) = index.find_partitions(&single_query)?;
+            partitions_per_query.push(Arc::new(partitions));
+            dists_per_query.push(Arc::new(q_c_dists));
+        }
+        Ok((partitions_per_query, dists_per_query))
+    })
+    .await
+    .map_err(|e| DataFusionError::Execution(format!("Failed to find partitions: {e}")))
+}
+
 fn normalize_query_for_index(index: &dyn VectorIndex, query: Query) -> DataFusionResult<Query> {
     if index.metric_type() != DistanceType::Cosine {
         return Ok(query);
@@ -2620,30 +2672,32 @@ impl ExecutionPlan for ANNIvfBatchExec {
                     dim,
                 )?;
 
-                let mut partitions_per_query = Vec::with_capacity(query_count);
-                let mut dists_per_query = Vec::with_capacity(query_count);
-                for query_index in 0..query_count {
-                    let mut single_query = normalized.clone();
-                    single_query.key = normalized.key.slice(query_index * dim, dim);
-                    // Probe a fixed number of partitions per query. The scanner
-                    // only routes here when `minimum_nprobes == maximum_nprobes`
-                    // and `minimum_nprobes > 0` (see
-                    // `Scanner::batch_index_search_supported`), so this is exactly
-                    // what the single-query path would search — no adaptive
-                    // `early_pruning` floor or late-search expansion applies,
-                    // making the batch result identical to repeated single-query
-                    // search. No clamp is needed: the gate rejects `nprobes(0)`
-                    // (which the single-query path treats as "probe nothing")
-                    // rather than silently searching one partition here.
-                    debug_assert!(
-                        single_query.minimum_nprobes > 0,
-                        "batch node reached with nprobes(0); the scanner gate should have fallen back"
-                    );
-                    single_query.maximum_nprobes = Some(single_query.minimum_nprobes);
-                    let (partitions, q_c_dists) = index.find_partitions(&single_query)?;
-                    partitions_per_query.push(Arc::new(partitions));
-                    dists_per_query.push(Arc::new(q_c_dists));
-                }
+                // Rank every query vector against the IVF centroids on the CPU
+                // runtime rather than inside this async future: the ranking is
+                // pure CPU and the batch width multiplies it, so a wide batch
+                // over a large centroid set could otherwise monopolize a Tokio
+                // worker. See `find_partitions_batch_on_cpu`.
+                let (partitions_per_query, dists_per_query) = find_partitions_batch_on_cpu(
+                    index.clone(),
+                    normalized.clone(),
+                    query_count,
+                    dim,
+                )
+                .await?;
+
+                // Record the partitions this delta actually reads. The batch
+                // node loads each probed partition once and scores every query
+                // that probes it, so the honest "partitions searched" count is
+                // the union across queries -- the shared I/O this node exists to
+                // save -- not the per-query sum. Mirrors the single-query
+                // ANNIvfSubIndexExec, which also records PARTITIONS_SEARCHED.
+                let distinct_partitions: RoaringBitmap = partitions_per_query
+                    .iter()
+                    .flat_map(|parts| parts.values().iter().copied())
+                    .collect();
+                metrics
+                    .partitions_searched
+                    .add(distinct_partitions.len() as usize);
 
                 let index_metrics: Arc<dyn MetricsCollector> =
                     Arc::new(metrics.index_metrics.clone());
@@ -3692,6 +3746,33 @@ mod tests {
         );
     }
 
+    // Batch analogue of `test_find_partitions_runs_on_cpu_runtime`: the batch
+    // node's per-query centroid ranking multiplies the CPU cost, so it must also
+    // run on the dedicated cpu runtime rather than a Tokio async worker.
+    #[tokio::test]
+    async fn test_find_partitions_batch_runs_on_cpu_runtime() {
+        let thread_name = Arc::new(Mutex::new(None));
+        let index: Arc<dyn VectorIndex> = Arc::new(ThreadCapturingIndex {
+            thread_name: thread_name.clone(),
+            row_ids: Vec::new(),
+        });
+
+        // Two query vectors of dim 1 concatenated into one key.
+        let mut query = base_query();
+        query.key = Arc::new(Float32Array::from(vec![0.0f32, 1.0f32]));
+        let (partitions, dists) = find_partitions_batch_on_cpu(index, query, 2, 1)
+            .await
+            .unwrap();
+        assert_eq!(partitions.len(), 2, "one partition list per query");
+        assert_eq!(dists.len(), 2, "one distance list per query");
+
+        let thread_name = thread_name.lock().unwrap().clone().unwrap();
+        assert!(
+            thread_name.contains("lance-cpu"),
+            "expected batch find_partitions to run on the dedicated cpu runtime, got thread {thread_name}",
+        );
+    }
+
     // All partitions fit in a single search batch, so they are searched in one
     // `spawn_cpu` dispatch and therefore share one cpu thread. The partition count
     // adapts to the configured batch size so the single-batch property holds under