diff --git a/Cargo.lock b/Cargo.lock index f99304b..7c861dd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1472,7 +1472,7 @@ dependencies = [ [[package]] name = "qdrant-client" version = "1.16.1-dev" -source = "git+https://github.com/qdrant/rust-client?branch=dev#18002abfd22c1d484a68eac3a1894b456f098dde" +source = "git+https://github.com/qdrant/rust-client?branch=dev#71143d873ffe1030bf712da8b0d67613b12ccee6" dependencies = [ "anyhow", "derive_builder", diff --git a/README.md b/README.md index ae5013e..ba63a53 100644 --- a/README.md +++ b/README.md @@ -346,6 +346,52 @@ 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. +#### BM25 over a text index + +A `text` request ranks points by BM25 over a payload field's text index, the +`Query.text` variant of Qdrant 1.19.3+. `using` names the payload field, whose +index needs `scoring: bm25` in the upload config: + +```yaml +requests: + - kind: text + using: body # payload field with a scoring text index + k: 1.2 # optional, BM25 term-frequency saturation + b: 0.75 # optional, BM25 length normalization, 0..1 + source: # random words, or `type: file` + type: random + vocab_size: 20000 + min_length: 2 + max_length: 5 + distribution: zipf +``` + +```yaml +fields: + - name: body + type: text + tokenizer: word + scoring: bm25 # rank by this field; implies phrase matching + lowercase: true # optional tokenizer options + stopwords: english + stemmer: english +``` + +A random source draws `word_` words, the vocabulary text payloads are generated +from, so the queries hit the corpus. `source: { type: file, path: queries.txt }` +reads one query per line instead (blank lines skipped; `strategy: from-start` or +`random-sample`), for a real query set. Matching is any-term: a point scores when it +holds one of the query's words, and terms it must or must not hold belong in +`filters`. BM25 over a text index is exact, so `--search-quality` has no exact +reference to compare against and is refused with a `text` request, and so is +`--prefetch`, which rescores a vector search. `lowercase`, `stopwords`, `stemmer` +and `scoring` apply to `type: text` fields only. + +```bash +bfb upload --example upload-bm25-text -n 200k --uri http://localhost:6334 +bfb search --example search-bm25-text -n 10k --search-limit 10 --uri http://localhost:6334 +``` + #### ACORN `--acorn` asks Qdrant to use ACORN for filtered HNSW search, and diff --git a/examples/search-bm25-text.yaml b/examples/search-bm25-text.yaml new file mode 100644 index 0000000..4cc411d --- /dev/null +++ b/examples/search-bm25-text.yaml @@ -0,0 +1,21 @@ +# Example bfb search config for BM25 over a text index (Qdrant 1.19.3+), matching +# upload-bm25-text. A point scores when it holds any of the query's words; terms +# a point must or must not hold belong in `filters`. +# +# bfb search --example search-bm25-text -n 10k --search-limit 10 -p 8 \ +# --uri http://localhost:6334 + +collection: + name: benchmark + +requests: + - kind: text + using: body # the payload field whose text index ranks points + # k: 1.2 # BM25 term-frequency saturation (server default) + # b: 0.75 # BM25 length normalization, 0..1 (server default) + source: + type: random # random | file (one query per line) + vocab_size: 20000 + min_length: 2 # words per query + max_length: 5 + distribution: zipf diff --git a/examples/upload-bm25-text.yaml b/examples/upload-bm25-text.yaml new file mode 100644 index 0000000..6af6fd8 --- /dev/null +++ b/examples/upload-bm25-text.yaml @@ -0,0 +1,30 @@ +# Example bfb upload config for BM25 over a text index (Qdrant 1.19.3+): every +# point carries a zipf-distributed text field indexed with `scoring: bm25`, so +# `kind: text` search requests can rank by it (see search-bm25-text). The dense +# vector is a placeholder the text query never reads. +# +# bfb upload --example upload-bm25-text -n 200k -b 256 -p 8 -t 4 \ +# --uri http://localhost:6334 + +collection: + name: benchmark + + vectors: + - size: 4 + distance: dot + source: random + + fields: + - name: body + type: text + tokenizer: word + scoring: bm25 # rank by BM25; implies phrase matching + # lowercase: true # tokenizer options, as for any text index + # stopwords: english + # stemmer: english + source: + type: random + vocab_size: 20000 + min_length: 20 # words per document + max_length: 200 + distribution: zipf diff --git a/src/collection/from_config.rs b/src/collection/from_config.rs index 2082aac..6338345 100644 --- a/src/collection/from_config.rs +++ b/src/collection/from_config.rs @@ -27,7 +27,7 @@ use crate::args::Args; use crate::client::random_client; use crate::config::{ ComparatorKind, DatatypeKind, DistanceKind, MemoryKind, ModifierKind, PayloadType, QuantKind, - TokenizerKind, UploadConfig, VectorConfig, + TextScoringKind, TokenizerKind, UploadConfig, VectorConfig, }; fn distance_to_grpc(distance: DistanceKind) -> Distance { @@ -442,6 +442,18 @@ async fn create_field_indices_from_config( if let Some(memory) = memory { params = params.memory(memory); } + if let Some(lowercase) = pc.lowercase { + params = params.lowercase(lowercase); + } + if let Some(language) = &pc.stopwords { + params = params.stopwords_language(language.clone()); + } + if let Some(language) = &pc.stemmer { + params = params.snowball_stemmer(language.clone()); + } + if let Some(TextScoringKind::Bm25) = pc.scoring { + params = params.bm25_scoring(); + } CreateFieldIndexCollectionBuilder::new(name, field, FieldType::Text) .field_index_params(params) } diff --git a/src/config/examples.rs b/src/config/examples.rs index 458952b..a21f35e 100644 --- a/src/config/examples.rs +++ b/src/config/examples.rs @@ -100,6 +100,16 @@ pub static EXAMPLES: &[Example] = &[ Search, "Search requests matching upload-config (dense, sparse, filters)" ), + example!( + "upload-bm25-text", + Upload, + "A text field indexed for BM25 ranking (Qdrant 1.19.3+)" + ), + example!( + "search-bm25-text", + Search, + "BM25 over a text index, matching upload-bm25-text" + ), example!( "search-dataset-accuracy", Search, diff --git a/src/config/mod.rs b/src/config/mod.rs index 1524c0c..7c24a5e 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -15,7 +15,7 @@ pub mod search; pub mod vector; pub use collection::{CollectionConfig, IdType, MemoryKind, QuantKind}; -pub use payload::{PayloadSource, PayloadSourceKind, PayloadType, TokenizerKind}; +pub use payload::{PayloadSource, PayloadSourceKind, PayloadType, TextScoringKind, TokenizerKind}; pub use vector::{ ComparatorKind, DatatypeKind, DistanceKind, DistributionKind, FileStrategy, ModifierKind, SparseKind, SparseSource, VectorConfig, VectorSource, @@ -136,6 +136,33 @@ impl UploadConfig { if !payload_names.insert(p.name.clone()) { bail!("duplicate payload field name {:?}", p.name); } + // Text index options describe a text index; anywhere else they would be + // silently dropped. + let text_options = p.lowercase.is_some() + || p.stopwords.is_some() + || p.stemmer.is_some() + || p.scoring.is_some(); + if text_options && p.kind != PayloadType::Text { + bail!( + "payload {:?}: `lowercase`, `stopwords`, `stemmer` and `scoring` apply to \ + `type: text` fields only", + p.name + ); + } + if text_options && !p.index { + bail!( + "payload {:?}: text index options need `index: true`", + p.name + ); + } + for (key, value) in [("stopwords", &p.stopwords), ("stemmer", &p.stemmer)] { + if value.as_deref().is_some_and(str::is_empty) { + bail!( + "payload {:?}: `{key}` names a language and must not be empty", + p.name + ); + } + } let Some(src) = &p.source else { // No per-field source: values come from `payload.source` (or, // absent that, are randomly generated). Nothing to validate. @@ -280,6 +307,50 @@ fn default_custom() -> String { mod tests { use super::*; + #[test] + fn text_index_options_apply_to_text_fields_only() { + let config = + |field: &str| format!("collection:\n vectors:\n - size: 4\n fields:\n{field}"); + let ok = parse( + &config(" - name: body\n type: text\n lowercase: true\n stopwords: english\n stemmer: english\n scoring: bm25\n source: random\n"), + "test", + ) + .unwrap(); + let body = &ok.collection.fields[0]; + assert_eq!(body.scoring, Some(TextScoringKind::Bm25)); + assert_eq!( + ( + body.stopwords.as_deref(), + body.stemmer.as_deref(), + body.lowercase + ), + (Some("english"), Some("english"), Some(true)) + ); + let err = parse( + &config(" - name: color\n type: keyword\n scoring: bm25\n source: random\n"), + "test", + ) + .unwrap_err(); + assert!( + format!("{err:#}").contains("`type: text` fields only"), + "{err:#}" + ); + let err = parse( + &config(" - name: body\n type: text\n index: false\n scoring: bm25\n source: random\n"), + "test", + ) + .unwrap_err(); + assert!(format!("{err:#}").contains("need `index: true`"), "{err:#}"); + let err = parse( + &config( + " - name: body\n type: text\n stemmer: ''\n source: random\n", + ), + "test", + ) + .unwrap_err(); + assert!(format!("{err:#}").contains("must not be empty"), "{err:#}"); + } + #[test] fn parses_minimal_config() { let yaml = r#" diff --git a/src/config/payload.rs b/src/config/payload.rs index 6457839..997bb89 100644 --- a/src/config/payload.rs +++ b/src/config/payload.rs @@ -36,6 +36,19 @@ pub struct PayloadConfig { pub prefix: bool, /// Text payloads: tokenizer. pub tokenizer: Option, + /// Text payloads: lowercase tokens. Server default: true. + #[serde(default)] + pub lowercase: Option, + /// Text payloads: drop the stopwords of this language, e.g. `english`. + #[serde(default)] + pub stopwords: Option, + /// Text payloads: stem tokens with this language's Snowball stemmer, e.g. `english`. + #[serde(default)] + pub stemmer: Option, + /// Text payloads: rank points by this field, so `kind: text` search requests + /// can query it (Qdrant 1.19.3+). Implies phrase matching. + #[serde(default)] + pub scoring: Option, /// Per-field value source. May be omitted when `collection.payload.source` /// provides the payload object — the entry is then index-only. #[serde(default, deserialize_with = "option_string_or_struct")] @@ -55,6 +68,13 @@ pub enum PayloadType { Datetime, } +/// How a text index ranks points for `kind: text` queries. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum TextScoringKind { + Bm25, +} + #[derive(Debug, Clone, Copy, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum TokenizerKind { diff --git a/src/config/schema.rs b/src/config/schema.rs index 4c2b715..d9a1a69 100644 --- a/src/config/schema.rs +++ b/src/config/schema.rs @@ -197,6 +197,10 @@ collection: prefix: false # bool default=false keyword payloads: enable prefix matching # (required for search `match_prefix` filters) tokenizer: null # enum optional (text) [word | whitespace | prefix | multilingual] + lowercase: null # bool optional (text) lowercase tokens (server default true) + stopwords: null # string optional (text) stopword language, e.g. english + stemmer: null # string optional (text) Snowball stemmer language, e.g. english + scoring: null # enum optional (text) [bm25] rank by this field (search `kind: text`) # Value source (optional when `payload.source` is set — then the entry is # index-only). Shorthand string `random` / `random-clusters` / `now`, or a # map. Which keys apply depends on the payload `type`; others ignored. @@ -303,34 +307,59 @@ mod tests { source: None, memory: Some(MemoryKind::Cached), }, - fields: vec![PayloadConfig { - name: "color".to_string(), - kind: PayloadType::Keyword, - index: true, - on_disk: false, - memory: Some(MemoryKind::Cached), - is_tenant: false, - is_principal: false, - range_index: true, - prefix: true, - tokenizer: Some(TokenizerKind::Word), - source: Some(PayloadSource { - kind: PayloadSourceKind::Random, - dataset: None, - field: None, - cardinality: Some(100), - length_multiplier: Some(1), - values_per_point: Some(1), - min: Some(0.0), - max: Some(1.0), - true_ratio: Some(0.5), - clusters: Some(10), - vocab_size: Some(1000), - min_length: Some(1), - max_length: Some(10), - distribution: DistributionKind::Uniform, - }), - }], + fields: vec![ + PayloadConfig { + name: "color".to_string(), + kind: PayloadType::Keyword, + index: true, + on_disk: false, + memory: Some(MemoryKind::Cached), + is_tenant: false, + is_principal: false, + range_index: true, + prefix: true, + tokenizer: Some(TokenizerKind::Word), + lowercase: None, + stopwords: None, + stemmer: None, + scoring: None, + source: Some(PayloadSource { + kind: PayloadSourceKind::Random, + dataset: None, + field: None, + cardinality: Some(100), + length_multiplier: Some(1), + values_per_point: Some(1), + min: Some(0.0), + max: Some(1.0), + true_ratio: Some(0.5), + clusters: Some(10), + vocab_size: Some(1000), + min_length: Some(1), + max_length: Some(10), + distribution: DistributionKind::Uniform, + }), + }, + // The text index options are refused on any other type, so they + // need a text field of their own. + PayloadConfig { + name: "body".to_string(), + kind: PayloadType::Text, + index: true, + on_disk: false, + memory: None, + is_tenant: false, + is_principal: false, + range_index: true, + prefix: false, + tokenizer: Some(TokenizerKind::Word), + lowercase: Some(true), + stopwords: Some("english".to_string()), + stemmer: Some("english".to_string()), + scoring: Some(TextScoringKind::Bm25), + source: None, + }, + ], }, } } diff --git a/src/config/search.rs b/src/config/search.rs index ccefec6..cf44241 100644 --- a/src/config/search.rs +++ b/src/config/search.rs @@ -8,8 +8,9 @@ use anyhow::{Context, Result, bail}; use serde::{Deserialize, Serialize}; use crate::config::{ - DatatypeKind, PayloadSource, PayloadType, SparseKind, SparseSource, VectorSource, - default_collection_name, string_or_struct, validate_file_source_path, + DatatypeKind, DistributionKind, FileStrategy, PayloadSource, PayloadType, SparseKind, + SparseSource, VectorSource, default_collection_name, string_or_struct, + validate_file_source_path, }; /// Top-level document: `{ collection: { name }, requests: [ … ] }`. @@ -78,6 +79,68 @@ pub enum SearchRequestConfig { #[serde(default)] idf_corpus: Vec, }, + /// BM25 over a payload field's text index (Qdrant 1.19.3+ `Query.text`). The + /// field's index needs `scoring: bm25`. A point scores when it holds any of the + /// query's terms; required or excluded terms belong in `filters`. + Text { + /// Payload field whose text index ranks the points. + using: String, + #[serde(default)] + source: TextQuerySource, + #[serde(default)] + filters: Vec, + /// BM25 term-frequency saturation (`k1`). Server default: 1.2. + #[serde(default)] + k: Option, + /// BM25 document-length normalization, from 0 (none) to 1 (full). Server + /// default: 0.75. + #[serde(default)] + b: Option, + }, +} + +/// Where a `kind: text` request's query strings come from. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "kebab-case", deny_unknown_fields)] +pub enum TextQuerySource { + /// Words drawn as `word_`, the vocabulary text payloads are generated from. + Random { + #[serde(default = "default_text_vocab_size")] + vocab_size: usize, + /// Words per query. + #[serde(default = "default_text_query_length")] + min_length: usize, + #[serde(default = "default_text_query_length")] + max_length: usize, + #[serde(default)] + distribution: DistributionKind, + }, + /// A UTF-8 file with one query per line; blank lines are skipped. A local path + /// or an http(s):// URL. + File { + path: String, + #[serde(default)] + strategy: FileStrategy, + }, +} + +impl Default for TextQuerySource { + fn default() -> Self { + TextQuerySource::Random { + vocab_size: default_text_vocab_size(), + min_length: default_text_query_length(), + max_length: default_text_query_length(), + distribution: DistributionKind::default(), + } + } +} + +fn default_text_vocab_size() -> usize { + crate::generators::random::DEFAULT_VOCAB_SIZE +} + +fn default_text_query_length() -> usize { + 2 } /// Shape of a multivector query. @@ -217,14 +280,15 @@ impl SearchRequestConfig { pub fn prefetches(&self) -> &[PrefetchConfig] { match self { SearchRequestConfig::Dense { prefetch, .. } => prefetch, - SearchRequestConfig::Sparse { .. } => &[], + SearchRequestConfig::Sparse { .. } | SearchRequestConfig::Text { .. } => &[], } } pub fn filters(&self) -> &[FilterPayloadConfig] { match self { SearchRequestConfig::Dense { filters, .. } - | SearchRequestConfig::Sparse { filters, .. } => filters, + | SearchRequestConfig::Sparse { filters, .. } + | SearchRequestConfig::Text { filters, .. } => filters, } } @@ -354,6 +418,49 @@ impl SearchRequestConfig { } } } + SearchRequestConfig::Text { + using, + source, + k, + b, + .. + } => { + if using.is_empty() { + bail!( + "requests[{index}]: text `using` names the payload field and must not be empty" + ); + } + if let Some(k) = k + && !(k.is_finite() && *k >= 0.0) + { + bail!("requests[{index}]: text `k` must be >= 0, got {k}"); + } + if let Some(b) = b + && !(0.0..=1.0).contains(b) + { + bail!("requests[{index}]: text `b` must be within [0, 1], got {b}"); + } + match source { + TextQuerySource::Random { + vocab_size, + min_length, + max_length, + .. + } => { + if *vocab_size == 0 { + bail!("requests[{index}]: text `vocab_size` must be > 0"); + } + if *min_length == 0 || max_length < min_length { + bail!( + "requests[{index}]: text `min_length` must be > 0 and <= \ + `max_length`" + ); + } + } + TextQuerySource::File { path, .. } => validate_file_source_path(path) + .with_context(|| format!("requests[{index}]"))?, + } + } SearchRequestConfig::Sparse { using, source, .. } => { if using.is_empty() { bail!("requests[{index}]: sparse `using` must not be empty"); @@ -387,6 +494,82 @@ mod tests { use super::*; use crate::config::DistributionKind; + #[test] + fn parses_text_request() { + let yaml = r#" +collection: + name: docs +requests: + - kind: text + using: body + k: 1.5 + b: 0.5 + source: { type: random, vocab_size: 100, min_length: 2, max_length: 4, distribution: zipf } + filters: + - name: color + type: keyword + source: { cardinality: 5 } + - kind: text + using: title +"#; + let cfg = parse(yaml, "test").unwrap(); + let SearchRequestConfig::Text { + using, + k, + b, + source, + filters, + } = &cfg.requests[0] + else { + panic!("expected a text request"); + }; + assert_eq!((using.as_str(), *k, *b), ("body", Some(1.5), Some(0.5))); + assert!(matches!( + source, + TextQuerySource::Random { + vocab_size: 100, + min_length: 2, + max_length: 4, + .. + } + )); + assert_eq!(filters.len(), 1); + // Defaults: generated two-word queries, server-side k and b. + let SearchRequestConfig::Text { source, k, b, .. } = &cfg.requests[1] else { + panic!("expected a text request"); + }; + assert!(matches!( + source, + TextQuerySource::Random { + min_length: 2, + max_length: 2, + .. + } + )); + assert_eq!((*k, *b), (None, None)); + } + + #[test] + fn rejects_bad_text_request() { + let bad = |body: &str| { + let yaml = format!("collection:\n name: x\nrequests:\n - kind: text\n{body}"); + format!("{:#}", parse(&yaml, "test").unwrap_err()) + }; + assert!(bad(" using: ''\n").contains("must not be empty")); + assert!(bad(" using: body\n k: -1.0\n").contains("`k` must be >= 0")); + assert!(bad(" using: body\n b: 1.5\n").contains("within [0, 1]")); + assert!( + bad(" using: body\n source: { type: random, min_length: 3, max_length: 2 }\n") + .contains("min_length") + ); + assert!( + bad(" using: body\n source: { type: file, path: /no/such/queries.txt }\n") + .contains("not found") + ); + // `using` is the field the index is on, so it cannot be left to a default. + assert!(bad(" k: 1.0\n").contains("using")); + } + const RESCORE_YAML: &str = r#" collection: name: bench diff --git a/src/generators/queries.rs b/src/generators/queries.rs index 7c8f9cd..2f43779 100644 --- a/src/generators/queries.rs +++ b/src/generators/queries.rs @@ -20,6 +20,7 @@ use super::random::{ }; use crate::config::search::{ FilterPayloadConfig, FusionKind, PrefetchKind, SearchConfig, SearchRequestConfig, + TextQuerySource, }; use crate::config::{ DatatypeKind, DistributionKind, FileStrategy, PayloadSourceKind, PayloadType, SparseKind, @@ -73,11 +74,23 @@ pub struct FusionSpec { pub weights: Option>, } +/// A BM25 query over a payload field's text index. +#[derive(Debug, Clone, PartialEq)] +pub struct GeneratedText { + pub query: String, + /// The payload field, sent as the request's `using`. + pub using: String, + pub k: Option, + pub b: Option, +} + /// One query vector plus optional filter, ready to be turned into a gRPC request. #[derive(Debug, Clone)] pub struct GeneratedQuery { pub dense: Option<(DenseQuery, Option)>, pub sparse: Option<(Vec, SparseIndices, String)>, + /// A text (BM25) query; the request then sends no vector. + pub text: Option, pub filter: Option, /// Sparse-vector IDF corpus: restricts which points the IDF statistics are /// computed over. `None` ⇒ collection-wide (global) statistics. @@ -160,6 +173,10 @@ struct RequestState { /// Reference dataset used as the query source (dense or sparse). query_dataset: Option, sparse_zipf: Option>, + /// Text requests with a file source: every query, read once at startup. + text_queries: Option>, + /// Text requests with a zipf-distributed random source. + text_zipf: Option>, filters: FilterGenerator, /// Sparse requests only: conditions defining the IDF corpus. idf_corpus: FilterGenerator, @@ -364,7 +381,26 @@ impl ConfigSearchGenerator { let mut per_request = Vec::with_capacity(config.requests.len()); for req in &config.requests { + let (text_queries, text_zipf) = match req { + SearchRequestConfig::Text { + source: TextQuerySource::File { path, .. }, + .. + } => (Some(Self::read_text_queries(datasets_dir, path)?), None), + SearchRequestConfig::Text { + source: + TextQuerySource::Random { + vocab_size, + distribution: DistributionKind::Zipf, + .. + }, + .. + } => (None, Some(create_zipf(*vocab_size))), + _ => (None, None), + }; let (dense_reader, query_dataset, sparse_zipf, filters, idf_corpus) = match req { + SearchRequestConfig::Text { filters, .. } => { + (None, None, None, filters, [].as_slice()) + } SearchRequestConfig::Dense { source, filters, @@ -429,6 +465,8 @@ impl ConfigSearchGenerator { dense_reader, query_dataset, sparse_zipf, + text_queries, + text_zipf, filters: FilterGenerator::new(filters, &mut rng), idf_corpus: FilterGenerator::new(idf_corpus, &mut rng), }); @@ -440,6 +478,25 @@ impl ConfigSearchGenerator { }) } + /// A text request's query file, one query per line, read once so file I/O + /// stays out of the timed path. Blank lines are skipped; an empty file fails + /// here rather than part-way through a benchmark. + fn read_text_queries(datasets_dir: &Path, path: &str) -> anyhow::Result> { + let local = ensure_local_file(datasets_dir, path)?; + let text = std::fs::read_to_string(&local) + .with_context(|| format!("failed to read text queries from {}", local.display()))?; + let queries: Vec = text + .lines() + .map(str::trim) + .filter(|line| !line.is_empty()) + .map(str::to_string) + .collect(); + if queries.is_empty() { + anyhow::bail!("text query file {} holds no queries", local.display()); + } + Ok(queries) + } + /// Open a reference dataset as a query source and read its entire query set /// into memory, requiring it to ship one (ann-benchmarks `test`/`neighbors`, /// or `tests.jsonl` / `queries.csr`+`results.gt`). @@ -660,6 +717,7 @@ impl ConfigSearchGenerator { // A fusion query ranks its prefetches and sends no vector itself. dense: fusion.is_none().then(|| (vector, using.clone())), sparse: None, + text: None, filter: dataset_filter.or_else(|| state.filters.build(rng)), idf_corpus: None, expected_ids, @@ -691,12 +749,68 @@ impl ConfigSearchGenerator { dense: None, fusion: None, sparse: Some((values, indices, using.clone())), + text: None, filter: dataset_filter.or_else(|| state.filters.build(rng)), idf_corpus: state.idf_corpus.build(rng), expected_ids, prefetch: Vec::new(), } } + SearchRequestConfig::Text { + using, + source, + filters: _, + k, + b, + } => { + let query = match source { + TextQuerySource::File { strategy, .. } => { + let queries = state + .text_queries + .as_ref() + .expect("a file source is read at startup"); + let idx = match strategy { + FileStrategy::FromStart => req_id % queries.len(), + FileStrategy::RandomSample => rng.random_range(0..queries.len()), + }; + queries[idx].clone() + } + TextQuerySource::Random { + vocab_size, + min_length, + max_length, + .. + } => { + let len = if max_length > min_length { + rng.random_range(*min_length..=*max_length) + } else { + *min_length + }; + match &state.text_zipf { + Some(zipf) => random_text(rng, len, zipf), + None => (0..len) + .map(|_| format!("word_{}", rng.random_range(0..*vocab_size))) + .collect::>() + .join(" "), + } + } + }; + GeneratedQuery { + dense: None, + sparse: None, + text: Some(GeneratedText { + query, + using: using.clone(), + k: *k, + b: *b, + }), + filter: state.filters.build(rng), + idf_corpus: None, + expected_ids: None, + prefetch: Vec::new(), + fusion: None, + } + } } } @@ -815,6 +929,63 @@ mod tests { ConfigSearchGenerator::new(&config).unwrap() } + #[test] + fn generates_random_text_queries() { + let generator = build_gen( + "collection:\n name: x\nrequests:\n - kind: text\n using: body\n k: 2.0\n source: { type: random, vocab_size: 50, min_length: 3, max_length: 5, distribution: zipf }\n", + ); + let mut rng = rand::rng(); + for _ in 0..20 { + let q = generator.make_query(0, &mut rng); + assert!(q.dense.is_none() && q.sparse.is_none() && q.fusion.is_none()); + let text = q.text.expect("a text request makes a text query"); + assert_eq!( + (text.using.as_str(), text.k, text.b), + ("body", Some(2.0), None) + ); + let words: Vec<&str> = text.query.split(' ').collect(); + assert!((3..=5).contains(&words.len()), "{}", text.query); + assert!(words.iter().all(|w| w.starts_with("word_"))); + } + } + + #[test] + fn reads_text_queries_from_a_file_in_order() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("queries.txt"); + std::fs::write( + &path, + "what is bm25\n\n how do stemmers work \nlast one\n", + ) + .unwrap(); + let yaml = format!( + "collection:\n name: x\nrequests:\n - kind: text\n using: body\n source: {{ type: file, path: {}, strategy: from-start }}\n", + path.display() + ); + let generator = build_gen(&yaml); + let mut rng = rand::rng(); + let got: Vec = (0..4) + .map(|req_id| generator.make_query(req_id, &mut rng).text.unwrap().query) + .collect(); + // Blank lines skipped, each line trimmed, and the set wraps around. + assert_eq!( + got, + [ + "what is bm25", + "how do stemmers work", + "last one", + "what is bm25" + ] + ); + + std::fs::write(&path, "\n \n").unwrap(); + let config: SearchConfig = serde_yaml::from_str(&yaml).unwrap(); + let err = ConfigSearchGenerator::new(&config) + .err() + .expect("an empty file fails"); + assert!(err.to_string().contains("holds no queries"), "{err}"); + } + #[test] fn generates_multivector_query_with_prefetch() { let generator = build_gen( diff --git a/src/search/from_config.rs b/src/search/from_config.rs index 2c50614..4c4303b 100644 --- a/src/search/from_config.rs +++ b/src/search/from_config.rs @@ -9,17 +9,25 @@ use qdrant_client::qdrant::shard_key::Key; use qdrant_client::qdrant::{ AcornSearchParamsBuilder, Fusion, IdfParamsBuilder, PrefetchQueryBuilder, QuantizationSearchParamsBuilder, Query, QueryBatchPointsBuilder, QueryPointsBuilder, Rrf, - SearchParams, SearchParamsBuilder, VectorInput, + SearchParams, SearchParamsBuilder, TextQueryBuilder, 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::{FusionKind, SearchConfig}; +use crate::config::search::{FusionKind, SearchConfig, SearchRequestConfig}; use crate::generators::ConfigSearchGenerator; -use crate::generators::queries::{FusionSpec, GeneratedPrefetch, PrefetchQuery}; +use crate::generators::queries::{FusionSpec, GeneratedPrefetch, GeneratedText, PrefetchQuery}; use crate::processor::{Processor, Timing}; +/// What a request ranks by, unless it fuses its prefetches. +pub(crate) enum MainQuery { + /// A dense or sparse vector, and the named vector it searches. + Vector(VectorInput, Option), + /// BM25 over a payload field's text index. + Text(GeneratedText), +} + /// YAML-config-driven search (`bfb search --file config.yaml`). pub struct ConfigSearchProcessor { args: Args, @@ -58,6 +66,25 @@ impl ConfigSearchProcessor { ); } } + let has_text = config + .requests + .iter() + .any(|r| matches!(r, SearchRequestConfig::Text { .. })); + if has_text { + // The comparison runs each query again with `exact: true`, and BM25 over a + // text index is exact already: it would compare the ranking with itself. + if args.search_quality { + anyhow::bail!( + "--search-quality has no exact reference for a `kind: text` request: \ + BM25 over a text index is exact already" + ); + } + if args.prefetch.is_some() { + anyhow::bail!( + "--prefetch rescores a vector search; a `kind: text` request has no vector" + ); + } + } // ACORN is a filtered-search path: without filters Qdrant never takes it, so a run // with the flag on would report the ordinary path under ACORN's name. if args.acorn && config.requests.iter().all(|r| r.filters().is_empty()) { @@ -127,7 +154,7 @@ impl ConfigSearchProcessor { &self, query_filter: Option, idf_corpus: Option, - vector: Option<(VectorInput, Option)>, + main: Option, prefetch: Vec, fusion: Option, mut search_params: SearchParamsBuilder, @@ -178,14 +205,11 @@ impl ConfigSearchProcessor { prefetch_params = prefetch_params.hnsw_ef(hnsw_ef as u64); } + let Some(MainQuery::Vector(input, _)) = main.as_ref() else { + panic!("--prefetch needs the request's own vector"); + }; let stage = PrefetchQueryBuilder::default() - .query(Query::new_nearest( - vector - .as_ref() - .expect("--prefetch needs the request's own vector") - .0 - .clone(), - )) + .query(Query::new_nearest(input.clone())) .params(prefetch_params) .limit(prefetch_limit as u64) .build(); @@ -207,13 +231,27 @@ impl ConfigSearchProcessor { }), }, }), - 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); + None => match main.expect("a query without fusion needs a vector or a text query") { + MainQuery::Text(text) => { + let mut query = TextQueryBuilder::new(text.query); + if let Some(k) = text.k { + query = query.k(k); + } + if let Some(b) = text.b { + query = query.b(b); + } + // `using` names the payload field whose text index ranks the points. + request_builder + .using(text.using) + .query(Query::new_text(query)) } - request_builder.query(Query::new_nearest(vector)) - } + MainQuery::Vector(vector, name) => { + 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 { @@ -265,23 +303,28 @@ impl ConfigSearchProcessor { .map(|generated| { let query_filter = generated.filter.clone(); let idf_corpus = generated.idf_corpus.clone(); - let vector = if let Some((values, indices, name)) = generated.sparse { - Some((VectorInput::new_sparse(indices.data, values), Some(name))) + let main = if let Some(text) = generated.text { + Some(MainQuery::Text(text)) + } else if let Some((values, indices, name)) = generated.sparse { + Some(MainQuery::Vector( + VectorInput::new_sparse(indices.data, values), + Some(name), + )) } else { // A fusion query has no vector of its own; anything else must have one. generated .dense - .map(|(query, name)| (query.into_vector_input(), name)) + .map(|(query, name)| MainQuery::Vector(query.into_vector_input(), name)) }; assert!( - vector.is_some() || generated.fusion.is_some(), - "a request without fusion must produce a dense or sparse vector" + main.is_some() || generated.fusion.is_some(), + "a request without fusion must produce a dense or sparse vector, or a text query" ); self.create_request_builder( query_filter, idf_corpus, - vector, + main, generated.prefetch, generated.fusion, search_params.clone(), @@ -489,6 +532,54 @@ 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"; + const TEXT_YAML: &str = "collection:\n name: x\nrequests:\n - kind: text\n using: body\n k: 1.5\n b: 0.25\n filters:\n - name: color\n type: keyword\n source: { cardinality: 5 }\n"; + + fn text_processor(extra: &[&str]) -> anyhow::Result { + let mut argv = vec!["bfb"]; + argv.extend_from_slice(extra); + ConfigSearchProcessor::new( + Args::parse_from(argv), + &parse(TEXT_YAML, "test").unwrap(), + Arc::new(AtomicBool::new(false)), + vec![], + ) + } + + /// A text request ranks by BM25 over the payload field named in `using`, with + /// its own k and b, and sends no vector. + #[test] + fn builds_a_text_request() { + let p = text_processor(&[]).unwrap(); + let q = p.generator.make_query(0, &mut rand::rng()); + let request = p + .create_request_builder( + q.filter, + None, + q.text.map(MainQuery::Text), + q.prefetch, + q.fusion, + SearchParamsBuilder::default(), + ) + .build(); + assert_eq!(request.using.as_deref(), Some("body")); + assert!(request.filter.is_some()); + let Some(QueryVariant::Text(text)) = request.query.and_then(|q| q.variant) else { + panic!("expected a text query"); + }; + assert_eq!((text.k, text.b), (Some(1.5), Some(0.25))); + assert_eq!(text.query.split(' ').count(), 2); + } + + #[test] + fn text_requests_refuse_flags_that_need_a_vector() { + let err = text_processor(&["--search-quality"]).err().unwrap(); + assert!(err.to_string().contains("exact already"), "{err}"); + let err = text_processor(&["--prefetch", "100"]).err().unwrap(); + assert!(err.to_string().contains("no vector"), "{err}"); + // Filters are allowed, so is ACORN on them. + assert!(text_processor(&["--acorn"]).is_ok()); + } + /// A fused request sends its stages and how to combine them, and no vector of its own. #[test] fn builds_a_hybrid_fusion_request() { @@ -642,7 +733,7 @@ mod tests { .create_request_builder( None, None, - Some((query.into_vector_input(), using)), + Some(MainQuery::Vector(query.into_vector_input(), using)), q.prefetch, q.fusion, // As search() builds it from --search-hnsw-ef. @@ -677,7 +768,7 @@ mod tests { .create_request_builder( None, None, - Some((query.into_vector_input(), using)), + Some(MainQuery::Vector(query.into_vector_input(), using)), q.prefetch, q.fusion, SearchParamsBuilder::default().exact(true).hnsw_ef(64),