diff --git a/chromadb/execution/executor/distributed.py b/chromadb/execution/executor/distributed.py
index 6479677d6ba..dc5eda48522 100644
--- a/chromadb/execution/executor/distributed.py
+++ b/chromadb/execution/executor/distributed.py
@@ -2,7 +2,6 @@
 # TODO: This will be refactored after grpc defs have been updated for query pushdown
 
 from typing import Optional, Sequence
-from uuid import UUID
 
 from chromadb.api.types import (
     GetResult,
@@ -16,7 +15,7 @@
 from chromadb.segment.impl.manager.distributed import DistributedSegmentManager
 from chromadb.segment.impl.metadata.grpc_segment import GrpcMetadataSegment
 from chromadb.segment.impl.vector.grpc_segment import GrpcVectorSegment
-from chromadb.types import VectorQuery, VectorQueryResult
+from chromadb.types import VectorQuery, VectorQueryResult, Collection
 from overrides import overrides
 
 
@@ -180,8 +179,8 @@ def knn(self, plan: KNNPlan) -> QueryResult:
             included=included,
         )
 
-    def _metadata_segment(self, collection_id: UUID) -> GrpcMetadataSegment:
-        return self._manager.get_segment(collection_id, GrpcMetadataSegment)
+    def _metadata_segment(self, collection: Collection) -> GrpcMetadataSegment:
+        return self._manager.get_segment(collection.id, GrpcMetadataSegment)
 
-    def _vector_segment(self, collection_id: UUID) -> GrpcVectorSegment:
-        return self._manager.get_segment(collection_id, GrpcVectorSegment)
+    def _vector_segment(self, collection: Collection) -> GrpcVectorSegment:
+        return self._manager.get_segment(collection.id, GrpcVectorSegment)
diff --git a/chromadb/execution/executor/local.py b/chromadb/execution/executor/local.py
index 444fbf1a8ef..1378138e671 100644
--- a/chromadb/execution/executor/local.py
+++ b/chromadb/execution/executor/local.py
@@ -1,5 +1,4 @@
 from typing import Optional, Sequence
-from uuid import UUID
 
 from chromadb.api.types import (
     GetResult,
@@ -13,7 +12,7 @@
 from chromadb.segment.impl.manager.local import LocalSegmentManager
 from chromadb.segment.impl.metadata.sqlite import SqliteMetadataSegment
 from chromadb.segment.impl.vector.local_hnsw import LocalHnswSegment
-from chromadb.types import VectorQuery, VectorQueryResult
+from chromadb.types import VectorQuery, VectorQueryResult, Collection
 from overrides import overrides
 
 
@@ -177,8 +176,8 @@ def knn(self, plan: KNNPlan) -> QueryResult:
             included=included,
         )
 
-    def _metadata_segment(self, collection_id: UUID) -> SqliteMetadataSegment:
-        return self._manager.get_segment(collection_id, SqliteMetadataSegment)
+    def _metadata_segment(self, collection: Collection) -> SqliteMetadataSegment:
+        return self._manager.get_segment(collection.id, SqliteMetadataSegment)
 
-    def _vector_segment(self, collection_id: UUID) -> LocalHnswSegment:
-        return self._manager.get_segment(collection_id, LocalHnswSegment)
+    def _vector_segment(self, collection: Collection) -> LocalHnswSegment:
+        return self._manager.get_segment(collection.id, LocalHnswSegment)