diff --git a/src/dataset/reader.rs b/src/dataset/reader.rs index 491c6e3..7be615c 100644 --- a/src/dataset/reader.rs +++ b/src/dataset/reader.rs @@ -6,7 +6,9 @@ use serde_json::Value; use super::config::{DatasetConfig, DatasetKind}; use super::download::ensure_downloaded; use super::parts::PartitionedReader; -use super::readers::{H5Reader, NpyReader, ParquetReader, SparseReader, TarReader}; +use super::readers::{ + H5Reader, NpyReader, ParquetReader, QueryEntry, SparseReader, SparseVector, TarReader, +}; use super::registry::load_registry; enum DatasetReaderInner { @@ -147,6 +149,37 @@ impl DatasetReader { } } + /// The whole dense query set with its ground truth, read in one pass. + /// + /// Text-backed formats get a sequential fast path; the rest fall back to + /// indexed reads, which for a binary format cost the same either way. + pub fn read_dense_query_set(&self) -> Result>>> { + if let DatasetReaderInner::Tar(reader) = &self.inner { + return reader.read_query_set(); + } + let mut rows = Vec::with_capacity(self.num_queries()); + for idx in 0..self.num_queries() { + rows.push(QueryEntry { + vector: self.query_dense_vector(idx)?, + ground_truth: self.query_ground_truth(idx)?, + }); + } + Ok(rows) + } + + /// The whole sparse query set with its ground truth. Sparse queries live in + /// binary CSR files, so there is nothing to gain from a bulk path. + pub fn read_sparse_query_set(&self) -> Result>> { + let mut rows = Vec::with_capacity(self.num_queries()); + for idx in 0..self.num_queries() { + rows.push(QueryEntry { + vector: self.query_sparse_vector(idx)?, + ground_truth: self.query_ground_truth(idx)?, + }); + } + Ok(rows) + } + /// Ground-truth nearest-neighbor ids for a query (indices into the corpus). pub fn query_ground_truth(&self, idx: usize) -> Result> { match &self.inner { diff --git a/src/dataset/readers/jsonl.rs b/src/dataset/readers/jsonl.rs index 1ce66de..f3cdbe6 100644 --- a/src/dataset/readers/jsonl.rs +++ b/src/dataset/readers/jsonl.rs @@ -3,6 +3,7 @@ use std::io::{BufRead, BufReader, Seek, SeekFrom}; use std::path::{Path, PathBuf}; use anyhow::{Context, Result}; +use serde::de::DeserializeOwned; use serde_json::Value; /// Line-oriented JSON store with byte-offset index for random access. @@ -40,6 +41,40 @@ impl JsonlStore { self.offsets.len() } + /// Deserialize every line, in order, into `T`. + /// + /// The bulk counterpart to [`Self::value_at`]: one open and one sequential + /// pass over the file, versus an open plus a seek per row. `T` should name + /// only the fields it needs, so serde can write them directly instead of + /// materializing a [`Value`] tree — for a row holding a 2048-element vector + /// that tree is 2048 separately boxed numbers. + pub fn deserialize_all(&self) -> Result> { + let file = File::open(&self.path) + .with_context(|| format!("failed to open {}", self.path.display()))?; + // Rows can be tens of KB; a large buffer keeps one row to one refill. + let mut reader = BufReader::with_capacity(1 << 20, file); + let mut rows = Vec::with_capacity(self.offsets.len()); + let mut line = String::new(); + loop { + line.clear(); + if reader.read_line(&mut line)? == 0 { + break; + } + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + rows.push(serde_json::from_str(trimmed).with_context(|| { + format!( + "failed to parse {} line {}", + self.path.display(), + rows.len() + ) + })?); + } + Ok(rows) + } + pub fn value_at(&self, idx: usize) -> Result> { let Some(&byte_offset) = self.offsets.get(idx) else { return Ok(None); diff --git a/src/dataset/readers/mod.rs b/src/dataset/readers/mod.rs index 0f29139..383c7d6 100644 --- a/src/dataset/readers/mod.rs +++ b/src/dataset/readers/mod.rs @@ -3,6 +3,7 @@ mod h5; mod jsonl; mod npy; mod parquet; +mod query; mod sparse; mod tar; @@ -11,5 +12,6 @@ pub use npy::{NpyReader, parse_npy_header}; pub use parquet::{ ParquetReader, parquet_footer_len, parquet_row_count, parquet_row_count_from_tail, }; +pub use query::{QueryEntry, SparseVector}; pub use sparse::SparseReader; pub use tar::TarReader; diff --git a/src/dataset/readers/query.rs b/src/dataset/readers/query.rs new file mode 100644 index 0000000..a8172f6 --- /dev/null +++ b/src/dataset/readers/query.rs @@ -0,0 +1,15 @@ +/// One entry of a dataset's query set: a query vector paired with the +/// ground-truth ids used to score its recall. +/// +/// Query sets are read in full at startup and returned as `Vec>`, +/// which keeps file I/O and JSON parsing off the timed request path. `V` is the +/// vector form the caller asked for — `Vec` for dense, `Vec<(u32, f32)>` +/// for sparse. +pub struct QueryEntry { + pub vector: V, + pub ground_truth: Vec, +} + +/// A sparse query vector as `(index, value)` pairs — the `V` of a sparse +/// [`QueryEntry`]. Named so the query-set return type stays legible. +pub type SparseVector = Vec<(u32, f32)>; diff --git a/src/dataset/readers/tar.rs b/src/dataset/readers/tar.rs index a7950e3..f6138f6 100644 --- a/src/dataset/readers/tar.rs +++ b/src/dataset/readers/tar.rs @@ -5,6 +5,7 @@ use serde_json::Value; use super::jsonl::JsonlStore; use super::npy::NpyMatrix; +use super::query::QueryEntry; /// An extracted `.tgz` bundle: `vectors.npy` plus optional `payloads.jsonl` and /// `tests.jsonl` (ann-filtering-benchmark-datasets layout). @@ -92,6 +93,27 @@ impl TarReader { parse_f32_array(query).context("tests.jsonl `query` is not an array of numbers") } + /// The whole query set: every query vector with its ground-truth ids. + /// + /// Reads each row exactly once. Going through [`Self::query_at`] and + /// [`Self::query_ground_truth`] per index would reopen the file and re-parse + /// the same row twice — once per field — which dominated startup on a query + /// set of 2048-d vectors. + pub fn read_query_set(&self) -> Result>>> { + let store = self + .queries + .as_ref() + .context("dataset has no tests.jsonl (no query set)")?; + let rows: Vec = store.deserialize_all()?; + Ok(rows + .into_iter() + .map(|row| QueryEntry { + vector: row.query, + ground_truth: row.closest_ids, + }) + .collect()) + } + /// Ground-truth nearest-neighbor ids for a query (`closest_ids` field). pub fn query_ground_truth(&self, idx: usize) -> Result> { let line = self.query_line(idx)?; @@ -109,6 +131,14 @@ 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. +#[derive(serde::Deserialize)] +struct QueryRow { + query: Vec, + closest_ids: Vec, +} + fn parse_f32_array(value: &Value) -> Option> { value .as_array()? @@ -152,6 +182,14 @@ mod tests { assert_eq!(reader.query_at(0).unwrap(), vec![1.0, 2.0, 3.0]); assert_eq!(reader.query_ground_truth(0).unwrap(), vec![1, 0]); assert_eq!(reader.query_ground_truth(1).unwrap(), vec![0]); + + // The bulk pass must agree with the indexed reads it replaces. + let entries = reader.read_query_set().unwrap(); + assert_eq!(entries.len(), 2); + assert_eq!(entries[0].vector, vec![1.0, 2.0, 3.0]); + assert_eq!(entries[0].ground_truth, vec![1, 0]); + assert_eq!(entries[1].vector, vec![4.0, 5.0, 6.0]); + assert_eq!(entries[1].ground_truth, vec![0]); } #[test] diff --git a/src/generators/queries.rs b/src/generators/queries.rs index 7ebd62a..77dc2e7 100644 --- a/src/generators/queries.rs +++ b/src/generators/queries.rs @@ -47,14 +47,36 @@ pub struct GeneratedQuery { pub expected_ids: Option>, } -/// A reference dataset used as a query source, together with a cursor that -/// hands out consecutive query indices so every query in the set is exercised. +/// A reference dataset's query set, held in memory, with a cursor that hands out +/// consecutive query indices so every query in the set is exercised. +/// +/// The whole set is read once when the dataset is opened rather than a row at a +/// time during the benchmark. Resolving a row on demand meant reopening the +/// file and parsing the JSON twice per search — for a 2048-d query that is +/// ~45 KB of text each way, enough client-side work to outweigh the search being +/// timed. Parsed form is far smaller than the file: a 10k x 2048-d set is ~82 MB. struct QueryDataset { - reader: DatasetReader, + vectors: QueryVectors, + /// Ground-truth nearest-neighbor ids per query, used to score recall. + ground_truth: Vec>, num_queries: usize, cursor: AtomicUsize, } +/// Query vectors in whichever form the request that opened the dataset needs. +enum QueryVectors { + Dense(Vec>), + /// Kept pre-split as (values, indices) so no per-request unzip is needed. + Sparse(Vec<(Vec, Vec)>), +} + +/// Which kind of query a request will draw from a dataset. +#[derive(Clone, Copy)] +enum QueryKind { + Dense, + Sparse, +} + impl QueryDataset { /// Next query index to use, wrapping around the query set. fn next_index(&self) -> usize { @@ -280,9 +302,14 @@ impl ConfigSearchGenerator { let local = ensure_local_file(datasets_dir, path)?; (Some(FBinReader::new(&local)?), None) } - VectorSource::Dataset { dataset } => { - (None, Some(Self::open_query_dataset(dataset, datasets_dir)?)) - } + VectorSource::Dataset { dataset } => ( + None, + Some(Self::open_query_dataset( + dataset, + datasets_dir, + QueryKind::Dense, + )?), + ), VectorSource::Random => (None, None), }; (dense_reader, query_dataset, None, filters, [].as_slice()) @@ -300,7 +327,11 @@ impl ConfigSearchGenerator { .context("sparse dataset query source is missing dataset fields")?; ( None, - Some(Self::open_query_dataset(dataset, datasets_dir)?), + Some(Self::open_query_dataset( + dataset, + datasets_dir, + QueryKind::Sparse, + )?), None, filters, idf_corpus.as_slice(), @@ -333,12 +364,17 @@ impl ConfigSearchGenerator { }) } - /// Open a reference dataset as a query source, requiring it to ship a query - /// set (ann-benchmarks `test`/`neighbors`, or `tests.jsonl` / - /// `queries.csr`+`results.gt`). + /// 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`). + /// + /// 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. fn open_query_dataset( dataset: &crate::dataset::DatasetConfig, datasets_dir: &Path, + kind: QueryKind, ) -> anyhow::Result { let reader = DatasetReader::open(datasets_dir, dataset)?; let num_queries = reader.num_queries(); @@ -348,8 +384,45 @@ impl ConfigSearchGenerator { dataset.name ); } + + let (vectors, ground_truth) = match kind { + QueryKind::Dense => { + let rows = reader.read_dense_query_set().with_context(|| { + format!( + "failed to read dense query set of dataset {:?}", + dataset.name + ) + })?; + let mut vectors = Vec::with_capacity(rows.len()); + let mut ground_truth = Vec::with_capacity(rows.len()); + for row in rows { + vectors.push(row.vector); + ground_truth.push(row.ground_truth); + } + (QueryVectors::Dense(vectors), ground_truth) + } + QueryKind::Sparse => { + let rows = reader.read_sparse_query_set().with_context(|| { + format!( + "failed to read sparse query set of dataset {:?}", + dataset.name + ) + })?; + let mut vectors = Vec::with_capacity(rows.len()); + let mut ground_truth = Vec::with_capacity(rows.len()); + for row in rows { + // 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); + } + (QueryVectors::Sparse(vectors), ground_truth) + } + }; + Ok(QueryDataset { - reader, + vectors, + ground_truth, num_queries, cursor: AtomicUsize::new(0), }) @@ -431,35 +504,39 @@ impl ConfigSearchGenerator { } } - /// Read the next dense query vector and its ground-truth ids from a dataset. + /// Take the next dense query vector and its ground-truth ids from a dataset. + /// + /// 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>) { let idx = query_dataset.next_index(); - let vector = query_dataset - .reader - .query_dense_vector(idx) - .unwrap_or_else(|e| panic!("failed to read dataset query vector at {idx}: {e}")); - let expected = query_dataset - .reader - .query_ground_truth(idx) - .unwrap_or_else(|e| panic!("failed to read dataset ground truth at {idx}: {e}")); - (vector, Some(expected)) + let QueryVectors::Dense(vectors) = &query_dataset.vectors else { + panic!("dense request drew from a query set opened as sparse"); + }; + ( + vectors[idx].clone(), + Some(query_dataset.ground_truth[idx].clone()), + ) } - /// Read the next sparse query vector and its ground-truth ids from a dataset. + /// Take the next sparse query vector and its ground-truth ids from a dataset. fn read_sparse_query( query_dataset: &QueryDataset, ) -> ((Vec, SparseIndices), Option>) { let idx = query_dataset.next_index(); - let pairs = query_dataset - .reader - .query_sparse_vector(idx) - .unwrap_or_else(|e| panic!("failed to read dataset sparse query at {idx}: {e}")); - let (indices, values): (Vec<_>, Vec<_>) = pairs.into_iter().unzip(); - let expected = query_dataset - .reader - .query_ground_truth(idx) - .unwrap_or_else(|e| panic!("failed to read dataset ground truth at {idx}: {e}")); - ((values, SparseIndices { data: indices }), Some(expected)) + let QueryVectors::Sparse(vectors) = &query_dataset.vectors else { + panic!("sparse request drew from a query set opened as dense"); + }; + let (values, indices) = &vectors[idx]; + ( + ( + values.clone(), + SparseIndices { + data: indices.clone(), + }, + ), + Some(query_dataset.ground_truth[idx].clone()), + ) } fn gen_dense_vector(