[ENH] Add hook to initialize collection before running verification (#4369)

## Description of changes

This adds a hook for workload executors to initialize the data sets
before starting the verification process. It resets the collection by
deleting + recreating it.

## Test plan

Tested locally, using same process as in
#4341, but without manual step
to create the test collection, since it's no longer necessary.

Author: jasonvigil

rust/load/src/data_sets.rs

@@ -972,6 +972,37 @@ impl DataSet for VerifyingDataSet {
         self.reference_data_set.cardinality()
     }
 
+    // Reset the test collection to an empty state by deleting and recreating it.
+    async fn initialize(
+        &self,
+        client: &ChromaClient,
+    ) -> Result<(), Box<dyn std::error::Error + Send>> {
+        // Attempt to delete the collection. If it doesn't exist, ignore the error.
+        match client.delete_collection(&self.test_data_set).await {
+            Ok(_) => (),
+            Err(err) => {
+                if !format!("{err:?}").contains("404") {
+                    return Err(Box::new(Error::InvalidRequest(format!(
+                        "failed to delete collection: {err:?}"
+                    ))));
+                }
+            }
+        };
+
+        // Create the collection.
+        match client
+            .create_collection(&self.test_data_set, None, true)
+            .await
+        {
+            Ok(_) => Ok(()),
+            Err(err) => {
+                return Err(Box::new(Error::InvalidRequest(format!(
+                    "failed to create collection: {err:?}"
+                ))));
+            }
+        }
+    }
+
     async fn get(
         &self,
         client: &ChromaClient,

rust/load/src/lib.rs

@@ -237,6 +237,11 @@ pub trait DataSet: std::fmt::Debug + Send + Sync {
         self.cardinality()
     }
 
+    // Hook to perform initialization of the data set, if necessary.
+    async fn initialize(&self, _: &ChromaClient) -> Result<(), Box<dyn std::error::Error + Send>> {
+        Ok(())
+    }
+
     /// 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(
@@ -1302,6 +1307,14 @@ impl LoadService {
                 }
             }
         });
+
+        // Initialize the data set.
+        let data_set = Arc::clone(&spec.data_set);
+        if let Err(err) = data_set.initialize(&client).await {
+            tracing::error!("failed to initialize data set: {err:?}");
+            return;
+        }
+
         let seq_no = Arc::new(TokioMutex::new(0u64));
         let start = Instant::now();
         while !done.load(std::sync::atomic::Ordering::Relaxed) {