feat(vector): add partition search parallelism (#6475)

## Feature

This PR adds query-time `query_parallelism` for vector search partition
execution and wires it through Rust, Python, and Java.

Callers can control how many IVF partitions a single query may search
concurrently:
- `0`: auto policy. The current implementation maps auto to the
single-worker sequential path.
- `1`: single-worker sequential partition search.
- `-1`: use the CPU pool size.
- `>= 2`: partition-parallel search, clamped to the CPU pool size.

The default is `0`, which currently preserves the optimized sequential
execution path.

## Performance Improvement

The performance issue is per-query worker fan-out. With many concurrent
queries, spawning a CPU task for every partition of every query
increases contention on the CPU worker pool and lengthens queueing time.
This is especially visible for fixed-`nprobes` IVF workloads where every
query searches the same number of partitions.

This PR improves that path by making query partition scheduling
configurable while keeping the optimized sequential model as the default
auto behavior. Sequential fixed-`nprobes` search prepares partitions
asynchronously, then searches prepared partitions on one CPU worker
using a query-level global top-k heap. Parallel execution remains
available for workloads that benefit from intra-query partition
parallelism.

## Implementation

- Rust, Python, and Java expose `query_parallelism` on vector search
APIs.
- `ANNIvfSubIndexExec` converts the configured value into an effective
partition concurrency.
- Effective partition concurrency is clamped to the CPU pool size.
- IVF v2 splits partition search into async prepare and sync execute
phases.
- Sequential fixed-`nprobes` search uses a query-level global top-k
heap.
- Sequential late-search keeps the existing per-partition output shape
so early-stop behavior is preserved.
- Parallel execution uses direct per-partition search tasks, matching
the original execution model.
- IVF_RQ/Flat sub-index search reuses caller-owned scratch buffers for
RQ distance-table quantization and top-k accumulation to reduce
allocation overhead.

## Developer Impact

- Rust callers can call `Scanner::query_parallelism(...)`.
- Python callers can pass `query_parallelism` in vector search APIs.
- Java callers can use `Query.Builder#setQueryParallelism(...)`.
- `parallel_mode` / `ParallelMode` have been replaced by the
concurrency-based API.

## Benchmark

GCP VM benchmark using the same index for all modes. All rows had
matching result checksums.

Configuration: `gist / IVF_RQ / target_partition_size=8192 / k=100 /
nprobes=20 / columns=[] / prewarm / max_queries=1000`. Percentages are
relative to `main`.

| Threads | Mode | Avg | P50 | P90 | P95 | P99 | QPS |
|---:|---|---:|---:|---:|---:|---:|---:|
| 8 | main | 4.98 ms | 4.92 ms | 6.28 ms | 6.86 ms | 7.56 ms | 1584.6 |
| 8 | sequential | 4.27 ms (-14.2%) | 4.21 ms (-14.5%) | 5.03 ms
(-19.9%) | 5.27 ms (-23.1%) | 5.97 ms (-21.0%) | 1838.9 (+16.1%) |
| 8 | parallel | 5.02 ms (+0.7%) | 4.97 ms (+0.9%) | 6.27 ms (-0.2%) |
6.60 ms (-3.7%) | 7.36 ms (-2.7%) | 1575.4 (-0.6%) |
| 16 | main | 9.95 ms | 9.74 ms | 13.12 ms | 14.11 ms | 16.83 ms |
1583.1 |
| 16 | sequential | 8.09 ms (-18.6%) | 7.98 ms (-18.0%) | 9.84 ms
(-25.0%) | 10.47 ms (-25.8%) | 11.67 ms (-30.6%) | 1936.8 (+22.3%) |
| 16 | parallel | 9.95 ms (+0.0%) | 9.68 ms (-0.6%) | 13.37 ms (+1.9%) |
14.42 ms (+2.2%) | 16.83 ms (+0.0%) | 1583.6 (+0.0%) |
| 32 | main | 18.68 ms | 18.16 ms | 26.66 ms | 28.77 ms | 33.12 ms |
1652.7 |
| 32 | sequential | 15.50 ms (-17.0%) | 15.15 ms (-16.6%) | 20.57 ms
(-22.9%) | 22.16 ms (-23.0%) | 25.32 ms (-23.6%) | 2000.6 (+21.1%) |
| 32 | parallel | 19.13 ms (+2.4%) | 18.58 ms (+2.3%) | 26.49 ms (-0.6%)
| 29.12 ms (+1.2%) | 33.92 ms (+2.4%) | 1625.9 (-1.6%) |
| 64 | main | 33.98 ms | 33.01 ms | 49.65 ms | 53.87 ms | 63.58 ms |
1718.4 |
| 64 | sequential | 29.95 ms (-11.8%) | 29.67 ms (-10.1%) | 42.44 ms
(-14.5%) | 46.81 ms (-13.1%) | 54.42 ms (-14.4%) | 1949.4 (+13.4%) |
| 64 | parallel | 35.17 ms (+3.5%) | 34.37 ms (+4.1%) | 50.70 ms (+2.1%)
| 55.04 ms (+2.2%) | 69.09 ms (+8.7%) | 1650.7 (-3.9%) |
| 128 | main | 38.73 ms | 36.07 ms | 68.28 ms | 76.78 ms | 96.81 ms |
1663.3 |
| 128 | sequential | 37.48 ms (-3.2%) | 35.90 ms (-0.5%) | 65.26 ms
(-4.4%) | 71.82 ms (-6.5%) | 87.48 ms (-9.6%) | 1915.3 (+15.2%) |
| 128 | parallel | 41.29 ms (+6.6%) | 39.38 ms (+9.2%) | 71.68 ms
(+5.0%) | 80.70 ms (+5.1%) | 96.20 ms (-0.6%) | 1577.2 (-5.2%) |

## Validation

- `cargo fmt --all --check`
- `cargo check -p lance --tests`
- `cd python && cargo check`
- `cd java && cargo check --manifest-path ./lance-jni/Cargo.toml`
- `uv run --with maturin maturin develop`
- `uv run pytest
python/tests/test_vector_index.py::test_vector_index_with_query_parallelism
python/tests/test_vector_index.py::test_vector_index_invalid_query_parallelism`
- Python ruff check / format check for touched Python files

Author: BubbleCal

java/lance-jni/src/blocking_scanner.rs

@@ -320,6 +320,11 @@ pub(crate) fn build_scanner_with_options<'a>(
 
         let use_index = env.get_boolean_from_method(&java_obj, "isUseIndex")?;
         scanner.use_index(use_index);
+
+        let query_parallelism = env
+            .call_method(&java_obj, "getQueryParallelism", "()I", &[])?
+            .i()?;
+        scanner.query_parallelism(query_parallelism);
         Ok(())
     })?;
 

java/lance-jni/src/utils.rs

@@ -207,6 +207,9 @@ pub fn get_query(env: &mut JNIEnv, query_obj: JObject) -> Result<Option<Query>>
         };
 
         let use_index = env.get_boolean_from_method(&java_obj, "isUseIndex")?;
+        let query_parallelism = env
+            .call_method(&java_obj, "getQueryParallelism", "()I", &[])?
+            .i()?;
 
         Ok(Query {
             column,
@@ -221,6 +224,7 @@ pub fn get_query(env: &mut JNIEnv, query_obj: JObject) -> Result<Option<Query>>
             metric_type: distance_type,
             use_index,
             dist_q_c: 0.0,
+            query_parallelism,
         })
     })?;
 

java/src/main/java/org/lance/ipc/Query.java

@@ -31,6 +31,7 @@ public class Query {
   private final Optional<Integer> refineFactor;
   private final Optional<DistanceType> distanceType;
   private final boolean useIndex;
+  private final int queryParallelism;
 
   private Query(Builder builder) {
     this.column = Preconditions.checkNotNull(builder.column, "Columns must be set");
@@ -50,6 +51,7 @@ private Query(Builder builder) {
     this.refineFactor = builder.refineFactor;
     this.distanceType = builder.distanceType;
     this.useIndex = builder.useIndex;
+    this.queryParallelism = builder.queryParallelism;
   }
 
   public String getColumn() {
@@ -92,6 +94,10 @@ public boolean isUseIndex() {
     return useIndex;
   }
 
+  public int getQueryParallelism() {
+    return queryParallelism;
+  }
+
   @Override
   public String toString() {
     return MoreObjects.toStringHelper(this)
@@ -104,6 +110,7 @@ public String toString() {
         .add("refineFactor", refineFactor.orElse(null))
         .add("distanceType", distanceType.orElse(null))
         .add("useIndex", useIndex)
+        .add("queryParallelism", queryParallelism)
         .toString();
   }
 
@@ -117,6 +124,7 @@ public static class Builder {
     private Optional<Integer> refineFactor = Optional.empty();
     private Optional<DistanceType> distanceType = Optional.empty();
     private boolean useIndex = true;
+    private int queryParallelism = 0;
 
     /**
      * Sets the column to be searched.
@@ -245,6 +253,24 @@ public Builder setUseIndex(boolean useIndex) {
       return this;
     }
 
+    /**
+     * Sets vector partition search concurrency for each query.
+     *
+     * <p>The default is 0. Value 0 uses the automatic policy, which currently maps to the
+     * single-worker sequential path. Value -1 uses the CPU pool size. Value 1 uses the
+     * single-worker sequential path. Values greater than or equal to 2 use the partition-parallel
+     * path and are clamped to the CPU pool size.
+     *
+     * @param queryParallelism The partition search concurrency policy.
+     * @return The Builder instance for method chaining.
+     */
+    public Builder setQueryParallelism(int queryParallelism) {
+      Preconditions.checkArgument(
+          queryParallelism >= -1, "Query parallelism must be greater than or equal to -1");
+      this.queryParallelism = queryParallelism;
+      return this;
+    }
+
     /**
      * Builds the Query object.
      *

java/src/test/java/org/lance/JNITest.java

@@ -59,6 +59,7 @@ public void testQuery() {
                 .setRefineFactor(40)
                 .setDistanceType(DistanceType.L2)
                 .setUseIndex(true)
+                .setQueryParallelism(-1)
                 .build()));
   }
 

protos/ann.proto

@@ -25,6 +25,7 @@ message VectorQueryProto {
   optional lance.index.pb.VectorMetricType metric_type = 10;
   bool use_index = 11;
   optional float dist_q_c = 12;
+  optional int32 query_parallelism = 13;
 }
 
 // Serializable form of ANNIvfSubIndexExec — the IVF sub-index search node.

python/python/lance/dataset.py

@@ -6,6 +6,7 @@
 import copy
 import dataclasses
 import json
+import operator
 import os
 import random
 import time
@@ -5631,8 +5632,21 @@ def nearest(
         refine_factor: Optional[int] = None,
         use_index: bool = True,
         ef: Optional[int] = None,
+        query_parallelism: Optional[int] = None,
         distance_range: Optional[tuple[Optional[float], Optional[float]]] = None,
     ) -> ScannerBuilder:
+        """Configure nearest neighbor search.
+
+        Parameters
+        ----------
+        query_parallelism: int, optional
+            Maximum partition-search concurrency for a single vector query.
+            The default is 0. Value 0 uses the automatic policy, which
+            currently maps to the single-worker sequential path. Value -1 uses
+            the CPU pool size. Value 1 uses the single-worker sequential path.
+            Values >= 2 use the partition-parallel path and are clamped to the
+            CPU pool size.
+        """
         self._nearest = _build_vector_search_query(
             column,
             q,
@@ -5645,6 +5659,7 @@ def nearest(
             refine_factor=refine_factor,
             use_index=use_index,
             ef=ef,
+            query_parallelism=query_parallelism,
             distance_range=distance_range,
         )
         return self
@@ -6760,6 +6775,7 @@ def _build_vector_search_query(
     refine_factor: Optional[int] = None,
     use_index: bool = True,
     ef: Optional[int] = None,
+    query_parallelism: Optional[int] = None,
     distance_range: Optional[tuple[Optional[float], Optional[float]]] = None,
 ) -> dict:
     """Configure nearest neighbor search.
@@ -6787,6 +6803,12 @@ def _build_vector_search_query(
         Whether to use the index for the search.
     ef: int, optional
         The ef parameter for HNSW search.
+    query_parallelism: int, optional
+        Maximum partition-search concurrency for a single vector query.
+        The default is 0. Value 0 uses the automatic policy, which currently
+        maps to the single-worker sequential path. Value -1 uses the CPU pool
+        size. Value 1 uses the single-worker sequential path. Values >= 2 use
+        the partition-parallel path and are clamped to the CPU pool size.
     distance_range: tuple[Optional[float], Optional[float]], optional
         A tuple of (lower_bound, upper_bound) to filter results by distance.
         Both bounds are optional. The lower bound is inclusive and the upper
@@ -6854,6 +6876,11 @@ def _build_vector_search_query(
         # `ef` should be >= `k`, but `k` could be None so we can't check it here
         # the rust code will check it
         raise ValueError(f"ef must be > 0 but got {ef}")
+    if query_parallelism is not None:
+        query_parallelism = operator.index(query_parallelism)
+
+    if query_parallelism is not None and query_parallelism < -1:
+        raise ValueError("query_parallelism must be >= -1")
 
     if distance_range is not None:
         if len(distance_range) != 2:
@@ -6871,6 +6898,7 @@ def _build_vector_search_query(
         "refine_factor": refine_factor,
         "use_index": use_index,
         "ef": ef,
+        "query_parallelism": query_parallelism,
         "distance_range": distance_range,
     }
 
@@ -7043,6 +7071,7 @@ def __init__(
         refine_factor: Optional[int] = None,
         use_index: bool = True,
         ef: Optional[int] = None,
+        query_parallelism: Optional[int] = None,
     ):
         self._inner = _build_vector_search_query(
             column,
@@ -7055,6 +7084,7 @@ def __init__(
             refine_factor=refine_factor,
             use_index=use_index,
             ef=ef,
+            query_parallelism=query_parallelism,
         )
 
     def inner(self):

python/python/tests/test_vector_index.py

@@ -1864,6 +1864,41 @@ def test_vector_index_with_nprobes(indexed_dataset):
     ).analyze_plan()
 
 
+def test_vector_index_with_query_parallelism(indexed_dataset):
+    q = np.random.randn(128)
+
+    sequential = indexed_dataset.to_table(
+        nearest={
+            "column": "vector",
+            "q": q,
+            "k": 10,
+            "query_parallelism": 0,
+        }
+    )
+    parallel = indexed_dataset.to_table(
+        nearest={
+            "column": "vector",
+            "q": q,
+            "k": 10,
+            "query_parallelism": -1,
+        }
+    )
+
+    assert sequential == parallel
+
+
+def test_vector_index_invalid_query_parallelism(indexed_dataset):
+    with pytest.raises(ValueError, match="query_parallelism"):
+        indexed_dataset.scanner(
+            nearest={
+                "column": "vector",
+                "q": np.random.randn(128),
+                "k": 10,
+                "query_parallelism": -2,
+            }
+        )
+
+
 def test_knn_deleted_rows(tmp_path):
     data = create_table()
     ds = lance.write_dataset(data, tmp_path)

python/src/dataset.rs

@@ -77,8 +77,8 @@ use lance_index::{
     progress::{IndexBuildProgress, NoopIndexBuildProgress},
     scalar::{FullTextSearchQuery, InvertedIndexParams, ScalarIndexParams},
     vector::{
-        Query as VectorQuery, hnsw::builder::HnswBuildParams, ivf::IvfBuildParams,
-        pq::PQBuildParams, sq::builder::SQBuildParams,
+        DEFAULT_QUERY_PARALLELISM, Query as VectorQuery, hnsw::builder::HnswBuildParams,
+        ivf::IvfBuildParams, pq::PQBuildParams, sq::builder::SQBuildParams,
     },
 };
 use lance_index::{
@@ -1251,6 +1251,7 @@ impl Dataset {
                 refine_factor,
                 use_index,
                 ef,
+                query_parallelism,
             ) = vector_query_params_from_dict(nearest, default_k)?;
 
             let (_, element_type) = get_vector_type(self_.ds.schema(), &column)
@@ -1311,6 +1312,7 @@ impl Dataset {
                     if let Some(ef) = ef {
                         s = s.ef(ef);
                     }
+                    s = s.query_parallelism(query_parallelism);
                     s.use_index(use_index);
                     if let Some((lower, upper)) = distance_range {
                         s.distance_range(lower, upper);
@@ -4309,8 +4311,28 @@ type VectorQueryParams = (
     Option<u32>,
     bool,
     Option<usize>,
+    i32,
 );
 
+fn extract_query_parallelism(value: &Bound<'_, PyAny>) -> PyResult<i32> {
+    let query_parallelism = value.extract()?;
+    if query_parallelism < -1 {
+        Err(PyValueError::new_err("query_parallelism must be >= -1"))
+    } else {
+        Ok(query_parallelism)
+    }
+}
+
+fn vector_query_query_parallelism_from_dict(dict: &Bound<'_, PyDict>) -> PyResult<i32> {
+    if let Some(query_parallelism) = dict.get_item("query_parallelism")?
+        && !query_parallelism.is_none()
+    {
+        extract_query_parallelism(&query_parallelism)
+    } else {
+        Ok(DEFAULT_QUERY_PARALLELISM)
+    }
+}
+
 fn vector_query_params_from_dict(
     dict: &Bound<'_, PyDict>,
     default_k: usize,
@@ -4416,6 +4438,8 @@ fn vector_query_params_from_dict(
         None
     };
 
+    let query_parallelism = vector_query_query_parallelism_from_dict(dict)?;
+
     Ok((
         column,
         key,
@@ -4426,6 +4450,7 @@ fn vector_query_params_from_dict(
         refine_factor,
         use_index,
         ef,
+        query_parallelism,
     ))
 }
 
@@ -4461,6 +4486,7 @@ impl PySearchFilter {
             refine_factor,
             use_index,
             ef,
+            query_parallelism,
         ) = vector_query_params_from_dict(query, default_k)?;
 
         let metric_type = Some(metric_type_opt.unwrap_or(MetricType::L2));
@@ -4477,6 +4503,7 @@ impl PySearchFilter {
             refine_factor,
             metric_type,
             use_index,
+            query_parallelism,
             dist_q_c: 0.0,
         };