diff --git a/README.md b/README.md index 38d1f6a..f62b922 100644 --- a/README.md +++ b/README.md @@ -267,12 +267,32 @@ Query set + ground truth are auto-detected from the dataset files: | `format` | Query vectors | Ground truth | |----------|---------------|--------------| | `h5` (ann-benchmarks) | `test` dataset | `neighbors` dataset | -| `tar` (ann-filtering-benchmark-datasets) | `tests.jsonl` `query` | `tests.jsonl` `closest_ids` | +| `tar` (ann-filtering-benchmark-datasets) | `tests.jsonl` `query` | `tests.jsonl` `closest_ids`, under `tests.jsonl` `conditions` | | `sparse` | `queries.csr` | `results.gt` | +**Filtered query sets.** A `tar` dataset may pair each query with the filter it +was answered under: + +```json +{"query": [..], "conditions": {"and": [{"similarity": {"range": {"gt": 0.34}}}]}, + "closest_ids": [..]} +``` + +Those `closest_ids` are the ground truth *for the filtered query*, so bfb +applies each query's own conditions — `and` as `must`, `or` as `should`, over +`match` / `range` / `geo`, the dialect vector-db-benchmark defines. Searching +such a set unfiltered scores against the answers to a different question: on +`laion-small-clip`, whose 5000 queries are half conditioned, that reads as +recall 0.63 instead of 0.99. bfb prints how many queries carry conditions when +it opens the query set. A `filters:` block on the same request is ignored when +the dataset supplies a filter; sets without conditions (`{}`, as in the +`-no-filters` datasets) are unaffected. + Accuracy only lines up when the corpus was uploaded with the default integer id scheme (point id == dataset row index), which is what `bfb upload` does for -`id: integer` collections. See +`id: integer` collections. The payload fields the conditions filter on must be +uploaded and indexed — the upload config for a filtered dataset needs its +`payload.source` and the matching `fields:` entries. See [`examples/search-dataset-accuracy.yaml`](examples/search-dataset-accuracy.yaml). ### `scroll` — run a scroll workload as its own phase diff --git a/examples/search-dataset-accuracy.yaml b/examples/search-dataset-accuracy.yaml index 2ebf7ff..0f1c241 100644 --- a/examples/search-dataset-accuracy.yaml +++ b/examples/search-dataset-accuracy.yaml @@ -13,6 +13,13 @@ # returned point ids against the dataset's ground-truth nearest neighbors. # Recall is reported under "--- Precision ---" as `|found ∩ expected[:k]| / k`. # +# If the query set carries per-query `conditions` (the ann-filtering-benchmark +# datasets: laion-small-clip, arxiv-titles-*-filters, h-and-m-*-filters, the +# random-*-filters family), its ground truth answers the FILTERED query, so bfb +# applies each query's own conditions. A `filters:` block on the same request is +# then ignored — the dataset defines the filter. bfb prints how many queries +# carry conditions when it opens the query set. +# # IMPORTANT: accuracy only lines up when the corpus was uploaded with the # default integer id scheme (point id == dataset row index), which is what # `bfb upload` does for `id: integer` collections. diff --git a/src/dataset/conditions.rs b/src/dataset/conditions.rs new file mode 100644 index 0000000..8554754 --- /dev/null +++ b/src/dataset/conditions.rs @@ -0,0 +1,274 @@ +//! Per-query filter conditions from a dataset's query set. +//! +//! `tests.jsonl` in the ann-filtering-benchmark-datasets layout pairs every +//! query with the conditions it was answered under: +//! +//! ```json +//! {"query": [..], "conditions": {"and": [{"similarity": {"range": {"gt": 0.34}}}]}, +//! "closest_ids": [..], "closest_scores": [..]} +//! ``` +//! +//! The `closest_ids` are the ground truth **for the filtered query**, so a +//! benchmark that ignores the conditions scores an unfiltered search against +//! the answers to a different question. On `laion-small-clip` — half of whose +//! queries carry a `range` condition — that reads as recall 0.64 instead of +//! 0.99. +//! +//! The dialect is the one vector-db-benchmark defines in +//! `engine/base_client/parser.py`, so numbers stay comparable with its results: +//! +//! ```text +//! conditions := { "and": [entry, ..], "or": [entry, ..] } both optional +//! entry := { field_name: { condition_type: criteria } } +//! condition := "match" { "value": string|number } +//! | "range" { "lt"?, "gt"?, "lte"?, "gte"? } +//! | "geo" { "lat", "lon", "radius" } +//! ``` +//! +//! `and` maps to `must`, `or` to `should`. + +use anyhow::{Context, Result, bail}; +use qdrant_client::qdrant::{Condition, Filter, GeoPoint, GeoRadius, Range}; +use serde_json::Value; + +/// Parse a `conditions` object into a filter. +/// +/// Returns `None` for the query sets that have no conditions at all — `null`, +/// or the `{}` that the `-no-filters` datasets write on every row — so an +/// unfiltered dataset costs nothing and stays byte-identical to before. +pub fn parse(conditions: Option<&Value>) -> Result> { + let Some(value) = conditions else { + return Ok(None); + }; + if value.is_null() { + return Ok(None); + } + let object = value.as_object().context("`conditions` is not an object")?; + if object.is_empty() { + return Ok(None); + } + + let mut filter = Filter { + should: vec![], + must: vec![], + must_not: vec![], + min_should: None, + }; + for (operator, entries) in object { + let target = match operator.as_str() { + "and" => &mut filter.must, + "or" => &mut filter.should, + other => bail!("unknown `conditions` operator {other:?}; expected \"and\" or \"or\""), + }; + let entries = entries + .as_array() + .with_context(|| format!("`conditions.{operator}` is not an array"))?; + for entry in entries { + push_entry(entry, target)?; + } + } + + // `{"and": []}` carries no conditions; an empty filter would still be sent + // as a filtered search, which is not the same request. + if filter.must.is_empty() && filter.should.is_empty() { + return Ok(None); + } + Ok(Some(filter)) +} + +/// One `{field: {condition_type: criteria}}` entry. A single entry may name +/// several fields, and a field several condition types; every one becomes its +/// own condition, as in vector-db-benchmark's `create_condition_subfilters`. +fn push_entry(entry: &Value, out: &mut Vec) -> Result<()> { + let fields = entry + .as_object() + .context("`conditions` entry is not an object")?; + for (field, by_type) in fields { + let by_type = by_type + .as_object() + .with_context(|| format!("`conditions` entry for {field:?} is not an object"))?; + for (condition_type, criteria) in by_type { + out.push( + build(field, condition_type, criteria) + .with_context(|| format!("`conditions` on field {field:?}"))?, + ); + } + } + Ok(()) +} + +fn build(field: &str, condition_type: &str, criteria: &Value) -> Result { + match condition_type { + "match" => { + let value = criteria.get("value").context("`match` has no `value`")?; + match value { + Value::String(text) => Ok(Condition::matches(field, text.clone())), + Value::Bool(flag) => Ok(Condition::matches(field, *flag)), + Value::Number(number) => { + let int = number + .as_i64() + .context("`match.value` is a non-integer number")?; + Ok(Condition::matches(field, int)) + } + other => bail!("`match.value` must be a string, integer or bool, got {other}"), + } + } + "range" => { + let bound = |name: &str| -> Result> { + match criteria.get(name) { + None | Some(Value::Null) => Ok(None), + Some(value) => { + Ok(Some(value.as_f64().with_context(|| { + format!("`range.{name}` is not a number") + })?)) + } + } + }; + let range = Range { + lt: bound("lt")?, + gt: bound("gt")?, + gte: bound("gte")?, + lte: bound("lte")?, + }; + if range.lt.is_none() + && range.gt.is_none() + && range.gte.is_none() + && range.lte.is_none() + { + bail!("`range` has no bounds"); + } + Ok(Condition::range(field, range)) + } + "geo" => { + let number = |name: &str| -> Result { + criteria + .get(name) + .and_then(Value::as_f64) + .with_context(|| format!("`geo` is missing a numeric `{name}`")) + }; + Ok(Condition::geo_radius( + field, + GeoRadius { + center: Some(GeoPoint { + lat: number("lat")?, + lon: number("lon")?, + }), + radius: number("radius")? as f32, + }, + )) + } + other => { + bail!("unknown condition type {other:?}; expected \"match\", \"range\" or \"geo\"") + } + } +} + +#[cfg(test)] +mod tests { + use qdrant_client::qdrant::condition::ConditionOneOf; + use serde_json::json; + + use super::*; + + fn field_conditions(filter: &Filter) -> (usize, usize) { + (filter.must.len(), filter.should.len()) + } + + /// The `-no-filters` datasets write `{}` on every row, and older query sets + /// omit the field. Both must stay an unfiltered search, not an empty filter. + #[test] + fn absent_and_empty_conditions_are_unfiltered() { + assert!(parse(None).unwrap().is_none()); + assert!(parse(Some(&Value::Null)).unwrap().is_none()); + assert!(parse(Some(&json!({}))).unwrap().is_none()); + assert!(parse(Some(&json!({"and": []}))).unwrap().is_none()); + } + + /// The exact shape half of laion-small-clip's query set carries. + #[test] + fn parses_the_laion_range_condition() { + let filter = parse(Some(&json!({ + "and": [{"similarity": {"range": {"gt": 0.3491986757069205}}}] + }))) + .unwrap() + .expect("a range condition is a filter"); + assert_eq!(field_conditions(&filter), (1, 0)); + + let Some(ConditionOneOf::Field(field)) = &filter.must[0].condition_one_of else { + panic!("expected a field condition, got {:?}", filter.must[0]); + }; + assert_eq!(field.key, "similarity"); + let range = field.range.as_ref().expect("range"); + assert_eq!(range.gt, Some(0.3491986757069205)); + assert_eq!(range.gte, None); + assert_eq!(range.lt, None); + assert_eq!(range.lte, None); + } + + /// `and` is `must`, `or` is `should` — the mapping vector-db-benchmark uses. + #[test] + fn and_is_must_or_is_should() { + let filter = parse(Some(&json!({ + "and": [{"a": {"match": {"value": "x"}}}], + "or": [{"b": {"match": {"value": 80}}}, {"b": {"match": {"value": 2}}}] + }))) + .unwrap() + .unwrap(); + assert_eq!(field_conditions(&filter), (1, 2)); + } + + #[test] + fn parses_geo_and_multi_bound_range() { + let filter = parse(Some(&json!({ + "and": [ + {"loc": {"geo": {"lat": 52.5, "lon": 13.4, "radius": 1000.0}}}, + {"n": {"range": {"gte": 1, "lt": 9}}} + ] + }))) + .unwrap() + .unwrap(); + assert_eq!(field_conditions(&filter), (2, 0)); + + let Some(ConditionOneOf::Field(geo)) = &filter.must[0].condition_one_of else { + panic!("expected a geo field condition"); + }; + let radius = geo.geo_radius.as_ref().expect("geo_radius"); + assert_eq!(radius.radius, 1000.0); + assert_eq!(radius.center.as_ref().unwrap().lat, 52.5); + + let Some(ConditionOneOf::Field(range)) = &filter.must[1].condition_one_of else { + panic!("expected a range field condition"); + }; + let range = range.range.as_ref().expect("range"); + assert_eq!((range.gte, range.lt), (Some(1.0), Some(9.0))); + } + + /// One entry may name several fields, and a field several condition types; + /// each becomes its own condition, matching vector-db-benchmark. + #[test] + fn one_entry_can_carry_several_conditions() { + let filter = parse(Some(&json!({ + "and": [{"a": {"match": {"value": "x"}}, "b": {"range": {"gt": 1}}}] + }))) + .unwrap() + .unwrap(); + assert_eq!(field_conditions(&filter), (2, 0)); + } + + /// Silently dropping something unrecognized would score the run against + /// ground truth for a filter that was never applied — the whole bug this + /// module exists to fix. + #[test] + fn unknown_shapes_are_errors_not_silent_drops() { + for bad in [ + json!({"nand": [{"a": {"match": {"value": 1}}}]}), + json!({"and": [{"a": {"regex": {"value": "x"}}}]}), + json!({"and": [{"a": {"range": {}}}]}), + json!({"and": [{"a": {"match": {}}}]}), + json!({"and": [{"a": {"geo": {"lat": 1.0, "lon": 2.0}}}]}), + json!({"and": "not-an-array"}), + ] { + assert!(parse(Some(&bad)).is_err(), "should have failed: {bad}"); + } + } +} diff --git a/src/dataset/mod.rs b/src/dataset/mod.rs index 24f5b7e..d33f8a8 100644 --- a/src/dataset/mod.rs +++ b/src/dataset/mod.rs @@ -1,3 +1,4 @@ +mod conditions; mod config; mod download; #[cfg(test)] @@ -10,6 +11,7 @@ mod registry; mod sources; mod upload; +pub use conditions::parse as parse_query_conditions; pub use config::DatasetConfig; pub use download::{ensure_local_file, is_remote_url}; pub use reader::DatasetReader; diff --git a/src/dataset/reader.rs b/src/dataset/reader.rs index 7be615c..94a868d 100644 --- a/src/dataset/reader.rs +++ b/src/dataset/reader.rs @@ -162,6 +162,8 @@ impl DatasetReader { rows.push(QueryEntry { vector: self.query_dense_vector(idx)?, ground_truth: self.query_ground_truth(idx)?, + // Only the tar layout has a `conditions` field. + conditions: None, }); } Ok(rows) @@ -175,6 +177,7 @@ impl DatasetReader { rows.push(QueryEntry { vector: self.query_sparse_vector(idx)?, ground_truth: self.query_ground_truth(idx)?, + conditions: None, }); } Ok(rows) diff --git a/src/dataset/readers/query.rs b/src/dataset/readers/query.rs index a8172f6..8bf1cb6 100644 --- a/src/dataset/readers/query.rs +++ b/src/dataset/readers/query.rs @@ -8,6 +8,10 @@ pub struct QueryEntry { pub vector: V, pub ground_truth: Vec, + /// Raw `conditions` object the query was answered under, when the query set + /// carries one. Kept as JSON here and turned into a filter once at startup; + /// the ground truth only holds for a search that applies it. + pub conditions: Option, } /// A sparse query vector as `(index, value)` pairs — the `V` of a sparse diff --git a/src/dataset/readers/tar.rs b/src/dataset/readers/tar.rs index f6138f6..eb5bd2d 100644 --- a/src/dataset/readers/tar.rs +++ b/src/dataset/readers/tar.rs @@ -110,6 +110,7 @@ impl TarReader { .map(|row| QueryEntry { vector: row.query, ground_truth: row.closest_ids, + conditions: row.conditions, }) .collect()) } @@ -131,12 +132,16 @@ impl TarReader { } } -/// The fields of a `tests.jsonl` row that a benchmark run needs. `conditions` -/// and `closest_scores` are deliberately absent so serde skips them. +/// The fields of a `tests.jsonl` row that a benchmark run needs. +/// `closest_scores` is deliberately absent so serde skips it. #[derive(serde::Deserialize)] struct QueryRow { query: Vec, closest_ids: Vec, + /// The filter the ground truth was computed under. Absent in query sets + /// that predate it, `{}` in the `-no-filters` datasets. + #[serde(default)] + conditions: Option, } fn parse_f32_array(value: &Value) -> Option> { diff --git a/src/generators/queries.rs b/src/generators/queries.rs index 77dc2e7..6d9d875 100644 --- a/src/generators/queries.rs +++ b/src/generators/queries.rs @@ -59,6 +59,11 @@ struct QueryDataset { vectors: QueryVectors, /// Ground-truth nearest-neighbor ids per query, used to score recall. ground_truth: Vec>, + /// Filter each query was answered under, parallel to `vectors`. Its ground + /// truth only holds for a search that applies it, so this is not optional + /// decoration — dropping it scores the search against the answers to a + /// different question. + filters: Vec>, num_queries: usize, cursor: AtomicUsize, } @@ -70,6 +75,11 @@ enum QueryVectors { Sparse(Vec<(Vec, Vec)>), } +/// A sparse query drawn from a dataset query set: the `(values, indices)` +/// vector, the ground truth it is scored against, and the filter that ground +/// truth assumes. +type SparseDatasetQuery = ((Vec, SparseIndices), Option>, Option); + /// Which kind of query a request will draw from a dataset. #[derive(Clone, Copy)] enum QueryKind { @@ -371,6 +381,17 @@ impl ConfigSearchGenerator { /// Reading it all here is what keeps file I/O and JSON parsing out of the /// timed request path, so a query set that is missing, truncated, or of the /// wrong kind fails at startup rather than part-way through a benchmark. + /// Turn one query row's raw `conditions` into a filter, naming the row on + /// failure so a bad line in a 10k-query set is findable. + fn parse_conditions( + conditions: &Option, + dataset: &str, + idx: usize, + ) -> anyhow::Result> { + crate::dataset::parse_query_conditions(conditions.as_ref()) + .with_context(|| format!("dataset {dataset:?}, query {idx}")) + } + fn open_query_dataset( dataset: &crate::dataset::DatasetConfig, datasets_dir: &Path, @@ -385,6 +406,9 @@ impl ConfigSearchGenerator { ); } + // Conditions are parsed once here, so a malformed query set fails at + // startup and the timed path only ever clones a ready-made filter. + let mut filters: Vec> = Vec::with_capacity(num_queries); let (vectors, ground_truth) = match kind { QueryKind::Dense => { let rows = reader.read_dense_query_set().with_context(|| { @@ -395,9 +419,10 @@ impl ConfigSearchGenerator { })?; let mut vectors = Vec::with_capacity(rows.len()); let mut ground_truth = Vec::with_capacity(rows.len()); - for row in rows { + for (idx, row) in rows.into_iter().enumerate() { vectors.push(row.vector); ground_truth.push(row.ground_truth); + filters.push(Self::parse_conditions(&row.conditions, &dataset.name, idx)?); } (QueryVectors::Dense(vectors), ground_truth) } @@ -410,19 +435,30 @@ impl ConfigSearchGenerator { })?; let mut vectors = Vec::with_capacity(rows.len()); let mut ground_truth = Vec::with_capacity(rows.len()); - for row in rows { + for (idx, row) in rows.into_iter().enumerate() { // Pre-split into (values, indices) so no per-request unzip is needed. let (indices, values): (Vec, Vec) = row.vector.into_iter().unzip(); vectors.push((values, indices)); ground_truth.push(row.ground_truth); + filters.push(Self::parse_conditions(&row.conditions, &dataset.name, idx)?); } (QueryVectors::Sparse(vectors), ground_truth) } }; + let filtered = filters.iter().filter(|f| f.is_some()).count(); + if filtered > 0 { + println!( + "Dataset {:?}: {filtered} of {num_queries} queries carry filter \ + conditions; applying them (their ground truth assumes it).", + dataset.name + ); + } + Ok(QueryDataset { vectors, ground_truth, + filters, num_queries, cursor: AtomicUsize::new(0), }) @@ -457,23 +493,24 @@ impl ConfigSearchGenerator { source, filters: _, } => { - let (vector, expected_ids) = if let Some(query_dataset) = &state.query_dataset { - Self::read_dense_query(query_dataset) - } else { - let vector = Self::gen_dense_vector( - rng, - *size as usize, - *datatype, - source, - state.dense_reader.as_ref(), - req_id, - ); - (vector, None) - }; + let (vector, expected_ids, dataset_filter) = + if let Some(query_dataset) = &state.query_dataset { + Self::read_dense_query(query_dataset) + } else { + let vector = Self::gen_dense_vector( + rng, + *size as usize, + *datatype, + source, + state.dense_reader.as_ref(), + req_id, + ); + (vector, None, None) + }; GeneratedQuery { dense: Some((vector, using.clone())), sparse: None, - filter: state.filters.build(rng), + filter: dataset_filter.or_else(|| state.filters.build(rng)), idf_corpus: None, expected_ids, } @@ -484,19 +521,20 @@ impl ConfigSearchGenerator { filters: _, idf_corpus: _, } => { - let ((values, indices), expected_ids) = + let ((values, indices), expected_ids, dataset_filter) = if let Some(query_dataset) = &state.query_dataset { Self::read_sparse_query(query_dataset) } else { ( Self::gen_sparse_vector(rng, source, state.sparse_zipf.as_ref()), None, + None, ) }; GeneratedQuery { dense: None, sparse: Some((values, indices, using.clone())), - filter: state.filters.build(rng), + filter: dataset_filter.or_else(|| state.filters.build(rng)), idf_corpus: state.idf_corpus.build(rng), expected_ids, } @@ -508,7 +546,9 @@ impl ConfigSearchGenerator { /// /// The kind mismatch cannot happen: the request template that opened the /// dataset is the same one reading from it here. - fn read_dense_query(query_dataset: &QueryDataset) -> (Vec, Option>) { + fn read_dense_query( + query_dataset: &QueryDataset, + ) -> (Vec, Option>, Option) { let idx = query_dataset.next_index(); let QueryVectors::Dense(vectors) = &query_dataset.vectors else { panic!("dense request drew from a query set opened as sparse"); @@ -516,13 +556,12 @@ impl ConfigSearchGenerator { ( vectors[idx].clone(), Some(query_dataset.ground_truth[idx].clone()), + query_dataset.filters[idx].clone(), ) } /// Take the next sparse query vector and its ground-truth ids from a dataset. - fn read_sparse_query( - query_dataset: &QueryDataset, - ) -> ((Vec, SparseIndices), Option>) { + fn read_sparse_query(query_dataset: &QueryDataset) -> SparseDatasetQuery { let idx = query_dataset.next_index(); let QueryVectors::Sparse(vectors) = &query_dataset.vectors else { panic!("sparse request drew from a query set opened as dense"); @@ -536,6 +575,7 @@ impl ConfigSearchGenerator { }, ), Some(query_dataset.ground_truth[idx].clone()), + query_dataset.filters[idx].clone(), ) }