Repository: chroma-core/chroma
Pull request #4341: [ENH] Add VerifyingDataSet that verifies loaded data
Author: login: jasonvigil
Created: 2025-04-21T23:53:07Z | Merged: 2025-04-23T02:29:06Z

## Description of changes

Upsert sequentially from the reference data set. Verify data by getting from the subset of already-loaded keys and checking results against the reference data set.

## Test plan

1. Create reference dataset using loader script:

```
export CHROMA_TENANT_ID="abc"
export CHROMA_DATABASE="chroma-load"
export CHROMA_API_KEY="xyz"

pipenv run -- python script/import_dataset.py \
--tenant $CHROMA_TENANT_ID \
--database $CHROMA_DATABASE \
--collection reference-all-MiniLM-L6-v2 \
--dataset roneneldan/TinyStories \
--dataset-split train \
--sentence-transformer-model all-MiniLM-L6-v2 \
--max-rows 1000000 \
--max-rps 120 \
--sequential-keys
```

2. (manually) Create (or reset via delete+recreate) the test collection: `test-all-MiniLM-L6-v2`
 
3.  Start the chroma-load service:

```
export CHROMA_HOST="https://staging.chroma.api:8000"
export CHROMA_DATABASE="chroma-load"
export CHROMA_TOKEN="xyz"

RUST_LOG=debug cargo run --bin chroma-load
```

4. Kick-off a verification workload, with 10 operations / second.

```
cargo run --bin chroma-load-start -- \
--host http://localhost:3001 \
--name my-test \
--expires "2min" \
--data-set "test-all-MiniLM-L6-v2" \
--workload verify \
--constant-throughput "10"
```

5. Ensure no failures are recorded by the workload.


Diff:
diff --git a/rust/load/src/data_sets.rs b/rust/load/src/data_sets.rs
index e9726d51f48..8822b293c2e 100644
--- a/rust/load/src/data_sets.rs
+++ b/rust/load/src/data_sets.rs
@@ -1,3 +1,5 @@
+use std::cmp::Reverse;
+use std::collections::BinaryHeap;
 use std::sync::atomic::AtomicUsize;
 use std::sync::Arc;
 
@@ -5,7 +7,9 @@ use chromadb::collection::{CollectionEntries, GetOptions, GetResult, QueryOption
 use chromadb::ChromaClient;
 use guacamole::combinators::*;
 use guacamole::Guacamole;
+use tokio::sync::Mutex;
 use tracing::Instrument;
+use tracing::{event, Level};
 
 use crate::{bit_difference, DataSet, Error, GetQuery, KeySelector, QueryQuery, UpsertQuery};
 
@@ -743,6 +747,10 @@ impl DataSet for ReferencingDataSet {
         self.cardinality
     }
 
+    fn reference_cardinality(&self) -> usize {
+        self.references.cardinality()
+    }
+
     async fn get(
         &self,
         client: &ChromaClient,
@@ -822,12 +830,12 @@ impl DataSet for ReferencingDataSet {
         &self,
         client: &ChromaClient,
         uq: UpsertQuery,
-        guac: &mut Guacamole,
+        _: &mut Guacamole,
     ) -> Result<(), Box<dyn std::error::Error + Send>> {
         let collection = client.get_collection(&self.operates_on).await?;
         let mut keys = vec![];
         for offset in 0..uq.batch_size {
-            keys.push(uq.key.select_with_offset(guac, self, offset));
+            keys.push(uq.key.select_from_reference(self, offset));
         }
         let keys = keys.iter().map(|k| k.as_str()).collect::<Vec<_>>();
         if let Some(res) = self.references.get_by_key(client, &keys).await? {
@@ -880,6 +888,349 @@ impl DataSet for ReferencingDataSet {
     }
 }
 
+//////////////////////////////////////// VerifyingDataSet /////////////////////////////////////////
+
+/// A verifying data set uses a reference data set to verify results from a test data set.
+#[derive(Debug)]
+pub struct VerifyingDataSet {
+    reference_data_set: Arc<dyn DataSet>,
+    test_data_set: String,
+    cardinality: AtomicUsize,
+    cardinality_heap: Arc<Mutex<BinaryHeap<Reverse<usize>>>>,
+}
+
+impl VerifyingDataSet {
+    /// Record a load of a chunk of data.
+    ///
+    /// This function updates the cardinality of the data set. It uses a min-heap to track the
+    /// positions for contiguous loaded chunks. For example, if the data set has a cardinality of
+    /// 1000, and a chunk is loaded from 1100 to 1199, then the heap will contain the value 1100.
+    /// The cardinality of the data set will not be updated in this case, because the chunk is not
+    /// contiguous with the data set cardinality. If the next chunk is loaded from 1000 to 1099,
+    /// then the heap will contain the values 1000 and 1100. After looping and reading these
+    /// contiguous chunks, the data set be updated with a cardinality of 1200, and the heap will
+    /// be emptied out.
+    async fn record_load(&self, loaded_chunk_start_idx: usize, batch_size: usize) {
+        let mut heap = self.cardinality_heap.lock().await;
+
+        let mut cardinality = self.cardinality.load(std::sync::atomic::Ordering::Relaxed);
+
+        heap.push(Reverse(loaded_chunk_start_idx));
+
+        let mut min_loaded_so_far = heap.peek().unwrap().0;
+
+        while min_loaded_so_far == cardinality {
+            cardinality = heap.pop().unwrap().0 + batch_size;
+
+            match heap.peek() {
+                Some(Reverse(next_min)) => min_loaded_so_far = *next_min,
+                None => break,
+            }
+        }
+
+        self.cardinality
+            .store(cardinality, std::sync::atomic::Ordering::Relaxed);
+
+        event!(
+            Level::INFO,
+            "Loaded {} keys, loaded chunk start idx: {}, new cardinality: {}",
+            batch_size,
+            loaded_chunk_start_idx,
+            cardinality
+        );
+    }
+}
+
+#[async_trait::async_trait]
+impl DataSet for VerifyingDataSet {
+    fn name(&self) -> String {
+        self.test_data_set.clone()
+    }
+
+    fn description(&self) -> String {
+        format!(
+            "referencing data set {}, testing on {}",
+            self.reference_data_set.name(),
+            self.test_data_set
+        )
+    }
+
+    fn json(&self) -> serde_json::Value {
+        serde_json::json! {
+            {
+                "reference_data_set": self.reference_data_set.json(),
+                "test_data_set": self.test_data_set,
+            }
+        }
+    }
+
+    fn cardinality(&self) -> usize {
+        self.cardinality.load(std::sync::atomic::Ordering::Relaxed)
+    }
+
+    fn reference_cardinality(&self) -> usize {
+        self.reference_data_set.cardinality()
+    }
+
+    async fn get(
+        &self,
+        client: &ChromaClient,
+        gq: GetQuery,
+        guac: &mut Guacamole,
+    ) -> Result<(), Box<dyn std::error::Error + Send>> {
+        let mut keys = vec![];
+        let num_keys = gq.limit.sample(guac);
+
+        if self.cardinality() < num_keys {
+            // If the cardinality is less than the number of keys requested, return Ok(())
+            // because the data set is not populated yet.
+            return Ok(());
+        }
+
+        for _ in 0..num_keys {
+            keys.push(KeySelector::Random(gq.skew).select(guac, self));
+        }
+
+        let reference_collection = client
+            .get_collection(&self.reference_data_set.name())
+            .await?;
+        let reference_results = reference_collection
+            .get(GetOptions {
+                ids: keys.clone(),
+                where_metadata: None,
+                limit: None,
+                offset: None,
+                where_document: None,
+                include: Some(vec![
+                    "embeddings".to_string(),
+                    "metadatas".to_string(),
+                    "documents".to_string(),
+                ]),
+            })
+            .await?;
+
+        let test_collection = client.get_collection(&self.test_data_set).await?;
+        let test_results = test_collection
+            .get(GetOptions {
+                ids: keys.clone(),
+                where_metadata: None,
+                limit: None,
+                offset: None,
+                where_document: None,
+                include: Some(vec![
+                    "embeddings".to_string(),
+                    "metadatas".to_string(),
+                    "documents".to_string(),
+                ]),
+            })
+            .await?;
+
+        event!(
+            Level::INFO,
+            "Querying for keys: {}, cardinality: {}, reference results: {}, test results: {}",
+            keys.iter()
+                .map(|k| k.as_str())
+                .collect::<Vec<_>>()
+                .join(", "),
+            self.cardinality(),
+            reference_results.ids.len(),
+            test_results.ids.len()
+        );
+
+        if test_results.ids.len() != reference_results.ids.len() {
+            return Err(Box::new(Error::InvalidRequest(format!(
+                "IDs lengths are different: got {}, expected {}",
+                test_results.ids.len(),
+                reference_results.ids.len()
+            ))));
+        }
+
+        let reference_embeddings = match reference_results.embeddings {
+            Some(embeddings) => embeddings,
+            None => {
+                return Err(Box::new(Error::InvalidRequest(
+                    "Reference results have no embeddings".into(),
+                )))
+            }
+        };
+        let test_embeddings = match test_results.embeddings {
+            Some(embeddings) => embeddings,
+            None => {
+                return Err(Box::new(Error::InvalidRequest(
+                    "Test results have no embeddings".into(),
+                )))
+            }
+        };
+        if reference_embeddings.len() != test_embeddings.len() {
+            return Err(Box::new(Error::InvalidRequest(format!(
+                "Embeddings lengths are different: got {}, expected {}",
+                test_embeddings.len(),
+                reference_embeddings.len()
+            ))));
+        }
+
+        let reference_metadatas = match reference_results.metadatas {
+            Some(metadatas) => metadatas,
+            None => {
+                return Err(Box::new(Error::InvalidRequest(
+                    "Reference results have no metadatas".into(),
+                )))
+            }
+        };
+        let test_metadatas = match test_results.metadatas {
+            Some(metadatas) => metadatas,
+            None => {
+                return Err(Box::new(Error::InvalidRequest(
+                    "Test results have no metadatas".into(),
+                )))
+            }
+        };
+        if reference_metadatas.len() != test_metadatas.len() {
+            return Err(Box::new(Error::InvalidRequest(format!(
+                "Metadatas lengths are different: got {}, expected {}",
+                test_metadatas.len(),
+                reference_metadatas.len()
+            ))));
+        }
+
+        for (reference_idx, reference_id) in reference_results.ids.iter().enumerate() {
+            let test_idx = test_results
+                .ids
+                .iter()
+                .position(|id| id == reference_id)
+                .unwrap();
+
+            let reference_embedding = match &reference_embeddings[reference_idx] {
+                Some(embedding) => embedding,
+                None => {
+                    return Err(Box::new(Error::InvalidRequest(format!(
+                        "Reference embedding for id {} is missing",
+                        reference_id
+                    ))));
+                }
+            };
+            let test_embedding = match &test_embeddings[test_idx] {
+                Some(embedding) => embedding,
+                None => {
+                    return Err(Box::new(Error::InvalidRequest(format!(
+                        "Test embedding for id {} is missing",
+                        reference_id
+                    ))));
+                }
+            };
+
+            if reference_embedding != test_embedding {
+                return Err(Box::new(Error::InvalidRequest(format!(
+                    "Embeddings for id {} are different: got {:?}, expected {:?}",
+                    reference_id, test_embedding, reference_embedding
+                ))));
+            }
+
+            let reference_metadata = &reference_metadatas[reference_idx];
+            let test_metadata = &test_metadatas[test_idx];
+
+            if reference_metadata != test_metadata {
+                return Err(Box::new(Error::InvalidRequest(format!(
+                    "Metadata for id {} are different: got {:?}, expected {:?}",
+                    reference_id, test_metadata, reference_metadata
+                ))));
+            }
+        }
+
+        Ok(())
+    }
+
+    async fn query(
+        &self,
+        _: &ChromaClient,
+        _: QueryQuery,
+        _: &mut Guacamole,
+    ) -> Result<(), Box<dyn std::error::Error + Send>> {
+        Err(Box::new(Error::InvalidRequest(
+            "Query not implemented (yet)".into(),
+        )))
+    }
+
+    async fn upsert(
+        &self,
+        client: &ChromaClient,
+        uq: UpsertQuery,
+        _: &mut Guacamole,
+    ) -> Result<(), Box<dyn std::error::Error + Send>> {
+        let collection = client.get_collection(&self.test_data_set).await?;
+        let mut keys = vec![];
+
+        let key_start_index = match &uq.key {
+            KeySelector::Index(key_start_index) => *key_start_index,
+            _ => {
+                return Err(Box::new(Error::InvalidRequest(
+                    "Key selector must be an index".into(),
+                )))
+            }
+        };
+
+        event!(
+            Level::INFO,
+            "Upserting {} keys, new cardinality: {}",
+            uq.batch_size,
+            key_start_index + uq.batch_size
+        );
+
+        for offset in 0..uq.batch_size {
+            keys.push(uq.key.select_from_reference(self, offset));
+        }
+        let keys = keys.iter().map(|k| k.as_str()).collect::<Vec<_>>();
+        if let Some(res) = self.reference_data_set.get_by_key(client, &keys).await? {
+            let mut keys = vec![];
+            for id in res.ids.iter() {
+                keys.push(id.as_str());
+            }
+            let mut documents = vec![];
+            if let Some(docs) = res.documents {
+                for (idx, doc) in docs.into_iter().enumerate() {
+                    if let Some(doc) = doc {
+                        documents.push(doc);
+                    } else {
+                        return Err(Box::new(Error::InvalidRequest(format!(
+                            "Missing document for {}",
+                            idx
+                        ))));
+                    }
+                }
+            } else {
+                return Err(Box::new(Error::InvalidRequest("No documents".into())));
+            }
+            let documents = documents.iter().map(|d| d.as_str()).collect::<Vec<_>>();
+            let mut embeddings = vec![];
+            if let Some(embeds) = res.embeddings {
+                for (idx, embed) in embeds.iter().enumerate() {
+                    if let Some(embed) = embed {
+                        embeddings.push(embed.clone());
+                    } else {
+                        return Err(Box::new(Error::InvalidRequest(format!(
+                            "Missing document for {}",
+                            idx
+                        ))));
+                    }
+                }
+            } else {
+                return Err(Box::new(Error::InvalidRequest("No documents".into())));
+            }
+            let entries = CollectionEntries {
+                ids: keys,
+                metadatas: res.metadatas,
+                documents: Some(documents),
+                embeddings: Some(embeddings),
+            };
+            collection.upsert(entries, None).await?;
+            self.record_load(key_start_index, uq.batch_size).await;
+        } else {
+            return Err(Box::new(Error::InvalidRequest("No results".into())));
+        }
+        Ok(())
+    }
+}
+
 //////////////////////////////////////////// RoundRobin ////////////////////////////////////////////
 
 /// A data set that round-robins between other data sets.
@@ -1046,6 +1397,19 @@ pub fn all_data_sets() -> Vec<Arc<dyn DataSet>> {
             }
         }
     }
+
+    let reference_data_set = Arc::new(TinyStoriesDataSet::new(TinyStoriesDataSetType::reference(
+        "reference",
+        ALL_MINILM_L6_V2,
+        1_000_000,
+    )));
+    data_sets.push(Arc::new(VerifyingDataSet {
+        reference_data_set,
+        test_data_set: "test-all-MiniLM-L6-v2".to_string(),
+        cardinality: AtomicUsize::new(0),
+        cardinality_heap: Arc::new(Mutex::new(BinaryHeap::new())),
+    }) as _);
+
     data_sets
 }
 
diff --git a/rust/load/src/lib.rs b/rust/load/src/lib.rs
index c5a6b04e001..f09fe494f7b 100644
--- a/rust/load/src/lib.rs
+++ b/rust/load/src/lib.rs
@@ -32,6 +32,7 @@ use guacamole::{Guacamole, Zipf};
 use opentelemetry::global;
 use opentelemetry::metrics::Counter;
 use opentelemetry::{Key, KeyValue, Value};
+use tokio::sync::Mutex as TokioMutex;
 use tower_http::trace::TraceLayer;
 use tracing::Instrument;
 use uuid::Uuid;
@@ -187,13 +188,14 @@ pub async fn client() -> ChromaClient {
 /// variable if set, or no authentication if unset.
 pub async fn client_for_url(url: String) -> ChromaClient {
     if let Ok(auth) = std::env::var("CHROMA_TOKEN") {
+        let db = std::env::var("CHROMA_DATABASE").unwrap_or_else(|_| "hf-tiny-stories".into());
         ChromaClient::new(ChromaClientOptions {
             url: Some(url),
             auth: ChromaAuthMethod::TokenAuth {
                 token: auth,
                 header: ChromaTokenHeader::XChromaToken,
             },
-            database: "hf-tiny-stories".to_string(),
+            database: db,
         })
         .await
         .unwrap()
@@ -230,6 +232,11 @@ pub trait DataSet: std::fmt::Debug + Send + Sync {
     /// The number of documents in the data set.
     fn cardinality(&self) -> usize;
 
+    /// The number of documents in the reference data set.
+    fn reference_cardinality(&self) -> usize {
+        self.cardinality()
+    }
+
     /// Get documents by key.  This is used when one workload references another.  Return None to
     /// indicate the data set does not support referencing by index.
     async fn get_by_key(
@@ -521,18 +528,14 @@ impl KeySelector {
         format!("{:0>16}", index)
     }
 
-    /// Select a key from the distribution.
-    pub fn select_with_offset(
-        &self,
-        guac: &mut Guacamole,
-        data_set: &dyn DataSet,
-        offset: usize,
-    ) -> String {
+    /// Select a key from the reference data set distribution.
+    /// Only supports index selection.
+    pub fn select_from_reference(&self, data_set: &dyn DataSet, offset: usize) -> String {
         let index = match self {
             KeySelector::Index(i) => *i,
-            KeySelector::Random(skew) => skew.sample(guac, data_set.cardinality()),
+            _ => panic!("Only index selection is supported for reference data sets"),
         };
-        let index = (index + offset) % data_set.cardinality();
+        let index = (index + offset) % data_set.reference_cardinality();
         format!("{:0>16}", index)
     }
 }
@@ -559,7 +562,7 @@ pub struct UpsertQuery {
 /// The state of a workload.
 #[derive(Clone)]
 pub struct WorkloadState {
-    seq_no: u64,
+    seq_no: Arc<TokioMutex<u64>>,
     guac: Guacamole,
 }
 
@@ -718,7 +721,16 @@ impl Workload {
                 Box::pin(wrap.step(client, metrics, data_set, state, done)).await
             }
             Workload::Load => {
-                if (state.seq_no * 100) >= data_set.cardinality() as u64 {
+                let load_start_idx: u64;
+                {
+                    // Get the seq_no, increment it, and immediately drop the lock (to allow for
+                    // concurrent loads)
+                    let mut seq_no = state.seq_no.lock().await;
+                    load_start_idx = *seq_no * 100;
+                    *seq_no += 1;
+                }
+
+                if load_start_idx >= data_set.reference_cardinality() as u64 {
                     done.store(true, std::sync::atomic::Ordering::Relaxed);
                     return Ok(());
                 }
@@ -733,7 +745,7 @@ impl Workload {
                     .upsert(
                         client,
                         UpsertQuery {
-                            key: KeySelector::Index((state.seq_no * 100) as usize),
+                            key: KeySelector::Index(load_start_idx as usize),
                             batch_size: 100,
                             // Associativity is the ratio of documents in a cluster to documents
                             // written by the workload.  It is ignored for load.
@@ -1290,10 +1302,9 @@ impl LoadService {
                 }
             }
         });
-        let mut seq_no = 0u64;
+        let seq_no = Arc::new(TokioMutex::new(0u64));
         let start = Instant::now();
         while !done.load(std::sync::atomic::Ordering::Relaxed) {
-            seq_no += 1;
             let throughput = match spec.throughput {
                 Throughput::Constant(throughput) => throughput,
                 Throughput::Sinusoidal {
@@ -1347,7 +1358,10 @@ impl LoadService {
                 let client = Arc::clone(&client);
                 let data_set = Arc::clone(&spec.data_set);
                 let guac = Guacamole::new(any(&mut guac));
-                let mut state = WorkloadState { seq_no, guac };
+                let mut state = WorkloadState {
+                    seq_no: Arc::clone(&seq_no),
+                    guac,
+                };
                 let done = Arc::clone(&done);
                 let fut = async move {
                     this.metrics.num_operations.add(1, &[]);
diff --git a/rust/load/src/workloads.rs b/rust/load/src/workloads.rs
index 06aca191c8f..69e811e92c1 100644
--- a/rust/load/src/workloads.rs
+++ b/rust/load/src/workloads.rs
@@ -109,5 +109,20 @@ pub fn all_workloads() -> HashMap<String, Workload> {
             "random-upsert".to_string(),
             Workload::RandomUpsert(KeySelector::Random(Skew::Zipf { theta: 0.999 })),
         ),
+        (
+            "verify".to_string(),
+            Workload::Hybrid(vec![
+                (
+                    0.9,
+                    Workload::Get(GetQuery {
+                        skew: Skew::Uniform,
+                        limit: Distribution::Constant(10),
+                        metadata: None,
+                        document: None,
+                    }),
+                ),
+                (0.1, Workload::Load),
+            ]),
+        ),
     ])
 }


Reviews:
user: login: rescrv
body: High level feedback:  We could simplify things a bit by assuming load has been run and then making verify compare.  You could even extend the existing reference data set code to do so (so, compare the returned result against the value used to generate the query).
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T15:09:41Z
user: login: jasonvigil
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T22:37:41Z
user: login: rescrv
body: Looks like a good start.  I left some comments, mostly to help you learn rust.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T23:27:46Z
user: login: jasonvigil
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T00:15:15Z
user: login: jasonvigil
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T00:17:03Z
user: login: jasonvigil
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T00:20:40Z
user: login: jasonvigil
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T00:20:46Z

Review comments:
diff_hunk: @@ -718,9 +719,10 @@ impl Workload {
                 Box::pin(wrap.step(client, metrics, data_set, state, done)).await
             }
             Workload::Load => {
-                if (state.seq_no * 100) >= data_set.cardinality() as u64 {
-                    done.store(true, std::sync::atomic::Ordering::Relaxed);
-                    return Ok(());
+                let mut key = (state.seq_no * 100) as usize;
+                if key >= data_set.cardinality() {
+                    // TODO: Probably should create another workload for this, so load can "finish".
path: rust/load/src/lib.rs
user: login: rescrv
body: Why this change?
created_at: 2025-04-22T15:07:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -718,9 +719,10 @@ impl Workload {
                 Box::pin(wrap.step(client, metrics, data_set, state, done)).await
             }
             Workload::Load => {
-                if (state.seq_no * 100) >= data_set.cardinality() as u64 {
-                    done.store(true, std::sync::atomic::Ordering::Relaxed);
-                    return Ok(());
+                let mut key = (state.seq_no * 100) as usize;
+                if key >= data_set.cardinality() {
+                    // TODO: Probably should create another workload for this, so load can "finish".
path: rust/load/src/lib.rs
user: login: jasonvigil
body: Chatted offline. Fixed this to keep the same behavior as before, by using an `Arc<Mutex<...>>` around the workload `state.seq_no`.
created_at: 2025-04-22T22:37:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -230,6 +232,9 @@ pub trait DataSet: std::fmt::Debug + Send + Sync {
     /// The number of documents in the data set.
     fn cardinality(&self) -> usize;
 
+    /// The number of documents in the reference data set.
+    fn reference_cardinality(&self) -> usize;
path: rust/load/src/lib.rs
user: login: rescrv
body: If you change this to `fn reference_cardinality(&self) -> usize { self.cardinality() }` you'll get the method for free everywhere and can override it when need be.
created_at: 2025-04-22T23:22:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -718,7 +719,13 @@ impl Workload {
                 Box::pin(wrap.step(client, metrics, data_set, state, done)).await
             }
             Workload::Load => {
-                if (state.seq_no * 100) >= data_set.cardinality() as u64 {
+                // Get the seq_no, increment it, and immediately drop the lock (to allow for concurrent loads)
+                let mut seq_no = state.seq_no.lock().await;
path: rust/load/src/lib.rs
user: login: rescrv
body: Another patter that's idiomatic is to:
```
{
    lock()
    do stuff()
}
```
so that you can avoid calling drop.  Just how I prefer to read mutexes.  Make your own way here once you decide which you like.
created_at: 2025-04-22T23:23:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -187,13 +188,14 @@ pub async fn client() -> ChromaClient {
 /// variable if set, or no authentication if unset.
 pub async fn client_for_url(url: String) -> ChromaClient {
     if let Ok(auth) = std::env::var("CHROMA_TOKEN") {
+        let db = std::env::var("CHROMA_DATABASE").unwrap_or_else(|_| "hf-tiny-stories".into());
path: rust/load/src/lib.rs
user: login: rescrv
body: I'm starting to think we should push these into a config rather than have new variables.  What do you think?
created_at: 2025-04-22T23:25:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -880,6 +896,322 @@ impl DataSet for ReferencingDataSet {
     }
 }
 
+//////////////////////////////////////// VerifyingDataSet /////////////////////////////////////////
+
+/// A verifying data set uses a reference data set to verify results from a test data set.
+#[derive(Debug)]
+pub struct VerifyingDataSet {
+    reference_data_set: Arc<dyn DataSet>,
+    test_data_set: String,
+    cardinality: AtomicUsize,
+    cardinality_heap: Arc<Mutex<BinaryHeap<Reverse<usize>>>>,
+}
+
+impl VerifyingDataSet {
+    async fn record_load(&self, new_cardinality: usize, batch_size: usize) {
path: rust/load/src/data_sets.rs
user: login: rescrv
body: This function needs a loop...
created_at: 2025-04-22T23:26:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -187,13 +188,14 @@ pub async fn client() -> ChromaClient {
 /// variable if set, or no authentication if unset.
 pub async fn client_for_url(url: String) -> ChromaClient {
     if let Ok(auth) = std::env::var("CHROMA_TOKEN") {
+        let db = std::env::var("CHROMA_DATABASE").unwrap_or_else(|_| "hf-tiny-stories".into());
path: rust/load/src/lib.rs
user: login: jasonvigil
body: Yeah, agreed. Will do this in a following PR.
created_at: 2025-04-23T00:15:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -230,6 +232,9 @@ pub trait DataSet: std::fmt::Debug + Send + Sync {
     /// The number of documents in the data set.
     fn cardinality(&self) -> usize;
 
+    /// The number of documents in the reference data set.
+    fn reference_cardinality(&self) -> usize;
path: rust/load/src/lib.rs
user: login: jasonvigil
body: Nice, done. Thanks
created_at: 2025-04-23T00:17:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -718,7 +719,13 @@ impl Workload {
                 Box::pin(wrap.step(client, metrics, data_set, state, done)).await
             }
             Workload::Load => {
-                if (state.seq_no * 100) >= data_set.cardinality() as u64 {
+                // Get the seq_no, increment it, and immediately drop the lock (to allow for concurrent loads)
+                let mut seq_no = state.seq_no.lock().await;
path: rust/load/src/lib.rs
user: login: jasonvigil
body: Ok, I did this. However, it felt a bit weird, because I also needed to assign a value for `load_start_idx` to be used in the outer scope. So, I defined the variable in the outer scope, then assigned the value in the inner scope. Not sure if what I did is the idiomatic way to do this, wdyt?
created_at: 2025-04-23T00:20:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -880,6 +896,322 @@ impl DataSet for ReferencingDataSet {
     }
 }
 
+//////////////////////////////////////// VerifyingDataSet /////////////////////////////////////////
+
+/// A verifying data set uses a reference data set to verify results from a test data set.
+#[derive(Debug)]
+pub struct VerifyingDataSet {
+    reference_data_set: Arc<dyn DataSet>,
+    test_data_set: String,
+    cardinality: AtomicUsize,
+    cardinality_heap: Arc<Mutex<BinaryHeap<Reverse<usize>>>>,
+}
+
+impl VerifyingDataSet {
+    async fn record_load(&self, new_cardinality: usize, batch_size: usize) {
path: rust/load/src/data_sets.rs
user: login: jasonvigil
body: Fixed!
created_at: 2025-04-23T00:20:46Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-21T23:53:18Z
body: # Reviewer Checklist
  Please leverage this checklist to ensure your code review is thorough before approving
  ## Testing, Bugs, Errors, Logs, Documentation
  - [ ] Can you think of any use case in which the code does not behave as intended? Have they been tested?
  - [ ] Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
  - [ ] If appropriate, are there adequate property based tests?
  - [ ] If appropriate, are there adequate unit tests?
  - [ ] Should any logging, debugging, tracing information be added or removed?
  - [ ] Are error messages user-friendly?
  - [ ] Have all documentation changes needed been made?
  - [ ] Have all non-obvious changes been commented?
  ## System Compatibility
  - [ ] Are there any potential impacts on other parts of the system or backward compatibility?
  - [ ] Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?
  ## Quality
  - [ ] Is this code of a unexpectedly high quality (Readability, Modularity, Intuitiveness)
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-21T23:53:21Z
body: Please tag your PR title with one of: \[ENH | BUG | DOC | TST | BLD | PERF | TYP | CLN | CHORE\]. See https://docs.trychroma.com/contributing#contributing-code-and-ideas
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-22T17:11:20Z
body: Please tag your PR title with one of: \[ENH | BUG | DOC | TST | BLD | PERF | TYP | CLN | CHORE\]. See https://docs.trychroma.com/contributing#contributing-code-and-ideas
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-22T18:29:29Z
body: Please tag your PR title with one of: \[ENH | BUG | DOC | TST | BLD | PERF | TYP | CLN | CHORE\]. See https://docs.trychroma.com/contributing#contributing-code-and-ideas
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-22T22:15:52Z
body: Please tag your PR title with one of: \[ENH | BUG | DOC | TST | BLD | PERF | TYP | CLN | CHORE\]. See https://docs.trychroma.com/contributing#contributing-code-and-ideas
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-22T22:20:10Z
body: Please tag your PR title with one of: \[ENH | BUG | DOC | TST | BLD | PERF | TYP | CLN | CHORE\]. See https://docs.trychroma.com/contributing#contributing-code-and-ideas
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-22T22:20:20Z
body: Please tag your PR title with one of: \[ENH | BUG | DOC | TST | BLD | PERF | TYP | CLN | CHORE\]. See https://docs.trychroma.com/contributing#contributing-code-and-ideas
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-22T22:20:22Z
body: Please tag your PR title with one of: \[ENH | BUG | DOC | TST | BLD | PERF | TYP | CLN | CHORE\]. See https://docs.trychroma.com/contributing#contributing-code-and-ideas
author_association: NONE