Remove deduplication for relation node/edge vectors (#3643)

* Remove deduplication

* Reorder protos to avoid storing field_id

Author: javitonino

nidx/nidx_protos/noderesources.proto

@@ -175,8 +175,16 @@ message Resource {
     bool skip_paragraphs = 19;
     bool skip_json = 24;
 
-    map<string, utils.RelationNodeVectors> relation_node_vectors = 20;
-    map<string, utils.RelationEdgeVectors> relation_edge_vectors = 21;
+    map<string, IndexFieldNodeVectors> field_node_vectors = 20;
+    map<string, IndexFieldEdgeVectors> field_edge_vectors = 21;
+}
+
+message IndexFieldNodeVectors {
+    map<string, utils.RelationNodeVectors> node_vectors = 1; // key: vectorset_id
+}
+
+message IndexFieldEdgeVectors {
+    map<string, utils.RelationEdgeVectors> edge_vectors = 1; // key: vectorset_id
 }
 
 message ShardMetadata {

nidx/nidx_vector/src/config.rs

@@ -176,12 +176,10 @@ impl VectorConfig {
             && matches!(&self.vector_type, VectorType::DenseF32 { dimension } if dimension.is_multiple_of(64))
     }
 
-    /// Whether to deduplicate paragraphs based on their key, using `metadata` to store the list of field_keys
-    pub fn deduplicate_keys(&self) -> bool {
-        match self.entity {
-            IndexEntity::Paragraph => false,
-            IndexEntity::RelationNode | IndexEntity::RelationEdge => true,
-        }
+    /// Whether this index uses relation-style inverted indexes (field mapping from metadata)
+    /// rather than paragraph-style inverted indexes (field mapping from key).
+    pub fn uses_relation_inverted_index(&self) -> bool {
+        matches!(self.entity, IndexEntity::RelationNode | IndexEntity::RelationEdge)
     }
 
     pub fn from_paragraph_proto(proto: VectorIndexConfig) -> VectorR<Self> {

nidx/nidx_vector/src/data_store/v2.rs

@@ -29,7 +29,7 @@ use crate::{
 use super::{DataStore, OpenReason, ParagraphAddr, VectorAddr};
 pub use paragraph_store::StoredParagraph;
 use paragraph_store::{ParagraphStore, ParagraphStoreWriter};
-use std::{collections::HashMap, path::Path};
+use std::path::Path;
 use vector_store::{VectorStore, VectorStoreWriter};
 
 mod paragraph_store;
@@ -88,7 +88,6 @@ impl DataStoreV2 {
         path: &Path,
         producers: Vec<(impl Iterator<Item = ParagraphAddr>, &dyn DataStore)>,
         config: &VectorConfig,
-        mut paragraph_deduplicator: Option<HashMap<String, Vec<u8>>>,
     ) -> VectorR<()> {
         let mut paragraphs = ParagraphStoreWriter::new(path)?;
         let mut vectors = VectorStoreWriter::new(path, &config.vector_type)?;
@@ -105,17 +104,6 @@ impl DataStoreV2 {
                 let paragraph = store.get_paragraph(paragraph_addr);
                 let p_vectors = paragraph.vectors(&paragraph_addr).map(|v| store.get_vector(v).vector());
 
-                let metadata = if let Some(paragraph_deduplicator) = &mut paragraph_deduplicator {
-                    // Entry is removed so if it appears in other segments it is not copied again
-                    let metadata = paragraph_deduplicator.remove(paragraph.id());
-                    if metadata.is_none() {
-                        continue;
-                    };
-                    metadata
-                } else {
-                    None
-                };
-
                 // Write to new store
                 let (first_vector, last_vector) = vectors.write(p_idx, p_vectors)?;
                 if let Some(quantized) = &mut quantized {
@@ -132,12 +120,7 @@ impl DataStoreV2 {
                     }
                 }
 
-                paragraphs.write_paragraph_ref(
-                    paragraph,
-                    first_vector,
-                    last_vector - first_vector + 1,
-                    metadata.as_deref(),
-                )?;
+                paragraphs.write_paragraph_ref(paragraph, first_vector, last_vector - first_vector + 1)?;
 
                 p_idx += 1;
             }

nidx/nidx_vector/src/data_store/v2/paragraph_store.rs

@@ -152,13 +152,12 @@ impl ParagraphStoreWriter {
         paragraph: ParagraphRef,
         first_vector: u32,
         num_vectors: u32,
-        override_metadata: Option<&[u8]>,
     ) -> VectorR<u32> {
         let labels = paragraph.labels();
         let paragraph = StoredParagraph {
             key: paragraph.id(),
             labels: labels.iter().map(|x| x.as_str()).collect(),
-            metadata: override_metadata.unwrap_or(paragraph.metadata()),
+            metadata: paragraph.metadata(),
             first_vector,
             num_vectors,
         };

nidx/nidx_vector/src/field_list_metadata.rs

@@ -18,7 +18,6 @@
 // along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 use crate::config::{VectorCardinality, VectorConfig};
-use crate::field_list_metadata::encode_field_list_metadata;
 use crate::multivector::extract_multi_vectors;
 use crate::segment::{self, Elem};
 use crate::utils::FieldKey;
@@ -152,13 +151,8 @@ pub fn index_resource(
     Ok(Some(segment.into_metadata()))
 }
 
-fn encode_metadata_field(rid: &str, fields: &HashSet<&String>) -> Vec<u8> {
-    let encoded_fields: Vec<_> = fields
-        .iter()
-        .map(|f| format!("{rid}/{f}"))
-        .filter_map(|f| FieldKey::from_field_id(&f))
-        .collect();
-    encode_field_list_metadata(&encoded_fields)
+fn encode_metadata_field(rid: &str, field: &str) -> Option<Vec<u8>> {
+    FieldKey::from_field_id(&format!("{rid}/{field}")).map(|k| k.bytes().to_vec())
 }
 
 pub fn index_relation_nodes(
@@ -169,50 +163,25 @@ pub fn index_relation_nodes(
 ) -> anyhow::Result<Option<VectorSegmentMetadata>> {
     debug!("Creating elements for the main index");
 
-    let mut entity_fields = HashMap::new();
     let Some(resource_id) = &resource.resource else {
         return Err(anyhow!("resource_id required"));
     };
     let rid = &resource_id.uuid;
 
-    for (field, relations) in &resource.field_relations {
-        // Find all copies of each relation node
-        for relation in &relations.relations {
-            let Some(relation) = &relation.relation else {
-                return Err(anyhow!("relation required"));
-            };
-            let Some(source) = &relation.source else {
-                return Err(anyhow!("relation source node required"));
-            };
-            entity_fields
-                .entry(source.value.clone())
-                .or_insert_with(HashSet::new)
-                .insert(field);
-
-            let Some(to) = &relation.to else {
-                return Err(anyhow!("relation to node required"));
-            };
-            entity_fields
-                .entry(to.value.clone())
-                .or_insert_with(HashSet::new)
-                .insert(field);
-        }
-    }
-
-    // Index each vector
     let mut elems = Vec::new();
-    if let Some(vectorset) = &resource.relation_node_vectors.get(index_name) {
+    for (field_id, field_data) in &resource.field_node_vectors {
+        let Some(vectorset) = field_data.node_vectors.get(index_name) else {
+            continue;
+        };
+        let Some(metadata) = encode_metadata_field(rid, field_id) else {
+            continue;
+        };
         for node_vector in &vectorset.vectors {
-            let vector = node_vector.vector.clone();
-            let fields = entity_fields.get(&node_vector.node_value);
-            let Some(fields) = fields else {
-                continue;
-            };
             elems.push(Elem::new(
                 node_vector.node_value.clone(),
-                vector,
+                node_vector.vector.clone(),
                 vec![],
-                Some(encode_metadata_field(rid, fields)),
+                Some(metadata.clone()),
             ));
         }
     }
@@ -235,40 +204,25 @@ pub fn index_relation_edges(
 ) -> anyhow::Result<Option<VectorSegmentMetadata>> {
     debug!("Creating elements for the main index");
 
-    let mut entity_fields = HashMap::new();
     let Some(resource_id) = &resource.resource else {
         return Err(anyhow!("resource_id required"));
     };
     let rid = &resource_id.uuid;
 
-    for (field, relations) in &resource.field_relations {
-        // Find all copies of each relation edge
-        for relation in &relations.relations {
-            let Some(relation) = &relation.relation else {
-                return Err(anyhow!("relation required"));
-            };
-
-            entity_fields
-                .entry(relation.relation_label.clone())
-                .or_insert_with(HashSet::new)
-                .insert(field);
-        }
-    }
-
-    // Index each vector
     let mut elems = Vec::new();
-    if let Some(vectorset) = &resource.relation_edge_vectors.get(index_name) {
+    for (field_id, field_data) in &resource.field_edge_vectors {
+        let Some(vectorset) = field_data.edge_vectors.get(index_name) else {
+            continue;
+        };
+        let Some(metadata) = encode_metadata_field(rid, field_id) else {
+            continue;
+        };
         for rel_vector in &vectorset.vectors {
-            let vector = rel_vector.vector.clone();
-            let fields = entity_fields.get(&rel_vector.relation_label);
-            let Some(fields) = fields else {
-                continue;
-            };
             elems.push(Elem::new(
                 rel_vector.relation_label.clone(),
-                vector,
+                rel_vector.vector.clone(),
                 vec![],
-                Some(encode_metadata_field(rid, fields)),
+                Some(metadata.clone()),
             ));
         }
     }

nidx/nidx_vector/src/indexer.rs

@@ -66,7 +66,7 @@ impl IndexBuilder {
 
 /// Build indexes from a DataStore.
 pub fn build_indexes(work_path: &Path, config: &VectorConfig, data_store: &impl DataStore) -> VectorR<()> {
-    if config.deduplicate_keys() {
+    if config.uses_relation_inverted_index() {
         RelationInvertedIndexes::build(work_path, data_store)
     } else {
         ParagraphInvertedIndexes::build(work_path, data_store)
@@ -84,7 +84,7 @@ pub enum InvertedIndexes {
 
 impl InvertedIndexes {
     pub fn open(config: &VectorConfig, work_path: &Path, records: usize, options: OpenOptions) -> VectorR<Self> {
-        if config.deduplicate_keys() {
+        if config.uses_relation_inverted_index() {
             Ok(InvertedIndexes::Relation(RelationInvertedIndexes::open(
                 work_path, options,
             )?))

nidx/nidx_vector/src/inverted_index.rs

@@ -27,7 +27,6 @@ use super::{
 use crate::{
     ParagraphAddr, VectorR,
     data_store::{DataStore, iter_paragraphs},
-    field_list_metadata::decode_field_list_metadata,
     utils::FieldKey,
 };
 
@@ -62,11 +61,8 @@ impl RelationInvertedIndexes {
 
         for paragraph_addr in iter_paragraphs(data_store) {
             let paragraph = data_store.get_paragraph(paragraph_addr);
-            let fields = decode_field_list_metadata(paragraph.metadata());
-
-            for field_key in fields {
-                field_builder.insert(field_key.bytes().to_vec(), paragraph_addr);
-            }
+            let field_key = FieldKey::from_bytes(paragraph.metadata());
+            field_builder.insert(field_key.bytes().to_vec(), paragraph_addr);
         }
 
         let mut map = InvertedMapWriter::new(&work_path.join(file::INDEX_MAP))?;

nidx/nidx_vector/src/inverted_index/relation.rs

@@ -21,7 +21,7 @@
 pub mod config;
 mod data_store;
 mod data_types;
-mod field_list_metadata;
+
 pub mod formula;
 mod hnsw;
 mod indexer;