diff --git a/chromadb/api/segment.py b/chromadb/api/segment.py index 381ec5b9a39..fc92856ef6a 100644 --- a/chromadb/api/segment.py +++ b/chromadb/api/segment.py @@ -384,11 +384,10 @@ def delete_collection( ) if existing: - segments = self._sysdb.get_segments(collection=existing[0].id) + self._manager.delete_segments(collection_id=existing[0].id) self._sysdb.delete_collection( existing[0].id, tenant=tenant, database=database ) - self._manager.delete_segments(segments) else: raise ValueError(f"Collection {name} does not exist.") @@ -895,11 +894,15 @@ def _get_collection(self, collection_id: UUID) -> t.Collection: @trace_method("SegmentAPI._scan", OpenTelemetryGranularity.ALL) def _scan(self, collection_id: UUID) -> Scan: - collection_and_segments = self._sysdb.get_collection_with_segments(collection_id) + collection_and_segments = self._sysdb.get_collection_with_segments( + collection_id + ) # For now collection should have exactly one segment per scope: # - Local scopes: vector, metadata - # - Distributed scopes: vector, metadata, record - scope_to_segment = {segment["scope"]: segment for segment in collection_and_segments["segments"]} + # - Distributed scopes: vector, metadata, record + scope_to_segment = { + segment["scope"]: segment for segment in collection_and_segments["segments"] + } return Scan( collection=collection_and_segments["collection"], knn=scope_to_segment[t.SegmentScope.VECTOR], diff --git a/chromadb/segment/__init__.py b/chromadb/segment/__init__.py index b41ecb32bbf..bd5fdc37151 100644 --- a/chromadb/segment/__init__.py +++ b/chromadb/segment/__init__.py @@ -112,7 +112,7 @@ def prepare_segments_for_new_collection( pass @abstractmethod - def delete_segments(self, segments: Sequence[Segment]) -> Sequence[UUID]: + def delete_segments(self, collection_id: UUID) -> Sequence[UUID]: """Delete any local state for all the segments associated with a collection, and returns a sequence of their IDs. Does not update the SysDB.""" pass diff --git a/chromadb/segment/impl/manager/distributed.py b/chromadb/segment/impl/manager/distributed.py index 4367ab1c44e..91b9d7e1304 100644 --- a/chromadb/segment/impl/manager/distributed.py +++ b/chromadb/segment/impl/manager/distributed.py @@ -78,10 +78,7 @@ def prepare_segments_for_new_collection( @override def delete_segments(self, collection_id: UUID) -> Sequence[UUID]: - # TODO: this should be a pass, delete_collection is expected to delete segments in - # distributed - segments = self._sysdb.get_segments(collection=collection_id) - return [s["id"] for s in segments] + return [] # noop @trace_method( "DistributedSegmentManager.get_endpoint", diff --git a/chromadb/segment/impl/manager/local.py b/chromadb/segment/impl/manager/local.py index b126dd8c267..042957f2512 100644 --- a/chromadb/segment/impl/manager/local.py +++ b/chromadb/segment/impl/manager/local.py @@ -153,7 +153,8 @@ def prepare_segments_for_new_collection( OpenTelemetryGranularity.OPERATION_AND_SEGMENT, ) @override - def delete_segments(self, segments: Sequence[Segment]) -> Sequence[UUID]: + def delete_segments(self, collection_id: UUID) -> Sequence[UUID]: + segments = self._sysdb.get_segments(collection=collection_id) for segment in segments: collection_id = segment["collection"] if segment["id"] in self._instances: