Only update new vectors

Author: blythed

plugins/qdrant/superduper_qdrant/qdrant.py

@@ -91,7 +91,7 @@ def __init__(
                 config_dict["timeout"] = 60  # 60 seconds timeout
         self.client = QdrantClient(**config_dict)
 
-        logging.info('Found these collections in the qdrant collection:')
+        logging.info('Found these collections in the qdrant connection:')
         collections = self.client.get_collections().collections
         for collection in collections:
             logging.info(f" - {collection.name}")
@@ -122,6 +122,45 @@ def initialize(self):
     def __len__(self):
         return self.client.get_collection(self.identifier).vectors_count
 
+    def list(self):
+        """List all ids in the index."""
+        if not self.client.collection_exists(self.identifier):
+            return []
+
+        ids: t.List[str] = []
+        next_page: t.Optional[models.ScrollCursor] = None
+        limit = max(256, self.batch_size)  # reasonable paging size
+
+        # Reuse the same retry pattern as upserts
+        retry = Retry(exception_types=QDRANT_RETRY_EXCEPTIONS)
+
+        @retry
+        def _do_scroll(offset):
+            return self.client.scroll(
+                collection_name=self.identifier,
+                limit=limit,
+                with_payload=[ID_PAYLOAD_KEY],
+                with_vectors=False,
+                offset=offset,
+            )
+
+        while True:
+            records, next_page = _do_scroll(next_page)
+
+            if not records:
+                break
+
+            for rec in records:
+                # We store original IDs in the payload under ID_PAYLOAD_KEY
+                if rec.payload and ID_PAYLOAD_KEY in rec.payload:
+                    ids.append(rec.payload[ID_PAYLOAD_KEY])
+
+            if next_page is None:
+                break
+
+        return ids
+
+
     def _create_collection(self):
         measure = (
             self.measure.name

superduper/backends/base/vector_search.py

@@ -187,6 +187,11 @@ def db(self) -> 'Datalayer':
     def db(self, value: 'Datalayer'):
         self._db = value
 
+    @abstractmethod
+    def list(self) -> t.List[str]:
+        """List all ids in the index."""
+        pass
+
     @classmethod
     def from_component(cls, index: 'VectorIndex'):
         """Create a vector searcher from a vector index.

superduper/backends/local/vector_search.py

@@ -50,15 +50,26 @@ def initialize(self):
             if t.is_component and t.cls is not None:
                 if issubclass(t.cls, VectorIndex):
                     components.append(t.identifier)
+                    
         for component in components:
             try:
                 for identifier in self.db.show(component):
                     try:
                         vector_index = self.db.load(component, identifier=identifier)
                         self.put_component(component, vector_index.uuid)
-                        vectors = vector_index.get_vectors()
-                        vectors = [VectorItem(**vector) for vector in vectors]
-                        self.get_tool(vector_index.uuid).add(vectors)
+                        all_ids = vector_index.list()
+                        tool = self.get_tool(vector_index.uuid)
+                        deployed_ids = tool.list()
+                        to_deploy_ids = list(set(all_ids) - set(deployed_ids))
+                        if to_deploy_ids:
+                            vectors = vector_index.get_vectors(ids=to_deploy_ids)
+                            vectors = [VectorItem(**vector) for vector in vectors]
+                            self.get_tool(vector_index.uuid).add(vectors)
+                        else:
+                            logging.info(
+                                'Skipping since have already deployed vectors for '
+                                f'component: {component}, vector_index: {vector_index.uuid}'
+                            )
 
                     except FileNotFoundError:
                         logging.error(
@@ -191,6 +202,9 @@ def _setup(self, h, index):
         self.index = index
         self.lookup = dict(zip(index, range(len(index))))
 
+    def list(self):
+        return self.index if self.index is not None else []
+
     def describe(self):
         """Describe the vector index.
 

superduper/components/vector_index.py

@@ -85,6 +85,9 @@ def setup(self):
             )
         return self
 
+    def list(self):
+        return self.db[self.indexing_listener.outputs].ids()
+
     @ensure_setup
     def get_vectors(self, ids: t.Sequence[str] | None = None):
         """Get vectors from the vector index.