Skip to content

Commit

Permalink
fix: Calling manager.delete_segments first
Browse files Browse the repository at this point in the history
  • Loading branch information
tazarov committed Jan 9, 2025
1 parent fa0e206 commit 788a07f
Show file tree
Hide file tree
Showing 4 changed files with 12 additions and 11 deletions.
13 changes: 8 additions & 5 deletions chromadb/api/segment.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.")

Expand Down Expand Up @@ -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],
Expand Down
2 changes: 1 addition & 1 deletion chromadb/segment/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
5 changes: 1 addition & 4 deletions chromadb/segment/impl/manager/distributed.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
3 changes: 2 additions & 1 deletion chromadb/segment/impl/manager/local.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down

0 comments on commit 788a07f

Please sign in to comment.