[ENH]: Rust frontend proptest

Cargo.toml

@@ -81,8 +81,8 @@ worker = { path = "rust/worker" }
 # Dev dependencies
 bincode = "1.3.3"
 indicatif = "0.17.9"
-proptest = "1.5.0"
-proptest-state-machine = "0.3.0"
+proptest = "1.6.0"
+proptest-state-machine = "0.3.1"
 proptest-derive = "0.5.1"
 rand = "0.8.5"
 rand_xorshift = "0.3.0"

rust/frontend/Cargo.toml

@@ -50,3 +50,8 @@ utoipa-swagger-ui = { version = "9", features = ["axum"] }
 [dev-dependencies]
 reqwest = { workspace = true }
 random-port = { workspace = true }
+proptest = { workspace = true }
+proptest-state-machine = { workspace = true }
+chroma-types = { workspace = true, features = ["testing"] }
+tempfile = { workspace = true }
+worker = { workspace = true }

rust/frontend/src/executor/config.rs

@@ -28,6 +28,21 @@ pub struct DistributedExecutorConfig {
     pub memberlist_provider: chroma_memberlist::config::MemberlistProviderConfig,
 }
 
+// todo
+impl Default for DistributedExecutorConfig {
+    fn default() -> Self {
+        DistributedExecutorConfig {
+            connections_per_node: 5,
+            replication_factor: 3,
+            connect_timeout_ms: 5000,
+            request_timeout_ms: 5000,
+            retry: RetryConfig::default(),
+            assignment: chroma_config::assignment::config::AssignmentPolicyConfig::default(),
+            memberlist_provider: chroma_memberlist::config::MemberlistProviderConfig::default(),
+        }
+    }
+}
+
 #[derive(Deserialize, Clone, Serialize, Debug)]
 pub struct LocalExecutorConfig {}
 

rust/frontend/src/executor/local.rs

@@ -126,7 +126,18 @@ impl LocalExecutor {
             .await?;
         if let Some(dimensionality) = collection_and_segments.collection.dimension {
             let allowed_user_ids = if plan.filter.where_clause.is_none() {
-                plan.filter.query_ids.unwrap_or_default()
+                match plan.filter.query_ids {
+                    Some(ids) => {
+                        if ids.is_empty() {
+                            return Ok(vec![Default::default(); plan.knn.embeddings.len()]);
+                        }
+
+                        ids
+                    }
+                    None => {
+                        vec![]
+                    }
+                }
             } else {
                 let filter_plan = Get {
                     scan: plan.scan.clone(),

rust/frontend/src/impls/in_memory_frontend.rs

@@ -26,7 +26,7 @@ struct Inner {
     collections: Vec<InMemoryCollection>,
 }
 
-#[derive(Clone, Default)]
+#[derive(Clone, Default, Debug)]
 pub struct InMemoryFrontend {
     inner: Inner,
 }
@@ -233,22 +233,25 @@ impl InMemoryFrontend {
             last_compaction_time_secs: 0,
         };
 
-        let reference_impl = TestReferenceSegment::default();
+        let metadata_segment = test_segment(
+            collection.collection_id,
+            chroma_types::SegmentScope::METADATA,
+        );
+        let vector_segment =
+            test_segment(collection.collection_id, chroma_types::SegmentScope::VECTOR);
+        let record_segment =
+            test_segment(collection.collection_id, chroma_types::SegmentScope::RECORD);
+
+        let mut reference_impl = TestReferenceSegment::default();
+        reference_impl.create_segment(metadata_segment.clone());
+        reference_impl.create_segment(vector_segment.clone());
+        reference_impl.create_segment(record_segment.clone());
 
         self.inner.collections.push(InMemoryCollection {
             collection: collection.clone(),
-            metadata_segment: test_segment(
-                collection.collection_id,
-                chroma_types::SegmentScope::METADATA,
-            ),
-            vector_segment: test_segment(
-                collection.collection_id,
-                chroma_types::SegmentScope::VECTOR,
-            ),
-            record_segment: test_segment(
-                collection.collection_id,
-                chroma_types::SegmentScope::RECORD,
-            ),
+            metadata_segment,
+            vector_segment,
+            record_segment,
             reference_impl,
         });
 
@@ -413,7 +416,7 @@ impl InMemoryFrontend {
             documents,
             uris,
             metadatas,
-            chroma_types::Operation::Add,
+            chroma_types::Operation::Upsert,
         )
         .map_err(|e| e.boxed())?;
 

rust/frontend/src/impls/service_based_frontend.rs

@@ -878,7 +878,7 @@ impl ServiceBasedFrontend {
         Ok((get_result, include).into())
     }
 
-    pub async fn get(&mut self, request: GetRequest) -> Result<GetResponse, QueryError> {
+    pub async fn get(&self, request: GetRequest) -> Result<GetResponse, QueryError> {
         let retries = Arc::new(AtomicUsize::new(0));
         let get_to_retry = || {
             let mut self_clone = self.clone();