Fix jsonl read bottleneck

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<Vec<QueryEntry<Vec<f32>>>> {
+        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<Vec<QueryEntry<SparseVector>>> {
+        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<Vec<u64>> {
         match &self.inner {

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<T: DeserializeOwned>(&self) -> Result<Vec<T>> {
+        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<Option<Value>> {
         let Some(&byte_offset) = self.offsets.get(idx) else {
             return Ok(None);

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;

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<QueryEntry<_>>`,
+/// which keeps file I/O and JSON parsing off the timed request path. `V` is the
+/// vector form the caller asked for — `Vec<f32>` for dense, `Vec<(u32, f32)>`
+/// for sparse.
+pub struct QueryEntry<V> {
+    pub vector: V,
+    pub ground_truth: Vec<u64>,
+}
+
+/// 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)>;

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<Vec<QueryEntry<Vec<f32>>>> {
+        let store = self
+            .queries
+            .as_ref()
+            .context("dataset has no tests.jsonl (no query set)")?;
+        let rows: Vec<QueryRow> = 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<Vec<u64>> {
         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<f32>,
+    closest_ids: Vec<u64>,
+}
+
 fn parse_f32_array(value: &Value) -> Option<Vec<f32>> {
     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]

src/generators/queries.rs

@@ -47,14 +47,36 @@ pub struct GeneratedQuery {
     pub expected_ids: Option<Vec<u64>>,
 }
 
-/// 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<Vec<u64>>,
     num_queries: usize,
     cursor: AtomicUsize,
 }
 
+/// Query vectors in whichever form the request that opened the dataset needs.
+enum QueryVectors {
+    Dense(Vec<Vec<f32>>),
+    /// Kept pre-split as (values, indices) so no per-request unzip is needed.
+    Sparse(Vec<(Vec<f32>, Vec<u32>)>),
+}
+
+/// 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<QueryDataset> {
         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<u32>, Vec<f32>) = 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<f32>, Option<Vec<u64>>) {
         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<f32>, SparseIndices), Option<Vec<u64>>) {
         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(