nidx-binding 6.3.1.post33__tar.gz → 6.3.1.post46__tar.gz

{nidx_binding-6.3.1.post33 → nidx_binding-6.3.1.post46}/Cargo.lock

@@ -2126,6 +2126,7 @@ dependencies = [
  "lazy_static",
  "nidx_protos",
  "nidx_tantivy",
+ "nidx_tests",
  "nidx_types",
  "tantivy",
  "tempfile",

{nidx_binding-6.3.1.post33 → nidx_binding-6.3.1.post46}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nidx_binding
-Version: 6.3.1.post33
+Version: 6.3.1.post46
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy

{nidx_binding-6.3.1.post33 → nidx_binding-6.3.1.post46}/nidx_protos/pyproject.toml

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

{nidx_binding-6.3.1.post33 → nidx_binding-6.3.1.post46}/nidx_protos/src/nidx.proto

@@ -21,6 +21,7 @@ service NidxApi {
 
 service NidxSearcher {
   rpc Search(nodereader.SearchRequest) returns (nodereader.SearchResponse) {}
+  rpc GraphSearch(nodereader.GraphSearchRequest) returns (nodereader.GraphSearchResponse) {}
   rpc Suggest(nodereader.SuggestRequest) returns (nodereader.SuggestResponse) {}
 
   // Streams

{nidx_binding-6.3.1.post33 → nidx_binding-6.3.1.post46}/nidx_relation/Cargo.toml

@@ -16,3 +16,4 @@ tracing = "0.1.41"
 [dev-dependencies]
 tempfile = "3"
 lazy_static = "1.4.0"
+nidx_tests = { path = "../nidx_tests" }

{nidx_binding-6.3.1.post33 → nidx_binding-6.3.1.post46}/nidx_relation/src/graph_query_parser.rs

@@ -418,3 +418,89 @@ impl From<&str> for Term {
         Self::Exact(value.to_string())
     }
 }
+
+impl TryFrom<&nidx_protos::graph_query::Query> for GraphQuery {
+    type Error = anyhow::Error;
+
+    fn try_from(value: &nidx_protos::graph_query::Query) -> Result<Self, Self::Error> {
+        let query = match value {
+            nidx_protos::graph_query::Query::Node(query_node) => {
+                let node = Node::try_from(query_node)?;
+                Self::NodeQuery(NodeQuery::Node(Expression::Value(node)))
+            }
+
+            nidx_protos::graph_query::Query::Relation(relation) => {
+                Self::RelationQuery(RelationQuery(Expression::Value(Relation::try_from(relation)?)))
+            }
+
+            nidx_protos::graph_query::Query::Path(path) => {
+                let source = match &path.source {
+                    Some(source_pb) => Node::try_from(source_pb)?,
+                    None => Node::default(),
+                };
+                let relation = match &path.relation {
+                    Some(relation_pb) => Relation::try_from(relation_pb)?,
+                    None => Relation::default(),
+                };
+                let destination = match &path.destination {
+                    Some(destination_pb) => Node::try_from(destination_pb)?,
+                    None => Node::default(),
+                };
+
+                if path.undirected {
+                    Self::PathQuery(PathQuery::UndirectedPath((
+                        Expression::Value(source),
+                        Expression::Value(relation),
+                        Expression::Value(destination),
+                    )))
+                } else {
+                    Self::PathQuery(PathQuery::DirectedPath((
+                        Expression::Value(source),
+                        Expression::Value(relation),
+                        Expression::Value(destination),
+                    )))
+                }
+            }
+        };
+
+        Ok(query)
+    }
+}
+
+impl TryFrom<&nidx_protos::graph_query::Node> for Node {
+    type Error = anyhow::Error;
+
+    fn try_from(node_pb: &nidx_protos::graph_query::Node) -> Result<Self, Self::Error> {
+        let value = node_pb.value.clone().map(|value| {
+            if node_pb.fuzzy_distance > 0 || node_pb.as_prefix {
+                Term::Fuzzy(FuzzyTerm {
+                    value: value,
+                    fuzzy_distance: node_pb.fuzzy_distance as u8,
+                    is_prefix: node_pb.as_prefix,
+                })
+            } else {
+                Term::Exact(value)
+            }
+        });
+        let node_type = node_pb.node_type.map(NodeType::try_from).transpose()?;
+        let node_subtype = node_pb.node_subtype.clone();
+
+        Ok(Node {
+            value,
+            node_type,
+            node_subtype,
+        })
+    }
+}
+
+impl TryFrom<&nidx_protos::graph_query::Relation> for Relation {
+    type Error = anyhow::Error;
+
+    fn try_from(relation_pb: &nidx_protos::graph_query::Relation) -> Result<Self, Self::Error> {
+        let value = relation_pb.value.clone();
+
+        Ok(Relation {
+            value,
+        })
+    }
+}

{nidx_binding-6.3.1.post33 → nidx_binding-6.3.1.post46}/nidx_relation/src/io_maps.rs

@@ -80,16 +80,17 @@ where
 
 pub fn source_to_relation_node(schema: &Schema, doc: &TantivyDocument) -> RelationNode {
     RelationNode {
-        value: schema.source_value(doc),
+        value: schema.source_value(doc).to_string(),
         ntype: u64_to_node_type::<i32>(schema.source_type(doc)),
-        subtype: schema.source_subtype(doc),
+        subtype: schema.source_subtype(doc).to_string(),
     }
 }
+
 pub fn target_to_relation_node(schema: &Schema, doc: &TantivyDocument) -> RelationNode {
     RelationNode {
-        value: schema.target_value(doc),
+        value: schema.target_value(doc).to_string(),
         ntype: u64_to_node_type::<i32>(schema.target_type(doc)),
-        subtype: schema.target_subtype(doc),
+        subtype: schema.target_subtype(doc).to_string(),
     }
 }
 
@@ -97,8 +98,16 @@ pub fn doc_to_relation(schema: &Schema, doc: &TantivyDocument) -> ProtosRelation
     ProtosRelation {
         metadata: decode_metadata(schema, doc),
         relation: u64_to_relation_type::<i32>(schema.relationship(doc)),
-        relation_label: schema.relationship_label(doc),
+        relation_label: schema.relationship_label(doc).to_string(),
         source: Some(source_to_relation_node(schema, doc)),
         to: Some(target_to_relation_node(schema, doc)),
     }
 }
+
+pub fn doc_to_graph_relation(schema: &Schema, doc: &TantivyDocument) -> nidx_protos::graph_search_response::Relation {
+    nidx_protos::graph_search_response::Relation {
+        relation_type: u64_to_relation_type::<i32>(schema.relationship(doc)),
+        label: schema.relationship_label(doc).to_string(),
+        metadata: decode_metadata(schema, doc),
+    }
+}

{nidx_binding-6.3.1.post33 → nidx_binding-6.3.1.post46}/nidx_relation/src/lib.rs

@@ -25,8 +25,9 @@ mod resource_indexer;
 mod schema;
 
 use nidx_protos::{
-    relation_node::NodeType, relation_prefix_search_request::Search, resource::ResourceStatus, RelationNode,
-    RelationNodeFilter, RelationPrefixSearchRequest, RelationSearchRequest, RelationSearchResponse,
+    relation_node::NodeType, relation_prefix_search_request::Search, resource::ResourceStatus, GraphSearchRequest,
+    GraphSearchResponse, RelationNode, RelationNodeFilter, RelationPrefixSearchRequest, RelationSearchRequest,
+    RelationSearchResponse,
 };
 use nidx_tantivy::{
     index_reader::{open_index_with_deletions, DeletionQueryBuilder},
@@ -135,6 +136,11 @@ impl RelationSearcher {
         self.reader.relation_search(request)
     }
 
+    #[instrument(name = "relation::graph_searcher", skip_all)]
+    pub fn graph_search(&self, request: &GraphSearchRequest) -> anyhow::Result<GraphSearchResponse> {
+        self.reader.graph_search(request)
+    }
+
     #[instrument(name = "relation::suggest", skip_all)]
     pub fn suggest(&self, prefixes: Vec<String>) -> Vec<RelationNode> {
         let requests =

{nidx_binding-6.3.1.post33 → nidx_binding-6.3.1.post46}/nidx_relation/src/reader.rs

@@ -23,11 +23,12 @@ use std::path::Path;
 
 use nidx_protos::relation_prefix_search_request::Search;
 use nidx_protos::{
-    EntitiesSubgraphResponse, RelationNode, RelationPrefixSearchResponse, RelationSearchRequest, RelationSearchResponse,
+    EntitiesSubgraphResponse, GraphSearchRequest, GraphSearchResponse, RelationNode, RelationPrefixSearchResponse,
+    RelationSearchRequest, RelationSearchResponse,
 };
 use tantivy::collector::TopDocs;
 use tantivy::query::{BooleanQuery, Occur, Query};
-use tantivy::{Index, IndexReader};
+use tantivy::{DocAddress, Index, IndexReader, Searcher};
 
 use crate::graph_query_parser::{
     Expression, FuzzyTerm, GraphQuery, GraphQueryParser, Node, NodeQuery, PathQuery, Relation, Term,
@@ -62,7 +63,72 @@ impl RelationsReaderService {
         })
     }
 
-    pub fn graph_search(&self, query: GraphQuery) -> anyhow::Result<Vec<nidx_protos::Relation>> {
+    pub fn graph_search(&self, request: &GraphSearchRequest) -> anyhow::Result<GraphSearchResponse> {
+        // No query? Empty graph
+        let Some(query) = &request.query else {
+            return Ok(GraphSearchResponse::default());
+        };
+        let Some(query) = &query.query else {
+            return Ok(GraphSearchResponse::default());
+        };
+
+        // Convert proto to tantivy query
+        let query = GraphQuery::try_from(query)?;
+        let parser = GraphQueryParser::new();
+        let index_query = parser.parse(query);
+
+        // Tantivy searcher query
+        let collector = TopDocs::with_limit(request.top_k as usize);
+
+        let searcher = self.reader.searcher();
+        let matching_docs = searcher.search(&index_query, &collector)?;
+
+        self.build_graph_response(&searcher, matching_docs.into_iter().map(|(_score, doc_address)| doc_address))
+    }
+
+    fn build_graph_response(
+        &self,
+        searcher: &Searcher,
+        docs: impl Iterator<Item = DocAddress>,
+    ) -> anyhow::Result<nidx_protos::GraphSearchResponse> {
+        // We are being very naive and writing everything to the proto response. We could be smarter
+        // and deduplicates nodes and relations. As paths are pointers, this would improve proto
+        // size and ser/de time at expenses of deduplication effort.
+
+        let mut nodes = Vec::new();
+        let mut relations = Vec::new();
+        let mut graph = Vec::new();
+
+        for doc_address in docs {
+            let doc = searcher.doc(doc_address)?;
+
+            let source = io_maps::source_to_relation_node(&self.schema, &doc);
+            let relation = io_maps::doc_to_graph_relation(&self.schema, &doc);
+            let destination = io_maps::target_to_relation_node(&self.schema, &doc);
+
+            let source_idx = nodes.len();
+            nodes.push(source);
+            let relation_idx = relations.len();
+            relations.push(relation);
+            let destination_idx = nodes.len();
+            nodes.push(destination);
+
+            graph.push(nidx_protos::graph_search_response::Path {
+                source: source_idx as u32,
+                relation: relation_idx as u32,
+                destination: destination_idx as u32,
+            })
+        }
+
+        let response = nidx_protos::GraphSearchResponse {
+            nodes,
+            relations,
+            graph,
+        };
+        Ok(response)
+    }
+
+    pub fn inner_graph_search(&self, query: GraphQuery) -> anyhow::Result<nidx_protos::GraphSearchResponse> {
         let parser = GraphQueryParser::new();
         let index_query: Box<dyn Query> = parser.parse(query);
 
@@ -71,13 +137,8 @@ impl RelationsReaderService {
 
         let searcher = self.reader.searcher();
         let matching_docs = searcher.search(&index_query, &collector)?;
-        let mut relations = Vec::with_capacity(matching_docs.len());
-        for (_, doc_addr) in matching_docs {
-            let doc = searcher.doc(doc_addr)?;
-            let relation = io_maps::doc_to_relation(&self.schema, &doc);
-            relations.push(relation);
-        }
-        Ok(relations)
+
+        self.build_graph_response(&searcher, matching_docs.into_iter().map(|(_score, doc_address)| doc_address))
     }
 }
 
@@ -238,7 +299,11 @@ impl RelationsReaderService {
         }
 
         let searcher = self.reader.searcher();
-        let topdocs = TopDocs::with_limit(NUMBER_OF_RESULTS_SUGGEST);
+        // FIXME: we are using a topdocs collector to get prefix results from source and target
+        // nodes. However, we are deduplicating afterwards, and this means we could end up with less
+        // results although results may exist. As a quick fix, we increase here the limit of the
+        // collector. The proper solution would be implementing a custom collector for this task
+        let topdocs = TopDocs::with_limit(NUMBER_OF_RESULTS_SUGGEST * 2);
         let schema = &self.schema;
 
         match search {
@@ -321,16 +386,17 @@ impl RelationsReaderService {
             let relation_node = io_maps::source_to_relation_node(schema, &source_res_doc);
             results.insert(HashedRelationNode(relation_node));
         }
-        for (_, source_res_addr) in searcher.search(&target_prefix_query, &topdocs)? {
-            let source_res_doc = searcher.doc(source_res_addr)?;
-            let relation_node = io_maps::target_to_relation_node(schema, &source_res_doc);
+        for (_, target_res_addr) in searcher.search(&target_prefix_query, &topdocs)? {
+            let target_res_doc = searcher.doc(target_res_addr)?;
+            let relation_node = io_maps::target_to_relation_node(schema, &target_res_doc);
             results.insert(HashedRelationNode(relation_node));
         }
-        response.nodes = results.into_iter().map(Into::into).collect();
+        response.nodes = results.into_iter().take(NUMBER_OF_RESULTS_SUGGEST).map(Into::into).collect();
         Ok(Some(response))
     }
 }
 
+#[derive(Debug)]
 pub struct HashedRelationNode(pub RelationNode);
 
 impl From<HashedRelationNode> for RelationNode {

{nidx_binding-6.3.1.post33 → nidx_binding-6.3.1.post46}/nidx_relation/src/schema.rs

@@ -101,51 +101,36 @@ impl Schema {
             .expect("Documents must have a resource id")
     }
 
-    pub fn source_value(&self, doc: &TantivyDocument) -> String {
-        doc.get_first(self.source_value)
-            .and_then(|i| i.as_str())
-            .map(String::from)
-            .expect("Documents must have a source value")
+    pub fn source_value<'a>(&self, doc: &'a TantivyDocument) -> &'a str {
+        doc.get_first(self.source_value).and_then(|i| i.as_str()).expect("Documents must have a source value")
     }
 
     pub fn source_type(&self, doc: &TantivyDocument) -> u64 {
         doc.get_first(self.source_type).and_then(|i| i.as_u64()).expect("Documents must have a source type")
     }
 
-    pub fn source_subtype(&self, doc: &TantivyDocument) -> String {
-        doc.get_first(self.source_subtype)
-            .and_then(|i| i.as_str())
-            .map(String::from)
-            .expect("Documents must have a source subtype")
+    pub fn source_subtype<'a>(&self, doc: &'a TantivyDocument) -> &'a str {
+        doc.get_first(self.source_subtype).and_then(|i| i.as_str()).expect("Documents must have a source subtype")
     }
 
-    pub fn target_value(&self, doc: &TantivyDocument) -> String {
-        doc.get_first(self.target_value)
-            .and_then(|i| i.as_str())
-            .map(String::from)
-            .expect("Documents must have a target value")
+    pub fn target_value<'a>(&self, doc: &'a TantivyDocument) -> &'a str {
+        doc.get_first(self.target_value).and_then(|i| i.as_str()).expect("Documents must have a target value")
     }
 
     pub fn target_type(&self, doc: &TantivyDocument) -> u64 {
         doc.get_first(self.target_type).and_then(|i| i.as_u64()).expect("Documents must have a target type")
     }
 
-    pub fn target_subtype(&self, doc: &TantivyDocument) -> String {
-        doc.get_first(self.target_subtype)
-            .and_then(|i| i.as_str())
-            .map(String::from)
-            .expect("Documents must have a target subtype")
+    pub fn target_subtype<'a>(&self, doc: &'a TantivyDocument) -> &'a str {
+        doc.get_first(self.target_subtype).and_then(|i| i.as_str()).expect("Documents must have a target subtype")
     }
 
     pub fn relationship(&self, doc: &TantivyDocument) -> u64 {
         doc.get_first(self.relationship).and_then(|i| i.as_u64()).expect("Documents must have a relationship type")
     }
 
-    pub fn relationship_label(&self, doc: &TantivyDocument) -> String {
-        doc.get_first(self.label)
-            .and_then(|i| i.as_str())
-            .map(String::from)
-            .expect("Documents must have a relationship label")
+    pub fn relationship_label<'a>(&self, doc: &'a TantivyDocument) -> &'a str {
+        doc.get_first(self.label).and_then(|i| i.as_str()).expect("Documents must have a relationship label")
     }
 
     pub fn metadata<'a>(&self, doc: &'a TantivyDocument) -> Option<&'a [u8]> {