Add ability to query with prefetch (#179)

Author: tellet-q

README.md

@@ -269,12 +269,33 @@ batch. Supported kinds:
 
 | `kind` | Fields | Notes |
 |--------|--------|-------|
-| `dense` | `size`, optional `using`, `source`, `filters` | Query a dense vector; `source: { type: dataset }` measures recall (see below) |
+| `dense` | `size`, optional `using`, `source`, `filters`, `multivector`, `prefetch` | Query a dense vector; `source: { type: dataset }` measures recall (see below) |
 | `sparse` | `using`, `source`, `filters` | Query a named sparse vector; `source: { type: dataset }` measures recall (see below) |
 
 Filter entries reuse the same payload `type` / `source` vocabulary as the
 upload config.
 
+A `dense` request can also query a multivector and run as a two-stage query,
+the usual ColBERT setup: an HNSW search on one named vector picks candidates,
+then the request's own vector rescores them.
+
+```yaml
+requests:
+  - kind: dense
+    using: colbert                 # the rescoring vector
+    size: 128
+    multivector: { count: 16 }     # query with 16 sub-vectors of 128
+    prefetch:                      # first stage
+      using: dense                 # a named dense vector
+      size: 128
+      limit: 500                   # candidates handed to the rescore
+```
+
+Both fields need `source: random`. The prefetch stage gets the run's search
+params (`--search-hnsw-ef`, `--search-exact`, quantization flags). `--prefetch`
+and `--search-quality` cannot be combined with a config that sets `prefetch`, and
+`prefetch.limit` must be at least `--search-limit`.
+
 #### Measuring accuracy against a reference dataset
 
 A `dense` or `sparse` request can draw its queries from a vector-db-benchmark

src/config/search.rs

@@ -47,6 +47,13 @@ pub enum SearchRequestConfig {
         source: VectorSource,
         #[serde(default)]
         filters: Vec<FilterPayloadConfig>,
+        /// Query with several sub-vectors of `size` each (a multivector). Random source only.
+        #[serde(default)]
+        multivector: Option<QueryMultivectorConfig>,
+        /// Two-stage query: this search picks candidates, then the request's own vector
+        /// rescores them.
+        #[serde(default)]
+        prefetch: Option<PrefetchConfig>,
     },
     Sparse {
         using: String,
@@ -63,6 +70,28 @@ pub enum SearchRequestConfig {
     },
 }
 
+/// Shape of a multivector query.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct QueryMultivectorConfig {
+    /// Sub-vectors per query.
+    pub count: usize,
+}
+
+/// The first stage of a two-stage query: a random dense query on its own named vector.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct PrefetchConfig {
+    /// Named dense vector the first stage searches.
+    pub using: String,
+    /// Its dimension.
+    pub size: u64,
+    #[serde(default)]
+    pub datatype: DatatypeKind,
+    /// Candidates handed to the second stage.
+    pub limit: u64,
+}
+
 /// Payload field used to build a filter condition for a search request.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(deny_unknown_fields)]
@@ -103,6 +132,32 @@ pub fn parse(text: &str, origin: &str) -> Result<SearchConfig> {
 }
 
 impl SearchConfig {
+    /// Whether any request runs as a two-stage (prefetch then rescore) query.
+    pub fn has_prefetch(&self) -> bool {
+        self.requests.iter().any(|r| {
+            matches!(
+                r,
+                SearchRequestConfig::Dense {
+                    prefetch: Some(_),
+                    ..
+                }
+            )
+        })
+    }
+
+    /// The smallest prefetch `limit` across requests, if any request prefetches.
+    pub fn min_prefetch_limit(&self) -> Option<u64> {
+        self.requests
+            .iter()
+            .filter_map(|r| match r {
+                SearchRequestConfig::Dense {
+                    prefetch: Some(p), ..
+                } => Some(p.limit),
+                _ => None,
+            })
+            .min()
+    }
+
     pub fn validate(&self) -> Result<()> {
         if self.requests.is_empty() {
             bail!("search config must define at least one request");
@@ -135,6 +190,37 @@ impl SearchRequestConfig {
             }
         }
 
+        if let SearchRequestConfig::Dense {
+            source,
+            multivector,
+            prefetch,
+            ..
+        } = self
+        {
+            if let Some(multivector) = multivector {
+                if multivector.count == 0 {
+                    bail!("requests[{index}]: `multivector.count` must be > 0");
+                }
+                if !matches!(source, VectorSource::Random) {
+                    bail!("requests[{index}]: `multivector` needs `source: random`");
+                }
+            }
+            if let Some(prefetch) = prefetch {
+                if prefetch.using.is_empty() {
+                    bail!("requests[{index}]: `prefetch.using` must not be empty");
+                }
+                if prefetch.size == 0 {
+                    bail!("requests[{index}]: `prefetch.size` must be > 0");
+                }
+                if prefetch.limit == 0 {
+                    bail!("requests[{index}]: `prefetch.limit` must be > 0");
+                }
+                if !matches!(source, VectorSource::Random) {
+                    bail!("requests[{index}]: `prefetch` needs `source: random`");
+                }
+            }
+        }
+
         match self {
             SearchRequestConfig::Dense { size, source, .. } => {
                 match source {
@@ -183,6 +269,63 @@ mod tests {
     use super::*;
     use crate::config::DistributionKind;
 
+    const RESCORE_YAML: &str = r#"
+collection:
+  name: bench
+requests:
+  - kind: dense
+    using: colbert
+    size: 128
+    multivector: { count: 16 }
+    prefetch: { using: dense, size: 128, limit: 500 }
+"#;
+
+    #[test]
+    fn parses_multivector_prefetch_request() {
+        let cfg: SearchConfig = serde_yaml::from_str(RESCORE_YAML).unwrap();
+        cfg.validate().unwrap();
+        assert!(cfg.has_prefetch());
+        match &cfg.requests[0] {
+            SearchRequestConfig::Dense {
+                multivector: Some(m),
+                prefetch: Some(p),
+                ..
+            } => {
+                assert_eq!(m.count, 16);
+                assert_eq!((p.using.as_str(), p.size, p.limit), ("dense", 128, 500));
+            }
+            other => panic!("unexpected request {other:?}"),
+        }
+    }
+
+    #[test]
+    fn rejects_bad_multivector_and_prefetch() {
+        for (from, to, message) in [
+            ("count: 16", "count: 0", "`multivector.count` must be > 0"),
+            ("limit: 500", "limit: 0", "`prefetch.limit` must be > 0"),
+            (
+                "{ using: dense,",
+                "{ using: \"\",",
+                "`prefetch.using` must not be empty",
+            ),
+            (
+                "size: 128, limit",
+                "size: 0, limit",
+                "`prefetch.size` must be > 0",
+            ),
+            (
+                "multivector: { count: 16 }",
+                "multivector: { count: 16 }\n    source: { type: file, path: q.fbin }",
+                "`multivector` needs `source: random`",
+            ),
+        ] {
+            let yaml = RESCORE_YAML.replace(from, to);
+            let cfg: SearchConfig = serde_yaml::from_str(&yaml).unwrap();
+            let err = cfg.validate().unwrap_err().to_string();
+            assert!(err.contains(message), "{from} -> {to}: {err}");
+        }
+    }
+
     #[test]
     fn parses_minimal_search_config() {
         let yaml = r#"

src/generators/queries.rs

@@ -6,7 +6,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 
 use anyhow::Context;
 use qdrant_client::qdrant::{
-    Condition, Filter, GeoPoint, GeoRadius, Range, RepeatedStrings, SparseIndices,
+    Condition, Filter, GeoPoint, GeoRadius, Range, RepeatedStrings, SparseIndices, VectorInput,
     r#match::MatchValue,
 };
 use rand::Rng;
@@ -32,10 +32,34 @@ const GEO_SPREAD_DEG: f64 = 1.0;
 const GEO_RADIUS_METERS_MIN: f64 = 1000.0;
 const GEO_RADIUS_METERS_MAX: f64 = 50000.0;
 
+/// A dense query: one vector, or several sub-vectors for a multivector.
+#[derive(Debug, Clone, PartialEq)]
+pub enum DenseQuery {
+    Single(Vec<f32>),
+    Multi(Vec<Vec<f32>>),
+}
+
+impl DenseQuery {
+    pub fn into_vector_input(self) -> VectorInput {
+        match self {
+            DenseQuery::Single(vector) => VectorInput::new_dense(vector),
+            DenseQuery::Multi(vectors) => VectorInput::new_multi(vectors),
+        }
+    }
+}
+
+/// The first stage of a two-stage query.
+#[derive(Debug, Clone)]
+pub struct GeneratedPrefetch {
+    pub vector: Vec<f32>,
+    pub using: String,
+    pub limit: u64,
+}
+
 /// One query vector plus optional filter, ready to be turned into a gRPC request.
 #[derive(Debug, Clone)]
 pub struct GeneratedQuery {
-    pub dense: Option<(Vec<f32>, Option<String>)>,
+    pub dense: Option<(DenseQuery, Option<String>)>,
     pub sparse: Option<(Vec<f32>, SparseIndices, String)>,
     pub filter: Option<Filter>,
     /// Sparse-vector IDF corpus: restricts which points the IDF statistics are
@@ -45,6 +69,8 @@ pub struct GeneratedQuery {
     /// when the request draws queries from a reference dataset. Used to measure
     /// search accuracy (recall) against the dataset's known answers.
     pub expected_ids: Option<Vec<u64>>,
+    /// Present for a two-stage query: its candidates are what the main query rescores.
+    pub prefetch: Option<GeneratedPrefetch>,
 }
 
 /// A reference dataset's query set, held in memory, with a cursor that hands out
@@ -492,10 +518,20 @@ impl ConfigSearchGenerator {
                 datatype,
                 source,
                 filters: _,
+                multivector,
+                prefetch,
             } => {
                 let (vector, expected_ids, dataset_filter) =
                     if let Some(query_dataset) = &state.query_dataset {
-                        Self::read_dense_query(query_dataset)
+                        let (vector, ids, filter) = Self::read_dense_query(query_dataset);
+                        (DenseQuery::Single(vector), ids, filter)
+                    } else if let Some(multivector) = multivector {
+                        // Validation allows multivector only with random queries.
+                        let is_uint = *datatype == DatatypeKind::Uint8;
+                        let vectors = (0..multivector.count)
+                            .map(|_| random_dense_vector(rng, *size as usize, is_uint))
+                            .collect();
+                        (DenseQuery::Multi(vectors), None, None)
                     } else {
                         let vector = Self::gen_dense_vector(
                             rng,
@@ -505,14 +541,24 @@ impl ConfigSearchGenerator {
                             state.dense_reader.as_ref(),
                             req_id,
                         );
-                        (vector, None, None)
+                        (DenseQuery::Single(vector), None, None)
                     };
+                let prefetch = prefetch.as_ref().map(|p| GeneratedPrefetch {
+                    vector: random_dense_vector(
+                        rng,
+                        p.size as usize,
+                        p.datatype == DatatypeKind::Uint8,
+                    ),
+                    using: p.using.clone(),
+                    limit: p.limit,
+                });
                 GeneratedQuery {
                     dense: Some((vector, using.clone())),
                     sparse: None,
                     filter: dataset_filter.or_else(|| state.filters.build(rng)),
                     idf_corpus: None,
                     expected_ids,
+                    prefetch,
                 }
             }
             SearchRequestConfig::Sparse {
@@ -537,6 +583,7 @@ impl ConfigSearchGenerator {
                     filter: dataset_filter.or_else(|| state.filters.build(rng)),
                     idf_corpus: state.idf_corpus.build(rng),
                     expected_ids,
+                    prefetch: None,
                 }
             }
         }
@@ -645,6 +692,26 @@ mod tests {
         ConfigSearchGenerator::new(&config).unwrap()
     }
 
+    #[test]
+    fn generates_multivector_query_with_prefetch() {
+        let generator = build_gen(
+            "collection:\n  name: x\nrequests:\n  - kind: dense\n    using: colbert\n    size: 8\n    multivector: { count: 4 }\n    prefetch: { using: dense, size: 6, limit: 50 }\n",
+        );
+        let q = generator.make_query(0, &mut rand::rng());
+        let (query, using) = q.dense.unwrap();
+        assert_eq!(using.as_deref(), Some("colbert"));
+        match query {
+            DenseQuery::Multi(vectors) => {
+                assert_eq!(vectors.len(), 4);
+                assert!(vectors.iter().all(|v| v.len() == 8));
+            }
+            other => panic!("expected a multivector query, got {other:?}"),
+        }
+        let prefetch = q.prefetch.unwrap();
+        assert_eq!((prefetch.using.as_str(), prefetch.limit), ("dense", 50));
+        assert_eq!(prefetch.vector.len(), 6);
+    }
+
     #[test]
     fn generates_dense_and_sparse_queries() {
         let generator = build_gen(
@@ -696,11 +763,17 @@ mod tests {
         let q0 = generator.make_query_for(0, 0, &mut rng);
         let q1 = generator.make_query_for(0, 0, &mut rng);
         let q2 = generator.make_query_for(0, 0, &mut rng);
-        assert_eq!(q0.dense.as_ref().unwrap().0, vec![0.0, 1.0, 2.0, 3.0]);
+        assert_eq!(
+            q0.dense.as_ref().unwrap().0,
+            DenseQuery::Single(vec![0.0, 1.0, 2.0, 3.0])
+        );
         assert_eq!(q0.expected_ids, Some(vec![0, 2]));
         assert_eq!(q1.expected_ids, Some(vec![1, 2]));
         // Wrapped back to the first query.
-        assert_eq!(q2.dense.as_ref().unwrap().0, vec![0.0, 1.0, 2.0, 3.0]);
+        assert_eq!(
+            q2.dense.as_ref().unwrap().0,
+            DenseQuery::Single(vec![0.0, 1.0, 2.0, 3.0])
+        );
         assert_eq!(q2.expected_ids, Some(vec![0, 2]));
     }