From b3b4a8838463c87d1a57195614a6202b7089223a Mon Sep 17 00:00:00 2001
From: Sicheng Pan <sicheng@trychroma.com>
Date: Mon, 16 Dec 2024 17:42:40 -0800
Subject: [PATCH] Make count backward compat

---
 rust/worker/src/server.rs | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/rust/worker/src/server.rs b/rust/worker/src/server.rs
index 44ed213f91e7..da550483fc7a 100644
--- a/rust/worker/src/server.rs
+++ b/rust/worker/src/server.rs
@@ -9,7 +9,7 @@ use chroma_types::{
         self, query_executor_server::QueryExecutor, CountPlan, CountResult, GetPlan, GetResult,
         KnnBatchResult, KnnPlan,
     },
-    CollectionUuid, Segment, SegmentUuid,
+    CollectionUuid, SegmentUuid,
 };
 use futures::{stream, StreamExt, TryStreamExt};
 use tokio::signal::unix::{signal, SignalKind};
@@ -191,12 +191,7 @@ impl WorkerServer {
 
         let count_orchestrator = CountQueryOrchestrator::new(
             self.clone_system()?,
-            &Segment::try_from(
-                scan.metadata
-                    .ok_or(Status::invalid_argument("Invalid metadata segment"))?,
-            )?
-            .id
-            .0,
+            &SegmentUuid::from_str(&scan.metadata.map(|seg| seg.id).unwrap_or(scan.metadata_id))?.0,
             &CollectionUuid::from_str(&collection.id)
                 .map_err(|e| Status::invalid_argument(e.to_string()))?,
             self.log.clone(),