nidx-binding 6.3.4.post141__tar.gz → 6.3.4.post146__tar.gz

{nidx_binding-6.3.4.post141 → nidx_binding-6.3.4.post146}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nidx_binding
-Version: 6.3.4.post141
+Version: 6.3.4.post146
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy

{nidx_binding-6.3.4.post141 → nidx_binding-6.3.4.post146}/nidx_protos/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "nidx_protos"
-version = "6.3.4.post141"
+version = "6.3.4.post146"
 license = { text = "AGPL" }
 description = "Protobuf definitions for nucliadb/nidx"
 authors = [ { name = "Nuclia", email = "nucliadb@nuclia.com" }]

{nidx_binding-6.3.4.post141 → nidx_binding-6.3.4.post146}/nidx_relation/src/graph_collector.rs

@@ -23,6 +23,7 @@ use std::collections::{HashMap, HashSet};
 use tantivy::{
     DocId, Score, SegmentOrdinal, SegmentReader, TantivyDocument,
     collector::{Collector, SegmentCollector},
+    columnar::Column,
     store::StoreReader,
 };
 use tracing::warn;
@@ -41,23 +42,42 @@ pub enum NodeSelector {
     DestinationNodes,
 }
 
+// Node collector for schema v1
+
 pub struct TopUniqueNodeCollector {
     limit: usize,
     selector: NodeSelector,
     schema: crate::schema::Schema,
 }
 
-pub struct TopUniqueRelationCollector {
+pub struct TopUniqueNodeSegmentCollector {
     limit: usize,
+    selector: NodeSelector,
+    unique: HashSet<NodeId>,
     schema: crate::schema::Schema,
+    store_reader: StoreReader,
 }
 
-pub struct TopUniqueNodeSegmentCollector {
+// Node collector for schema v2
+//
+// We can now use fast fields to uniquely identify nodes.
+
+pub struct TopUniqueNodeCollector2 {
     limit: usize,
     selector: NodeSelector,
-    unique: HashSet<NodeId>,
+}
+
+pub struct TopUniqueNodeSegmentCollector2 {
+    limit: usize,
+    unique: HashSet<Vec<u64>>,
+    encoded_node_reader: Column<u64>,
+}
+
+// Relations collector
+
+pub struct TopUniqueRelationCollector {
+    limit: usize,
     schema: crate::schema::Schema,
-    store_reader: StoreReader,
 }
 
 pub struct TopUniqueRelationSegmentCollector {
@@ -67,6 +87,18 @@ pub struct TopUniqueRelationSegmentCollector {
     store_reader: StoreReader,
 }
 
+// Relations collector for schema v2
+
+pub struct TopUniqueRelationCollector2 {
+    limit: usize,
+}
+
+pub struct TopUniqueRelationSegmentCollector2 {
+    limit: usize,
+    unique: HashSet<Vec<u64>>,
+    encoded_relation_reader: Column<u64>,
+}
+
 impl TopUniqueNodeCollector {
     pub fn new(schema: crate::schema::Schema, selector: NodeSelector, limit: usize) -> Self {
         Self {
@@ -77,12 +109,6 @@ impl TopUniqueNodeCollector {
     }
 }
 
-impl TopUniqueRelationCollector {
-    pub fn new(schema: crate::schema::Schema, limit: usize) -> Self {
-        Self { limit, schema }
-    }
-}
-
 impl Collector for TopUniqueNodeCollector {
     type Fruit = HashSet<NodeId>;
     type Child = TopUniqueNodeSegmentCollector;
@@ -117,6 +143,111 @@ impl Collector for TopUniqueNodeCollector {
     }
 }
 
+impl SegmentCollector for TopUniqueNodeSegmentCollector {
+    type Fruit = HashSet<NodeId>;
+
+    fn collect(&mut self, doc_id: DocId, _score: Score) {
+        // we already have all unique results we need
+        if self.unique.len() >= self.limit {
+            return;
+        }
+
+        // log and skip documents not in the store. This should not happen
+        let doc = match self.store_reader.get::<TantivyDocument>(doc_id) {
+            Ok(doc) => doc,
+            Err(error) => {
+                warn!("Error while getting document from store: {error:?}");
+                return;
+            }
+        };
+
+        let node = match self.selector {
+            NodeSelector::SourceNodes => {
+                let source_value = self.schema.source_value(&doc).to_string();
+                let source_type = io_maps::u64_to_node_type::<i32>(self.schema.source_type(&doc));
+                let source_subtype = self.schema.source_subtype(&doc).to_string();
+                (source_value, source_type, source_subtype)
+            }
+            NodeSelector::DestinationNodes => {
+                let destination_value = self.schema.target_value(&doc).to_string();
+                let destination_type = io_maps::u64_to_node_type::<i32>(self.schema.target_type(&doc));
+                let destination_subtype = self.schema.target_subtype(&doc).to_string();
+                (destination_value, destination_type, destination_subtype)
+            }
+        };
+        self.unique.insert(node);
+    }
+
+    fn harvest(self) -> Self::Fruit {
+        self.unique
+    }
+}
+
+impl TopUniqueNodeCollector2 {
+    pub fn new(selector: NodeSelector, limit: usize) -> Self {
+        Self { limit, selector }
+    }
+}
+
+impl Collector for TopUniqueNodeCollector2 {
+    type Fruit = HashSet<Vec<u64>>;
+    type Child = TopUniqueNodeSegmentCollector2;
+
+    fn requires_scoring(&self) -> bool {
+        false
+    }
+
+    fn for_segment(&self, _segment_local_id: SegmentOrdinal, segment: &SegmentReader) -> tantivy::Result<Self::Child> {
+        let fast_field_reader = match self.selector {
+            NodeSelector::SourceNodes => segment.fast_fields().u64("encoded_source_id")?,
+            NodeSelector::DestinationNodes => segment.fast_fields().u64("encoded_target_id")?,
+        };
+        Ok(TopUniqueNodeSegmentCollector2 {
+            limit: self.limit,
+            unique: HashSet::new(),
+            encoded_node_reader: fast_field_reader,
+        })
+    }
+
+    fn merge_fruits(
+        &self,
+        segment_fruits: Vec<<Self::Child as SegmentCollector>::Fruit>,
+    ) -> tantivy::Result<Self::Fruit> {
+        let mut unique = HashSet::new();
+        let mut fruits = segment_fruits.into_iter().flatten();
+        let mut fruit = fruits.next();
+
+        while fruit.is_some() && unique.len() < self.limit {
+            unique.insert(fruit.unwrap());
+            fruit = fruits.next();
+        }
+        Ok(unique)
+    }
+}
+
+impl SegmentCollector for TopUniqueNodeSegmentCollector2 {
+    type Fruit = HashSet<Vec<u64>>;
+
+    fn collect(&mut self, doc_id: DocId, _score: Score) {
+        // we already have all unique results we need
+        if self.unique.len() >= self.limit {
+            return;
+        }
+        let encoded_node = self.encoded_node_reader.values_for_doc(doc_id).collect::<Vec<u64>>();
+        self.unique.insert(encoded_node);
+    }
+
+    fn harvest(self) -> Self::Fruit {
+        self.unique
+    }
+}
+
+impl TopUniqueRelationCollector {
+    pub fn new(schema: crate::schema::Schema, limit: usize) -> Self {
+        Self { limit, schema }
+    }
+}
+
 impl Collector for TopUniqueRelationCollector {
     type Fruit = Vec<TantivyDocument>;
     type Child = TopUniqueRelationSegmentCollector;
@@ -154,8 +285,8 @@ impl Collector for TopUniqueRelationCollector {
     }
 }
 
-impl SegmentCollector for TopUniqueNodeSegmentCollector {
-    type Fruit = HashSet<NodeId>;
+impl SegmentCollector for TopUniqueRelationSegmentCollector {
+    type Fruit = HashMap<RelationId, TantivyDocument>;
 
     fn collect(&mut self, doc_id: DocId, _score: Score) {
         // we already have all unique results we need
@@ -172,21 +303,9 @@ impl SegmentCollector for TopUniqueNodeSegmentCollector {
             }
         };
 
-        let node = match self.selector {
-            NodeSelector::SourceNodes => {
-                let source_value = self.schema.source_value(&doc).to_string();
-                let source_type = io_maps::u64_to_node_type::<i32>(self.schema.source_type(&doc));
-                let source_subtype = self.schema.source_subtype(&doc).to_string();
-                (source_value, source_type, source_subtype)
-            }
-            NodeSelector::DestinationNodes => {
-                let destination_value = self.schema.target_value(&doc).to_string();
-                let destination_type = io_maps::u64_to_node_type::<i32>(self.schema.target_type(&doc));
-                let destination_subtype = self.schema.target_subtype(&doc).to_string();
-                (destination_value, destination_type, destination_subtype)
-            }
-        };
-        self.unique.insert(node);
+        let relation_label = self.schema.relationship_label(&doc).to_string();
+        let relation_type = self.schema.relationship(&doc);
+        self.unique.insert((relation_label, relation_type), doc);
     }
 
     fn harvest(self) -> Self::Fruit {
@@ -194,8 +313,46 @@ impl SegmentCollector for TopUniqueNodeSegmentCollector {
     }
 }
 
-impl SegmentCollector for TopUniqueRelationSegmentCollector {
-    type Fruit = HashMap<RelationId, TantivyDocument>;
+impl TopUniqueRelationCollector2 {
+    pub fn new(limit: usize) -> Self {
+        Self { limit }
+    }
+}
+
+impl Collector for TopUniqueRelationCollector2 {
+    type Fruit = HashSet<Vec<u64>>;
+    type Child = TopUniqueRelationSegmentCollector2;
+
+    fn requires_scoring(&self) -> bool {
+        false
+    }
+
+    fn for_segment(&self, _segment_local_id: SegmentOrdinal, segment: &SegmentReader) -> tantivy::Result<Self::Child> {
+        Ok(TopUniqueRelationSegmentCollector2 {
+            limit: self.limit,
+            unique: HashSet::new(),
+            encoded_relation_reader: segment.fast_fields().u64("encoded_relation_id")?,
+        })
+    }
+
+    fn merge_fruits(
+        &self,
+        segment_fruits: Vec<<Self::Child as SegmentCollector>::Fruit>,
+    ) -> tantivy::Result<Self::Fruit> {
+        let mut unique = HashSet::new();
+        let mut fruits = segment_fruits.into_iter().flat_map(|map| map.into_iter());
+        let mut fruit = fruits.next();
+
+        while fruit.is_some() && unique.len() < self.limit {
+            unique.insert(fruit.unwrap());
+            fruit = fruits.next();
+        }
+        Ok(unique)
+    }
+}
+
+impl SegmentCollector for TopUniqueRelationSegmentCollector2 {
+    type Fruit = HashSet<Vec<u64>>;
 
     fn collect(&mut self, doc_id: DocId, _score: Score) {
         // we already have all unique results we need
@@ -203,18 +360,11 @@ impl SegmentCollector for TopUniqueRelationSegmentCollector {
             return;
         }
 
-        // log and skip documents not in the store. This should not happen
-        let doc = match self.store_reader.get::<TantivyDocument>(doc_id) {
-            Ok(doc) => doc,
-            Err(error) => {
-                warn!("Error while getting document from store: {error:?}");
-                return;
-            }
-        };
-
-        let relation_label = self.schema.relationship_label(&doc).to_string();
-        let relation_type = self.schema.relationship(&doc);
-        self.unique.insert((relation_label, relation_type), doc);
+        let relation = self
+            .encoded_relation_reader
+            .values_for_doc(doc_id)
+            .collect::<Vec<u64>>();
+        self.unique.insert(relation);
     }
 
     fn harvest(self) -> Self::Fruit {

{nidx_binding-6.3.4.post141 → nidx_binding-6.3.4.post146}/nidx_relation/src/reader.rs

@@ -34,11 +34,14 @@ use tantivy::schema::Field;
 use tantivy::{DocAddress, Index, IndexReader, Searcher};
 use uuid::Uuid;
 
-use crate::graph_collector::{NodeSelector, TopUniqueNodeCollector, TopUniqueRelationCollector};
+use crate::graph_collector::{
+    NodeSelector, TopUniqueNodeCollector, TopUniqueNodeCollector2, TopUniqueRelationCollector,
+    TopUniqueRelationCollector2,
+};
 use crate::graph_query_parser::{
     BoolGraphQuery, BoolNodeQuery, Expression, FuzzyTerm, GraphQuery, GraphQueryParser, Node, NodeQuery, Term,
 };
-use crate::schema::{Schema, encode_field_id};
+use crate::schema::{Schema, decode_node, decode_relation, encode_field_id};
 use crate::{RelationConfig, io_maps};
 
 const FUZZY_DISTANCE: u8 = 1;
@@ -166,33 +169,62 @@ impl RelationsReaderService {
         let source_query = self.apply_prefilter(source_query, prefilter);
         let destination_query = self.apply_prefilter(destination_query, prefilter);
 
-        let mut unique_nodes = HashSet::new();
-
-        let collector = TopUniqueNodeCollector::new(self.schema.clone(), NodeSelector::SourceNodes, top_k);
-        let searcher = self.reader.searcher();
-        let mut source_nodes = searcher.search(&source_query, &collector)?;
-        unique_nodes.extend(source_nodes.drain());
-
-        let collector = TopUniqueNodeCollector::new(self.schema.clone(), NodeSelector::DestinationNodes, top_k);
         let searcher = self.reader.searcher();
-        let mut destination_nodes = searcher.search(&destination_query, &collector)?;
-        unique_nodes.extend(destination_nodes.drain());
-
-        let nodes = unique_nodes
-            .into_iter()
-            .map(|(value, node_type, node_subtype)| RelationNode {
-                value,
-                ntype: node_type,
-                subtype: node_subtype,
-            })
-            .take(top_k)
-            .collect();
 
-        let response = nidx_protos::GraphSearchResponse {
-            nodes,
-            ..Default::default()
-        };
-        Ok(response)
+        if self.schema.version == 1 {
+            let mut unique_nodes = HashSet::new();
+
+            let collector = TopUniqueNodeCollector::new(self.schema.clone(), NodeSelector::SourceNodes, top_k);
+            let mut source_nodes = searcher.search(&source_query, &collector)?;
+            unique_nodes.extend(source_nodes.drain());
+
+            let collector = TopUniqueNodeCollector::new(self.schema.clone(), NodeSelector::DestinationNodes, top_k);
+            let mut destination_nodes = searcher.search(&destination_query, &collector)?;
+            unique_nodes.extend(destination_nodes.drain());
+
+            let nodes = unique_nodes
+                .into_iter()
+                .map(|(value, node_type, node_subtype)| RelationNode {
+                    value,
+                    ntype: node_type,
+                    subtype: node_subtype,
+                })
+                .take(top_k)
+                .collect();
+
+            let response = nidx_protos::GraphSearchResponse {
+                nodes,
+                ..Default::default()
+            };
+            Ok(response)
+        } else {
+            let mut unique_nodes = HashSet::new();
+
+            let collector = TopUniqueNodeCollector2::new(NodeSelector::SourceNodes, top_k);
+            let mut source_nodes = searcher.search(&source_query, &collector)?;
+            unique_nodes.extend(source_nodes.drain());
+
+            let collector = TopUniqueNodeCollector2::new(NodeSelector::DestinationNodes, top_k);
+            let mut destination_nodes = searcher.search(&destination_query, &collector)?;
+            unique_nodes.extend(destination_nodes.drain());
+
+            let nodes = unique_nodes
+                .into_iter()
+                .map(|encoded_node| decode_node(&encoded_node))
+                .map(|(value, node_type, node_subtype)| RelationNode {
+                    value,
+                    ntype: io_maps::u64_to_node_type(node_type),
+                    subtype: node_subtype,
+                })
+                .take(top_k)
+                .collect();
+
+            let response = nidx_protos::GraphSearchResponse {
+                nodes,
+                ..Default::default()
+            };
+            Ok(response)
+        }
     }
 
     fn relations_graph_search(
@@ -206,20 +238,48 @@ impl RelationsReaderService {
         let index_query = parser.parse_bool(query);
         let index_query = self.apply_prefilter(index_query, prefilter);
 
-        let collector = TopUniqueRelationCollector::new(self.schema.clone(), top_k);
         let searcher = self.reader.searcher();
-        let matching_docs = searcher.search(&index_query, &collector)?;
 
-        let relations = matching_docs
-            .into_iter()
-            .map(|doc| io_maps::doc_to_graph_relation(&self.schema, &doc))
-            .collect();
+        if self.schema.version == 1 {
+            let collector = TopUniqueRelationCollector::new(self.schema.clone(), top_k);
+            let matching_docs = searcher.search(&index_query, &collector)?;
+
+            let relations = matching_docs
+                .into_iter()
+                .map(|doc| nidx_protos::graph_search_response::Relation {
+                    relation_type: io_maps::u64_to_relation_type::<i32>(self.schema.relationship(&doc)),
+                    label: self.schema.relationship_label(&doc).to_string(),
+                    metadata: None,
+                })
+                .collect();
+
+            let response = nidx_protos::GraphSearchResponse {
+                relations,
+                ..Default::default()
+            };
+            Ok(response)
+        } else {
+            let collector = TopUniqueRelationCollector2::new(top_k);
+            let matching_docs = searcher.search(&index_query, &collector)?;
+
+            let relations = matching_docs
+                .into_iter()
+                .map(|encoded_relation| {
+                    let (relation_type, relation_label) = decode_relation(&encoded_relation);
+                    nidx_protos::graph_search_response::Relation {
+                        relation_type: io_maps::u64_to_relation_type::<i32>(relation_type),
+                        label: relation_label,
+                        metadata: None,
+                    }
+                })
+                .collect();
 
-        let response = nidx_protos::GraphSearchResponse {
-            relations,
-            ..Default::default()
-        };
-        Ok(response)
+            let response = nidx_protos::GraphSearchResponse {
+                relations,
+                ..Default::default()
+            };
+            Ok(response)
+        }
     }
 
     fn apply_prefilter(&self, query: Box<dyn Query>, prefilter: &PrefilterResult) -> Box<dyn Query> {

{nidx_binding-6.3.4.post141 → nidx_binding-6.3.4.post146}/nidx_relation/src/resource_indexer.rs

@@ -19,7 +19,7 @@
 //
 
 use crate::io_maps;
-use crate::schema::{Schema, encode_field_id};
+use crate::schema::{Schema, encode_field_id, encode_node, encode_relation};
 use anyhow::anyhow;
 use nidx_protos::noderesources::IndexRelation;
 use nidx_protos::prost::*;
@@ -99,12 +99,24 @@ pub fn index_relations(
             let field = field_key.ok_or(anyhow!("Field ID required for v2"))?;
             new_doc.add_bytes(schema.resource_id, rid.as_bytes());
             new_doc.add_bytes(schema.resource_field_id.unwrap(), encode_field_id(rid, field));
+
             for facet in &index_relation.facets {
                 new_doc.add_facet(schema.facets.unwrap(), Facet::from_text(facet)?);
             }
-            // encoded_source_id = Some(builder.add_u64_field("encoded_source_id", FAST));
-            // encoded_target_id = Some(builder.add_u64_field("encoded_target_id", FAST));
-            // todo!();
+
+            // Encode source and target nodes and relation for faster retrieval
+            let encoded_source_id = schema.encoded_source_id.unwrap();
+            for b in encode_node(source_value, source_type, source_subtype) {
+                new_doc.add_u64(encoded_source_id, b);
+            }
+            let encoded_target_id = schema.encoded_target_id.unwrap();
+            for b in encode_node(target_value, target_type, target_subtype) {
+                new_doc.add_u64(encoded_target_id, b);
+            }
+            let encoded_relation_id = schema.encoded_relation_id.unwrap();
+            for b in encode_relation(relationship, label) {
+                new_doc.add_u64(encoded_relation_id, b);
+            }
         }
 
         if let Some(metadata) = relation.metadata.as_ref() {