Make sure broker messages from vectorset migrations are properly ingested (#2764)

* Fix has_field usage

* Get fields before applying extracted vectors

* Add test

* add test

* Replace sentences_to_delete for vector_prefixes_to_delete

* Implement nidx vector deletion by vectorset

* Implement the same for node

* Populate vector_prefixes_to_delete if vectorsets are used

* Remove useless datamanager vectorset creation

* Linter

* Use inject_message to mark nidx dirty

* Extend test to validate keyword search

* Fix paragraphs_to_delete

* Pass replace_field parameter

---------

Co-authored-by: Joan Antoni RE <joanantoni.re16@gmail.com>

Author: lferran

nidx/nidx_binding/pyproject.toml

@@ -11,5 +11,6 @@ classifiers = [
     "Programming Language :: Python :: Implementation :: PyPy",
 ]
 dynamic = ["version"]
+
 [tool.maturin]
 features = ["pyo3/extension-module"]

nidx/nidx_vector/src/data_point_provider/reader.rs

@@ -395,16 +395,10 @@ mod tests {
         };
         let resource = Resource {
             resource: Some(resource_id),
-            metadata: None,
             texts: HashMap::with_capacity(0),
             status: ResourceStatus::Processed as i32,
             labels: vec!["2".to_string()],
             paragraphs: HashMap::from([("DOC/KEY".to_string(), paragraphs)]),
-            paragraphs_to_delete: vec![],
-            sentences_to_delete: vec![],
-            relations: vec![],
-            vectors: HashMap::default(),
-            vectors_to_delete: HashMap::default(),
             shard_id: "DOC".to_string(),
             ..Default::default()
         };
@@ -486,16 +480,10 @@ mod tests {
         };
         let resource = Resource {
             resource: Some(resource_id),
-            metadata: None,
             texts: HashMap::with_capacity(0),
             status: ResourceStatus::Processed as i32,
             labels: vec!["2".to_string()],
             paragraphs: HashMap::from([("DOC/KEY".to_string(), paragraphs)]),
-            paragraphs_to_delete: vec![],
-            sentences_to_delete: vec![],
-            relations: vec![],
-            vectors: HashMap::default(),
-            vectors_to_delete: HashMap::default(),
             shard_id: "DOC".to_string(),
             ..Default::default()
         };
@@ -608,16 +596,10 @@ mod tests {
         };
         let resource = Resource {
             resource: Some(resource_id),
-            metadata: None,
             texts: HashMap::with_capacity(0),
             status: ResourceStatus::Processed as i32,
             labels: vec!["2".to_string()],
             paragraphs: HashMap::from([("DOC/KEY".to_string(), paragraphs)]),
-            paragraphs_to_delete: vec![],
-            sentences_to_delete: vec![],
-            relations: vec![],
-            vectors: HashMap::default(),
-            vectors_to_delete: HashMap::default(),
             shard_id: "DOC".to_string(),
             ..Default::default()
         };

nidx/nidx_vector/src/lib.rs

@@ -67,8 +67,14 @@ impl VectorIndexer {
         index_resource(vectorset_resource, output_dir, config)
     }
 
-    pub fn deletions_for_resource(&self, resource: &Resource) -> Vec<String> {
-        resource.sentences_to_delete.clone()
+    pub fn deletions_for_resource(&self, resource: &Resource, index_name: &str) -> Vec<String> {
+        if let Some(prefixes) = resource.vector_prefixes_to_delete.get(index_name) {
+            prefixes.items.clone()
+        } else {
+            // DEPRECATED: Bw/c while moving from sentences_to_delete to vector_prefixes_to_delete
+            #[allow(deprecated)]
+            resource.sentences_to_delete.clone()
+        }
     }
 
     #[instrument(name = "vector::merge", skip_all)]

nidx/src/indexer.rs

@@ -263,7 +263,7 @@ fn index_resource_to_index(
     };
 
     let deletions = match index.kind {
-        IndexKind::Vector => nidx_vector::VectorIndexer.deletions_for_resource(resource),
+        IndexKind::Vector => nidx_vector::VectorIndexer.deletions_for_resource(resource, &index.name),
         IndexKind::Text => nidx_text::TextIndexer.deletions_for_resource(resource),
         IndexKind::Paragraph => nidx_paragraph::ParagraphIndexer.deletions_for_resource(resource),
         IndexKind::Relation => nidx_relation::RelationIndexer.deletions_for_resource(resource),
@@ -277,6 +277,7 @@ fn index_resource_to_index(
 mod tests {
     use std::io::{Seek, Write};
 
+    use nidx_protos::StringList;
     use nidx_vector::config::{Similarity, VectorConfig, VectorType};
     use tempfile::tempfile;
     use uuid::Uuid;
@@ -354,7 +355,12 @@ mod tests {
             format!("{}/a/title/0-15", resource.resource.as_ref().unwrap().uuid),
             format!("{}/a/summary/0-150", resource.resource.as_ref().unwrap().uuid),
         ];
-        resource.sentences_to_delete.append(&mut keys.clone());
+        resource.vector_prefixes_to_delete.insert(
+            index.name,
+            StringList {
+                items: keys.clone(),
+            },
+        );
         index_resource(
             &meta,
             storage.clone(),

nidx/tests/test_date_range_search.rs

@@ -92,14 +92,8 @@ async fn populate(fixture: &mut NidxFixture, shard_id: String, metadata: IndexMe
         resource: Some(resource_id),
         metadata: Some(metadata),
         status: ResourceStatus::Processed as i32,
-        labels: vec![],
         texts: HashMap::from([(field_id.clone(), text_content)]),
         paragraphs: HashMap::from([(format!("{raw_resource_id}/{field_id}"), paragraphs)]),
-        paragraphs_to_delete: vec![],
-        sentences_to_delete: vec![],
-        relations: vec![],
-        vectors: HashMap::default(),
-        vectors_to_delete: HashMap::default(),
         ..Default::default()
     };
 

nucliadb/src/nucliadb/common/external_index_providers/pinecone.py

@@ -441,6 +441,7 @@ def get_index_host(self, vectorset_id: str, rollover: bool = False) -> str:
 
     def get_prefixes_to_delete(self, index_data: Resource) -> set[str]:
         prefixes_to_delete = set()
+        # TODO: migrate to vector_prefixes_to_delete
         for field_id in index_data.sentences_to_delete:
             try:
                 delete_vid = VectorId.from_string(field_id)

nucliadb/src/nucliadb/ingest/orm/brain.py

@@ -100,6 +100,8 @@ def apply_field_metadata(
         page_positions: Optional[FilePagePositions],
         extracted_text: Optional[ExtractedText],
         basic_user_field_metadata: Optional[UserFieldMetadata] = None,
+        *,
+        replace_field: bool = False,
     ):
         # To check for duplicate paragraphs
         unique_paragraphs: set[str] = set()
@@ -224,6 +226,11 @@ def apply_field_metadata(
 
             self.brain.paragraphs[field_key].paragraphs[key].CopyFrom(p)
 
+        if replace_field:
+            field_type, field_name = field_key.split("/")
+            full_field_id = ids.FieldId(rid=self.rid, type=field_type, key=field_name).full()
+            self.brain.paragraphs_to_delete.append(full_field_id)
+
         for relations in metadata.metadata.relations:
             for relation in relations.relations:
                 self.brain.relations.append(relation)
@@ -301,8 +308,11 @@ def apply_field_vectors(
 
         if replace_field:
             full_field_id = ids.FieldId(rid=self.rid, type=fid.type, key=fid.key).full()
-            self.brain.sentences_to_delete.append(full_field_id)
-            self.brain.paragraphs_to_delete.append(full_field_id)
+            if vectorset is None:
+                # DEPRECATED
+                self.brain.sentences_to_delete.append(full_field_id)
+            else:
+                self.brain.vector_prefixes_to_delete[vectorset].items.append(full_field_id)
 
     def _apply_field_vector(
         self,

nucliadb/src/nucliadb/ingest/orm/processor/processor.py

@@ -275,7 +275,6 @@ async def txn(
 
                     if message.source == writer_pb2.BrokerMessage.MessageSource.WRITER:
                         resource = await kb.get(uuid)
-
                         if resource is None:
                             # It's a new resource
                             resource = await kb.add_resource(uuid, message.slug, message.basic)
@@ -737,7 +736,11 @@ def has_vectors_operation(index_message: PBBrainResource) -> bool:
     """
     Returns True if the index message has any vectors to index or to delete.
     """
-    if len(index_message.sentences_to_delete) > 0 or len(index_message.paragraphs_to_delete) > 0:
+    if (
+        len(index_message.sentences_to_delete) > 0
+        or len(index_message.paragraphs_to_delete) > 0
+        or any([len(deletions.items) for deletions in index_message.vector_prefixes_to_delete.values()])
+    ):
         return True
     for field_paragraphs in index_message.paragraphs.values():
         for paragraph in field_paragraphs.paragraphs.values():

nucliadb/src/nucliadb/ingest/orm/resource.py

@@ -226,6 +226,7 @@ async def set_basic(
                             page_positions=page_positions,
                             extracted_text=await field_obj.get_extracted_text(),
                             basic_user_field_metadata=user_field_metadata,
+                            replace_field=True,
                         )
 
         # Some basic fields are computed off field metadata.
@@ -336,6 +337,7 @@ async def generate_index_message(self, reindex: bool = False) -> ResourceBrain:
                     page_positions=page_positions,
                     extracted_text=await field.get_extracted_text(),
                     basic_user_field_metadata=user_field_metadata,
+                    replace_field=reindex,
                 )
 
             if self.disable_vectors is False:
@@ -584,6 +586,7 @@ async def apply_extracted(self, message: BrokerMessage):
         # Upload to binary storage
         # Vector indexing
         if self.disable_vectors is False:
+            await self.get_fields(force=True)
             for field_vectors in message.field_vectors:
                 await self._apply_extracted_vectors(field_vectors)
 
@@ -723,6 +726,7 @@ async def _apply_field_computed_metadata(self, field_metadata: FieldComputedMeta
             page_positions=page_positions,
             extracted_text=extracted_text,
             basic_user_field_metadata=user_field_metadata,
+            replace_field=True,
         )
         loop = asyncio.get_running_loop()
         await loop.run_in_executor(_executor, apply_field_metadata)

nucliadb/src/nucliadb/writer/api/v1/knowledgebox.py

@@ -68,6 +68,7 @@ async def create_kb_endpoint(request: Request, item: KnowledgeBoxConfig) -> Know
     except ExternalIndexCreationError as exc:
         raise HTTPException(status_code=502, detail=str(exc))
     except Exception:
+        logger.exception("Could not create KB")
         raise HTTPException(status_code=500, detail="Error creating knowledge box")
     else:
         return KnowledgeBoxObj(uuid=kbid, slug=slug)