diff --git a/README.md b/README.md index c74cb11..ae5013e 100644 --- a/README.md +++ b/README.md @@ -287,7 +287,7 @@ batch. Supported kinds: | `kind` | Fields | Notes | |--------|--------|-------| -| `dense` | `size`, optional `using`, `source`, `filters`, `multivector`, `prefetch` | Query a dense vector; `source: { type: dataset }` measures recall (see below) | +| `dense` | `size`, optional `using`, `source`, `filters`, `multivector`, `prefetch`, `fusion` | 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 @@ -317,6 +317,35 @@ prefetches from a dataset without it is refused, rather than pairing a query's sub-vectors with an unrelated vector — `prefetch.size` is then ignored, like the request's own `size`. A `file` source cannot drive a prefetch. +#### Hybrid search: fusing several prefetches + +`prefetch` also takes a list, and a `fusion` on the request combines what the +stages return instead of rescoring them. Stages may be dense or sparse, which is +what makes a dense-plus-sparse hybrid query: + +```yaml +requests: + - kind: dense + size: 128 + fusion: rrf # rrf | dbsf + rrf_k: 60 # optional, rrf only + weights: [2.0, 1.0] # optional, rrf only; one per stage + prefetch: + - using: dense # stage 1: dense vector + size: 128 + limit: 200 + - using: bm25 # stage 2: sparse vector + kind: sparse + limit: 200 + source: { vocab_size: 100000, length: 20 } +``` + +A fused request sends no vector of its own — the stages carry the queries and +`fusion` decides the ranking — so `multivector` cannot apply to it. Fusion needs +at least two stages, `weights` needs one entry per stage, and `rrf_k` and +`weights` need `fusion: rrf`. A single prefetch without `fusion` stays the +two-stage rescore described above. + #### ACORN `--acorn` asks Qdrant to use ACORN for filtered HNSW search, and diff --git a/src/config/search.rs b/src/config/search.rs index 7364492..ccefec6 100644 --- a/src/config/search.rs +++ b/src/config/search.rs @@ -50,10 +50,20 @@ pub enum SearchRequestConfig { /// Query with several sub-vectors of `size` each (a multivector). Random source only. #[serde(default)] multivector: Option, - /// Two-stage query: this search picks candidates, then the request's own vector - /// rescores them. + /// Candidate stages. One prefetch feeds a rescore by this request's own vector; + /// several feed a `fusion` query. Accepts a single entry or a list. + #[serde(default, deserialize_with = "one_or_many")] + prefetch: Vec, + /// Combine several prefetches instead of rescoring them: `rrf` or `dbsf`. + /// The request then sends no vector of its own. #[serde(default)] - prefetch: Option, + fusion: Option, + /// Reciprocal rank fusion's `k`. Needs `fusion: rrf`. + #[serde(default)] + rrf_k: Option, + /// Per-prefetch weights, in the order the prefetches are listed. Needs `fusion: rrf`. + #[serde(default)] + weights: Option>, }, Sparse { using: String, @@ -82,14 +92,57 @@ pub struct QueryMultivectorConfig { #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct PrefetchConfig { - /// Named dense vector the first stage searches. + /// Named vector this stage searches. pub using: String, - /// Its dimension. + /// Dense only: its dimension. + #[serde(default)] pub size: u64, #[serde(default)] pub datatype: DatatypeKind, - /// Candidates handed to the second stage. + /// Candidates this stage hands on. pub limit: u64, + /// What the stage queries. Dense unless said otherwise. + #[serde(default)] + pub kind: PrefetchKind, + /// Sparse only: how its query vectors are generated. + #[serde(default)] + pub source: Option, +} + +/// Which kind of vector a prefetch stage searches. +#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum PrefetchKind { + #[default] + Dense, + Sparse, +} + +/// How several prefetches are combined into one ranking. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum FusionKind { + /// Reciprocal Rank Fusion. + Rrf, + /// Distribution-Based Score Fusion. + Dbsf, +} + +/// A prefetch written as a single mapping, or several as a list. +fn one_or_many<'de, D>(deserializer: D) -> std::result::Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + #[derive(Deserialize)] + #[serde(untagged)] + enum OneOrMany { + One(Box), + Many(Vec), + } + Ok(match OneOrMany::deserialize(deserializer)? { + OneOrMany::One(one) => vec![*one], + OneOrMany::Many(many) => many, + }) } /// Payload field used to build a filter condition for a search request. @@ -132,29 +185,16 @@ pub fn parse(text: &str, origin: &str) -> Result { } impl SearchConfig { - /// Whether any request runs as a two-stage (prefetch then rescore) query. + /// Whether any request runs its query in stages: prefetch then rescore, or fusion. pub fn has_prefetch(&self) -> bool { - self.requests.iter().any(|r| { - matches!( - r, - SearchRequestConfig::Dense { - prefetch: Some(_), - .. - } - ) - }) + self.requests.iter().any(|r| !r.prefetches().is_empty()) } /// The smallest prefetch `limit` across requests, if any request prefetches. pub fn min_prefetch_limit(&self) -> Option { self.requests .iter() - .filter_map(|r| match r { - SearchRequestConfig::Dense { - prefetch: Some(p), .. - } => Some(p.limit), - _ => None, - }) + .flat_map(|r| r.prefetches().iter().map(|p| p.limit)) .min() } @@ -173,6 +213,14 @@ impl SearchConfig { impl SearchRequestConfig { /// Payload conditions applied to the query itself. + /// The request's prefetch stages; empty when it queries in one step. + pub fn prefetches(&self) -> &[PrefetchConfig] { + match self { + SearchRequestConfig::Dense { prefetch, .. } => prefetch, + SearchRequestConfig::Sparse { .. } => &[], + } + } + pub fn filters(&self) -> &[FilterPayloadConfig] { match self { SearchRequestConfig::Dense { filters, .. } @@ -194,6 +242,9 @@ impl SearchRequestConfig { source, multivector, prefetch, + fusion, + rrf_k, + weights, .. } = self { @@ -205,21 +256,85 @@ impl SearchRequestConfig { 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"); + for (slot, stage) in prefetch.iter().enumerate() { + let at = format!("requests[{index}].prefetch[{slot}]"); + if stage.using.is_empty() { + bail!("{at}: `using` must not be empty"); } - if prefetch.size == 0 { - bail!("requests[{index}]: `prefetch.size` must be > 0"); + if stage.limit == 0 { + bail!("{at}: `limit` must be > 0"); + } + match stage.kind { + PrefetchKind::Dense => { + // A misplaced `source` is the likelier mistake, so say that first. + if stage.source.is_some() { + bail!("{at}: `source` describes a sparse query; add `kind: sparse`"); + } + if stage.size == 0 { + bail!("{at}: dense `size` must be > 0"); + } + } + PrefetchKind::Sparse => { + let sparse = stage + .source + .as_ref() + .with_context(|| format!("{at}: a sparse stage needs `source`"))?; + if sparse.kind != SparseKind::Random { + bail!( + "{at}: sparse prefetch queries must be generated: `type: random`" + ); + } + if sparse.length == 0 || sparse.length > sparse.vocab_size { + bail!("{at}: sparse `length` must be > 0 and <= `vocab_size`"); + } + } } - if prefetch.limit == 0 { - bail!("requests[{index}]: `prefetch.limit` must be > 0"); + } + if !prefetch.is_empty() && matches!(source, VectorSource::File { .. }) { + bail!( + "requests[{index}]: `prefetch` needs `source: random` or a dataset \ + source whose query set carries prefetch vectors" + ); + } + match fusion { + None => { + if prefetch.len() > 1 { + bail!( + "requests[{index}]: {} prefetches need `fusion` to combine them; \ + a rescore takes one", + prefetch.len() + ); + } + if rrf_k.is_some() || weights.is_some() { + bail!("requests[{index}]: `rrf_k` and `weights` need `fusion: rrf`"); + } } - if matches!(source, VectorSource::File { .. }) { - bail!( - "requests[{index}]: `prefetch` needs `source: random` or a dataset \ - source whose query set carries prefetch vectors" - ); + Some(kind) => { + if prefetch.len() < 2 { + bail!( + "requests[{index}]: `fusion` combines prefetches, so it needs at \ + least 2; got {}", + prefetch.len() + ); + } + if multivector.is_some() { + bail!( + "requests[{index}]: a `fusion` query sends no vector of its own, \ + so `multivector` cannot apply" + ); + } + if *kind == FusionKind::Dbsf && (rrf_k.is_some() || weights.is_some()) { + bail!("requests[{index}]: `rrf_k` and `weights` need `fusion: rrf`"); + } + if let Some(weights) = weights + && weights.len() != prefetch.len() + { + bail!( + "requests[{index}]: {} `weights` for {} prefetches; one each", + weights.len(), + prefetch.len() + ); + } } } } @@ -283,6 +398,96 @@ requests: prefetch: { using: dense, size: 128, limit: 500 } "#; + const FUSION_YAML: &str = "collection:\n name: x\nrequests:\n - kind: dense\n size: 128\n fusion: rrf\n prefetch:\n - using: dense\n size: 128\n limit: 200\n - using: bm25\n kind: sparse\n limit: 200\n source: { vocab_size: 1000, length: 10 }\n"; + + #[test] + fn parses_hybrid_fusion_request() { + let cfg: SearchConfig = serde_yaml::from_str(FUSION_YAML).unwrap(); + cfg.validate().unwrap(); + assert!(cfg.has_prefetch()); + assert_eq!(cfg.min_prefetch_limit(), Some(200)); + match &cfg.requests[0] { + SearchRequestConfig::Dense { + prefetch, fusion, .. + } => { + assert_eq!(*fusion, Some(FusionKind::Rrf)); + assert_eq!(prefetch.len(), 2); + assert_eq!(prefetch[0].kind, PrefetchKind::Dense); + assert_eq!(prefetch[1].kind, PrefetchKind::Sparse); + assert_eq!(prefetch[1].source.as_ref().unwrap().length, 10); + } + other => panic!("unexpected request {other:?}"), + } + } + + #[test] + fn parses_rrf_parameters() { + let yaml = FUSION_YAML.replace( + "fusion: rrf", + "fusion: rrf\n rrf_k: 60\n weights: [2.0, 1.0]", + ); + let cfg: SearchConfig = serde_yaml::from_str(&yaml).unwrap(); + cfg.validate().unwrap(); + match &cfg.requests[0] { + SearchRequestConfig::Dense { rrf_k, weights, .. } => { + assert_eq!(*rrf_k, Some(60)); + assert_eq!(weights.as_deref(), Some(&[2.0f32, 1.0][..])); + } + other => panic!("unexpected request {other:?}"), + } + } + + #[test] + fn rejects_bad_fusion() { + for (from, to, message) in [ + // Several stages with nothing to combine them. + (" fusion: rrf\n", "", "need `fusion` to combine them"), + // Fusion ranks stages against each other, so one is not a fusion. + ( + " - using: bm25\n kind: sparse\n limit: 200\n source: { vocab_size: 1000, length: 10 }\n", + "", + "needs at least 2", + ), + // One weight per stage, or the ranking silently means something else. + ( + "fusion: rrf", + "fusion: rrf\n weights: [1.0]", + "1 `weights` for 2 prefetches", + ), + ( + "fusion: rrf", + "fusion: dbsf\n rrf_k: 60", + "need `fusion: rrf`", + ), + ( + " kind: sparse\n", + "", + "`source` describes a sparse query; add `kind: sparse`", + ), + ( + " source: { vocab_size: 1000, length: 10 }\n", + "", + "a sparse stage needs `source`", + ), + ( + " size: 128\n fusion: rrf", + " size: 128\n multivector: { count: 4 }\n fusion: rrf", + "sends no vector of its own", + ), + ] { + let yaml = FUSION_YAML.replace(from, to); + assert_ne!(yaml, FUSION_YAML, "replacement {from:?} did not apply"); + let err = match serde_yaml::from_str::(&yaml) { + Ok(cfg) => match cfg.validate() { + Err(err) => err.to_string(), + Ok(()) => panic!("accepted {message:?}"), + }, + Err(err) => err.to_string(), + }; + assert!(err.contains(message), "expected {message:?}, got {err}"); + } + } + #[test] fn parses_multivector_prefetch_request() { let cfg: SearchConfig = serde_yaml::from_str(RESCORE_YAML).unwrap(); @@ -291,10 +496,13 @@ requests: match &cfg.requests[0] { SearchRequestConfig::Dense { multivector: Some(m), - prefetch: Some(p), + prefetch, .. } => { assert_eq!(m.count, 16); + let [p] = &prefetch[..] else { + panic!("expected one prefetch, got {}", prefetch.len()) + }; assert_eq!((p.using.as_str(), p.size, p.limit), ("dense", 128, 500)); } other => panic!("unexpected request {other:?}"), @@ -305,16 +513,16 @@ requests: 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"), + ("limit: 500", "limit: 0", "`limit` must be > 0"), ( "{ using: dense,", "{ using: \"\",", - "`prefetch.using` must not be empty", + "`using` must not be empty", ), ( "size: 128, limit", "size: 0, limit", - "`prefetch.size` must be > 0", + "dense `size` must be > 0", ), ( "multivector: { count: 16 }", diff --git a/src/generators/queries.rs b/src/generators/queries.rs index dc8ab7d..7c8f9cd 100644 --- a/src/generators/queries.rs +++ b/src/generators/queries.rs @@ -18,7 +18,9 @@ use super::random::{ DEFAULT_VOCAB_SIZE, create_zipf, random_dense_vector, random_keyword, random_sparse_vector, random_text, }; -use crate::config::search::{FilterPayloadConfig, SearchConfig, SearchRequestConfig}; +use crate::config::search::{ + FilterPayloadConfig, FusionKind, PrefetchKind, SearchConfig, SearchRequestConfig, +}; use crate::config::{ DatatypeKind, DistributionKind, FileStrategy, PayloadSourceKind, PayloadType, SparseKind, VectorSource, @@ -51,11 +53,26 @@ impl DenseQuery { /// The first stage of a two-stage query. #[derive(Debug, Clone)] pub struct GeneratedPrefetch { - pub vector: Vec, + pub query: PrefetchQuery, pub using: String, pub limit: u64, } +/// What one prefetch stage searches with. +#[derive(Debug, Clone)] +pub enum PrefetchQuery { + Dense(Vec), + Sparse(Vec, SparseIndices), +} + +/// How a query's prefetches are combined, when it fuses them instead of rescoring. +#[derive(Debug, Clone)] +pub struct FusionSpec { + pub kind: FusionKind, + pub rrf_k: Option, + pub weights: Option>, +} + /// One query vector plus optional filter, ready to be turned into a gRPC request. #[derive(Debug, Clone)] pub struct GeneratedQuery { @@ -70,7 +87,10 @@ pub struct GeneratedQuery { /// search accuracy (recall) against the dataset's known answers. pub expected_ids: Option>, /// Present for a two-stage query: its candidates are what the main query rescores. - pub prefetch: Option, + pub prefetch: Vec, + /// Set when the prefetches are fused rather than rescored; the query then + /// carries no vector of its own. + pub fusion: Option, } /// A reference dataset's query set, held in memory, with a cursor that hands out @@ -362,7 +382,7 @@ impl ConfigSearchGenerator { dataset, datasets_dir, QueryKind::Dense, - prefetch.is_some(), + !prefetch.is_empty(), )?), ), VectorSource::Random => (None, None), @@ -585,6 +605,9 @@ impl ConfigSearchGenerator { filters: _, multivector, prefetch, + fusion, + rrf_k, + weights, } => { let (vector, expected_ids, dataset_filter, dataset_prefetch) = if let Some(query_dataset) = &state.query_dataset { @@ -607,22 +630,45 @@ impl ConfigSearchGenerator { ); (DenseQuery::Single(vector), None, None, None) }; - let prefetch = prefetch.as_ref().map(|p| GeneratedPrefetch { - // A dataset supplies the prefetch's own query; only a - // generated query needs a generated prefetch vector. - vector: dataset_prefetch.unwrap_or_else(|| { - random_dense_vector(rng, p.size as usize, p.datatype == DatatypeKind::Uint8) - }), - using: p.using.clone(), - limit: p.limit, - }); + // A dataset supplies the prefetch's own query, for the first dense stage; + // every other stage generates one. + let mut dataset_prefetch = dataset_prefetch; + let prefetch = prefetch + .iter() + .map(|p| GeneratedPrefetch { + query: match p.kind { + PrefetchKind::Dense => { + PrefetchQuery::Dense(dataset_prefetch.take().unwrap_or_else(|| { + random_dense_vector( + rng, + p.size as usize, + p.datatype == DatatypeKind::Uint8, + ) + })) + } + PrefetchKind::Sparse => { + let source = p.source.as_ref().expect("validated sparse source"); + let (values, indices) = Self::gen_sparse_vector(rng, source, None); + PrefetchQuery::Sparse(values, indices) + } + }, + using: p.using.clone(), + limit: p.limit, + }) + .collect(); GeneratedQuery { - dense: Some((vector, using.clone())), + // A fusion query ranks its prefetches and sends no vector itself. + dense: fusion.is_none().then(|| (vector, using.clone())), sparse: None, filter: dataset_filter.or_else(|| state.filters.build(rng)), idf_corpus: None, expected_ids, prefetch, + fusion: fusion.map(|kind| FusionSpec { + kind, + rrf_k: *rrf_k, + weights: weights.clone(), + }), } } SearchRequestConfig::Sparse { @@ -643,11 +689,12 @@ impl ConfigSearchGenerator { }; GeneratedQuery { dense: None, + fusion: None, sparse: Some((values, indices, using.clone())), filter: dataset_filter.or_else(|| state.filters.build(rng)), idf_corpus: state.idf_corpus.build(rng), expected_ids, - prefetch: None, + prefetch: Vec::new(), } } } @@ -783,9 +830,11 @@ mod tests { } other => panic!("expected a multivector query, got {other:?}"), } - let prefetch = q.prefetch.unwrap(); + let [prefetch] = &q.prefetch[..] else { + panic!("expected one prefetch stage") + }; assert_eq!((prefetch.using.as_str(), prefetch.limit), ("dense", 50)); - assert_eq!(prefetch.vector.len(), 6); + assert!(matches!(&prefetch.query, PrefetchQuery::Dense(v) if v.len() == 6)); } #[test] @@ -892,11 +941,13 @@ mod tests { assert_eq!(q1.expected_ids, Some(vec![1, 2])); // The prefetch vector is the dataset's, and it moves with the query. - let p0 = q0.prefetch.unwrap(); - let p1 = q1.prefetch.unwrap(); + let (p0, p1) = (&q0.prefetch[0], &q1.prefetch[0]); assert_eq!((p0.using.as_str(), p0.limit), ("dense", 50)); - assert_eq!(p0.vector.len(), 6); - assert_ne!(p0.vector, p1.vector); + let (PrefetchQuery::Dense(v0), PrefetchQuery::Dense(v1)) = (&p0.query, &p1.query) else { + panic!("expected dense prefetch queries") + }; + assert_eq!(v0.len(), 6); + assert_ne!(v0, v1); } /// Without prefetch queries the two stages would search unrelated vectors, diff --git a/src/search/from_config.rs b/src/search/from_config.rs index 8eb6627..2c50614 100644 --- a/src/search/from_config.rs +++ b/src/search/from_config.rs @@ -7,17 +7,17 @@ use indicatif::ProgressBar; use qdrant_client::Qdrant; use qdrant_client::qdrant::shard_key::Key; use qdrant_client::qdrant::{ - AcornSearchParamsBuilder, IdfParamsBuilder, PrefetchQueryBuilder, - QuantizationSearchParamsBuilder, Query, QueryBatchPointsBuilder, QueryPointsBuilder, + AcornSearchParamsBuilder, Fusion, IdfParamsBuilder, PrefetchQueryBuilder, + QuantizationSearchParamsBuilder, Query, QueryBatchPointsBuilder, QueryPointsBuilder, Rrf, SearchParams, SearchParamsBuilder, VectorInput, }; use super::{SearchStats, compare_batch_search_results, recall_against_ground_truth}; use crate::args::Args; use crate::client::retry_with_clients; -use crate::config::search::SearchConfig; +use crate::config::search::{FusionKind, SearchConfig}; use crate::generators::ConfigSearchGenerator; -use crate::generators::queries::GeneratedPrefetch; +use crate::generators::queries::{FusionSpec, GeneratedPrefetch, PrefetchQuery}; use crate::processor::{Processor, Timing}; /// YAML-config-driven search (`bfb search --file config.yaml`). @@ -127,9 +127,9 @@ impl ConfigSearchProcessor { &self, query_filter: Option, idf_corpus: Option, - vector: VectorInput, - vector_name: Option, - prefetch: Option, + vector: Option<(VectorInput, Option)>, + prefetch: Vec, + fusion: Option, mut search_params: SearchParamsBuilder, ) -> QueryPointsBuilder { // Sparse IDF statistics are computed over this sub-corpus instead of the @@ -143,28 +143,33 @@ impl ConfigSearchProcessor { .with_vectors(self.args.search_with_vectors) .limit(self.args.search_limit as u64); - if let Some(vector_name) = vector_name.filter(|n| !n.is_empty()) { - request_builder = request_builder.using(vector_name); - } - if let Some(filter) = query_filter { request_builder = request_builder.filter(filter); } - let query = Query::new_nearest(vector); - - if let Some(prefetch) = prefetch { - // Two stages: the prefetch searches its own vector with the run's search params, - // the main query rescores its candidates. - let prefetch = PrefetchQueryBuilder::default() - .query(Query::new_nearest(VectorInput::new_dense(prefetch.vector))) - .using(prefetch.using) - .params(search_params.clone()) - .limit(prefetch.limit) - .build(); - request_builder = request_builder - .prefetch(vec![prefetch]) - .params(search_params); + if !prefetch.is_empty() { + // Each stage searches its own vector with the run's search params; the outer + // query then rescores one stage's candidates, or fuses several stages' rankings. + let stages: Vec<_> = prefetch + .into_iter() + .map(|stage| { + let query = match stage.query { + PrefetchQuery::Dense(vector) => { + Query::new_nearest(VectorInput::new_dense(vector)) + } + PrefetchQuery::Sparse(values, indices) => { + Query::new_nearest(VectorInput::new_sparse(indices.data, values)) + } + }; + PrefetchQueryBuilder::default() + .query(query) + .using(stage.using) + .params(search_params.clone()) + .limit(stage.limit) + .build() + }) + .collect(); + request_builder = request_builder.prefetch(stages).params(search_params); } else if let Some(prefetch_limit) = self.args.prefetch { let mut prefetch_params = SearchParamsBuilder::default() .quantization(QuantizationSearchParamsBuilder::default().rescore(false)); @@ -173,19 +178,43 @@ impl ConfigSearchProcessor { prefetch_params = prefetch_params.hnsw_ef(hnsw_ef as u64); } - let prefetch = PrefetchQueryBuilder::default() - .query(query.clone()) + let stage = PrefetchQueryBuilder::default() + .query(Query::new_nearest( + vector + .as_ref() + .expect("--prefetch needs the request's own vector") + .0 + .clone(), + )) .params(prefetch_params) .limit(prefetch_limit as u64) .build(); - request_builder = request_builder.prefetch(vec![prefetch]); + request_builder = request_builder.prefetch(vec![stage]); request_builder = request_builder.params(search_params); } else { request_builder = request_builder.params(search_params); } - request_builder = request_builder.query(query); + request_builder = match fusion { + Some(spec) => request_builder.query(match spec.kind { + FusionKind::Dbsf => Query::new_fusion(Fusion::Dbsf), + FusionKind::Rrf => match (spec.rrf_k, spec.weights) { + (None, None) => Query::new_fusion(Fusion::Rrf), + (k, weights) => Query::new_rrf(Rrf { + k, + weights: weights.unwrap_or_default(), + }), + }, + }), + None => { + let (vector, name) = vector.expect("a query without fusion needs a vector"); + if let Some(name) = name.filter(|n| !n.is_empty()) { + request_builder = request_builder.using(name); + } + request_builder.query(Query::new_nearest(vector)) + } + }; if let Some(read_consistency) = self.args.read_consistency { request_builder = request_builder.read_consistency(read_consistency); @@ -236,21 +265,25 @@ impl ConfigSearchProcessor { .map(|generated| { let query_filter = generated.filter.clone(); let idf_corpus = generated.idf_corpus.clone(); - let (vector, vector_name) = if let Some((values, indices, name)) = generated.sparse - { - (VectorInput::new_sparse(indices.data, values), Some(name)) - } else if let Some((query, name)) = generated.dense { - (query.into_vector_input(), name) + let vector = if let Some((values, indices, name)) = generated.sparse { + Some((VectorInput::new_sparse(indices.data, values), Some(name))) } else { - panic!("search config request must produce a dense or sparse vector"); + // A fusion query has no vector of its own; anything else must have one. + generated + .dense + .map(|(query, name)| (query.into_vector_input(), name)) }; + assert!( + vector.is_some() || generated.fusion.is_some(), + "a request without fusion must produce a dense or sparse vector" + ); self.create_request_builder( query_filter, idf_corpus, vector, - vector_name, generated.prefetch, + generated.fusion, search_params.clone(), ) .build() @@ -454,6 +487,87 @@ mod tests { ) } + const FUSION_YAML: &str = "collection:\n name: x\nrequests:\n - kind: dense\n size: 8\n fusion: rrf\n rrf_k: 60\n weights: [2.0, 1.0]\n prefetch:\n - using: dense\n size: 8\n limit: 200\n - using: bm25\n kind: sparse\n limit: 200\n source: { vocab_size: 1000, length: 10 }\n"; + + /// A fused request sends its stages and how to combine them, and no vector of its own. + #[test] + fn builds_a_hybrid_fusion_request() { + let config = parse(FUSION_YAML, "test").unwrap(); + let p = ConfigSearchProcessor::new( + Args::parse_from(vec!["bfb", "--search-hnsw-ef", "64"]), + &config, + Arc::new(AtomicBool::new(false)), + vec![], + ) + .unwrap(); + let q = p.generator.make_query(0, &mut rand::rng()); + assert!( + q.dense.is_none(), + "a fusion query carries no vector of its own" + ); + + let request = p + .create_request_builder( + None, + None, + None, + q.prefetch, + q.fusion, + SearchParamsBuilder::default().hnsw_ef(64), + ) + .build(); + + assert!(request.using.is_none()); + let Some(QueryVariant::Rrf(rrf)) = request.query.and_then(|q| q.variant) else { + panic!("expected a parameterized rrf query"); + }; + assert_eq!((rrf.k, rrf.weights), (Some(60), vec![2.0, 1.0])); + + assert_eq!(request.prefetch.len(), 2); + assert_eq!(request.prefetch[0].using.as_deref(), Some("dense")); + assert_eq!(request.prefetch[1].using.as_deref(), Some("bm25")); + // Both stages run under the run's own search params. + for stage in &request.prefetch { + assert_eq!(stage.limit, Some(200)); + assert_eq!(stage.params.as_ref().and_then(|p| p.hnsw_ef), Some(64)); + } + let Some(QueryVariant::Nearest(sparse)) = + request.prefetch[1].query.clone().and_then(|q| q.variant) + else { + panic!("expected a nearest query on the sparse stage"); + }; + assert!(matches!(sparse.variant, Some(InputVariant::Sparse(s)) if !s.values.is_empty())); + } + + /// Plain `fusion: rrf` with no parameters uses the server's own defaults. + #[test] + fn plain_fusion_sends_no_parameters() { + let yaml = FUSION_YAML.replace(" rrf_k: 60\n weights: [2.0, 1.0]\n", ""); + let config = parse(&yaml, "test").unwrap(); + let p = ConfigSearchProcessor::new( + Args::parse_from(vec!["bfb"]), + &config, + Arc::new(AtomicBool::new(false)), + vec![], + ) + .unwrap(); + let q = p.generator.make_query(0, &mut rand::rng()); + let request = p + .create_request_builder( + None, + None, + None, + q.prefetch, + q.fusion, + SearchParamsBuilder::default(), + ) + .build(); + assert!(matches!( + request.query.and_then(|q| q.variant), + Some(QueryVariant::Fusion(f)) if f == Fusion::Rrf as i32 + )); + } + #[test] fn acorn_flags_reach_the_search_params() { let plain = filtered_processor(&[]).unwrap().search_params().build(); @@ -528,9 +642,9 @@ mod tests { .create_request_builder( None, None, - query.into_vector_input(), - using, + Some((query.into_vector_input(), using)), q.prefetch, + q.fusion, // As search() builds it from --search-hnsw-ef. SearchParamsBuilder::default().hnsw_ef(64), ) @@ -563,9 +677,9 @@ mod tests { .create_request_builder( None, None, - query.into_vector_input(), - using, + Some((query.into_vector_input(), using)), q.prefetch, + q.fusion, SearchParamsBuilder::default().exact(true).hnsw_ef(64), ) .build(); diff --git a/src/serverless/query.rs b/src/serverless/query.rs index d25923f..07f3846 100644 --- a/src/serverless/query.rs +++ b/src/serverless/query.rs @@ -61,9 +61,12 @@ impl ServerlessQueryProcessor { generated: GeneratedQuery, collection: &str, ) -> Result { - if generated.prefetch.is_some() { + if !generated.prefetch.is_empty() { bail!("prefetch is not supported for serverless queries"); } + if generated.fusion.is_some() { + bail!("fusion is not supported for serverless queries"); + } let (vector, using) = if let Some((values, indices, name)) = generated.sparse { (VectorInput::new_sparse(indices.data, values), Some(name)) } else if let Some((query, name)) = generated.dense { @@ -212,7 +215,10 @@ fn infer_search_config(name: &str, config: &CollectionConfig) -> Result