nidx-binding 6.3.3.post92__tar.gz → 6.3.3.post96__tar.gz

{nidx_binding-6.3.3.post92 → nidx_binding-6.3.3.post96}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nidx_binding
-Version: 6.3.3.post92
+Version: 6.3.3.post96
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy

{nidx_binding-6.3.3.post92 → nidx_binding-6.3.3.post96}/nidx_protos/pyproject.toml

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

nidx_binding-6.3.3.post96/nidx_relation/src/graph_collector.rs

@@ -0,0 +1,226 @@
+// Copyright (C) 2021 Bosutech XXI S.L.
+//
+// nucliadb is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at info@nuclia.com.
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see <http://www.gnu.org/licenses/>.
+//
+
+use std::collections::{HashMap, HashSet};
+
+use tantivy::{
+    DocId, Score, SegmentOrdinal, SegmentReader, TantivyDocument,
+    collector::{Collector, SegmentCollector},
+    store::StoreReader,
+};
+use tracing::warn;
+
+use crate::io_maps;
+
+// XXX: this is a completely arbitrary number. Feel free to change for a more adequate value
+const STORE_READER_CACHED_BLOCKS: usize = 10;
+
+pub type NodeId = (String, i32, String);
+pub type RelationId = (String, u64);
+
+#[derive(Clone, Copy)]
+pub enum NodeSelector {
+    SourceNodes,
+    DestinationNodes,
+}
+
+pub struct TopUniqueNodeCollector {
+    limit: usize,
+    selector: NodeSelector,
+    schema: crate::schema::Schema,
+}
+
+pub struct TopUniqueRelationCollector {
+    limit: usize,
+    schema: crate::schema::Schema,
+}
+
+pub struct TopUniqueNodeSegmentCollector {
+    limit: usize,
+    selector: NodeSelector,
+    unique: HashSet<NodeId>,
+    schema: crate::schema::Schema,
+    store_reader: StoreReader,
+}
+
+pub struct TopUniqueRelationSegmentCollector {
+    limit: usize,
+    unique: HashMap<RelationId, TantivyDocument>,
+    schema: crate::schema::Schema,
+    store_reader: StoreReader,
+}
+
+impl TopUniqueNodeCollector {
+    pub fn new(selector: NodeSelector, limit: usize) -> Self {
+        Self {
+            limit,
+            selector,
+            schema: crate::schema::Schema::new(),
+        }
+    }
+}
+
+impl TopUniqueRelationCollector {
+    pub fn new(limit: usize) -> Self {
+        Self {
+            limit,
+            schema: crate::schema::Schema::new(),
+        }
+    }
+}
+
+impl Collector for TopUniqueNodeCollector {
+    type Fruit = HashSet<NodeId>;
+    type Child = TopUniqueNodeSegmentCollector;
+
+    fn requires_scoring(&self) -> bool {
+        false
+    }
+
+    fn for_segment(&self, _segment_local_id: SegmentOrdinal, segment: &SegmentReader) -> tantivy::Result<Self::Child> {
+        Ok(TopUniqueNodeSegmentCollector {
+            limit: self.limit,
+            selector: self.selector,
+            unique: HashSet::new(),
+            store_reader: segment.get_store_reader(STORE_READER_CACHED_BLOCKS)?,
+            schema: self.schema.clone(),
+        })
+    }
+
+    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 Collector for TopUniqueRelationCollector {
+    type Fruit = Vec<TantivyDocument>;
+    type Child = TopUniqueRelationSegmentCollector;
+
+    fn requires_scoring(&self) -> bool {
+        false
+    }
+
+    fn for_segment(&self, _segment_local_id: SegmentOrdinal, segment: &SegmentReader) -> tantivy::Result<Self::Child> {
+        Ok(TopUniqueRelationSegmentCollector {
+            limit: self.limit,
+            unique: HashMap::new(),
+            store_reader: segment.get_store_reader(STORE_READER_CACHED_BLOCKS)?,
+            schema: self.schema.clone(),
+        })
+    }
+
+    fn merge_fruits(
+        &self,
+        segment_fruits: Vec<<Self::Child as SegmentCollector>::Fruit>,
+    ) -> tantivy::Result<Self::Fruit> {
+        let mut unique = HashSet::new();
+        let mut docs = Vec::new();
+        let mut fruits = segment_fruits.into_iter().flatten();
+        let mut fruit = fruits.next();
+
+        while fruit.is_some() && unique.len() < self.limit {
+            let (relation_id, doc) = fruit.unwrap();
+            if unique.insert(relation_id) {
+                docs.push(doc);
+            }
+            fruit = fruits.next();
+        }
+        Ok(docs)
+    }
+}
+
+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 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
+        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 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 {
+        self.unique
+    }
+}

{nidx_binding-6.3.3.post92 → nidx_binding-6.3.3.post96}/nidx_relation/src/graph_query_parser.rs

@@ -17,6 +17,7 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program. If not, see <http://www.gnu.org/licenses/>.
 //
+use anyhow::anyhow;
 use nidx_protos::relation_node::NodeType;
 use nidx_types::query_language::{BooleanExpression, BooleanOperation, Operator};
 use tantivy::query::{AllQuery, BooleanQuery, FuzzyTermQuery, Occur, Query, TermQuery};
@@ -107,6 +108,13 @@ pub enum GraphQuery {
     PathQuery(PathQuery),
 }
 
+#[derive(Clone, Copy)]
+enum NodePosition {
+    Source,
+    Destination,
+}
+
+pub struct BoolNodeQuery(BooleanExpression<Node>);
 pub struct BoolGraphQuery(BooleanExpression<GraphQuery>);
 
 #[derive(Clone, Copy)]
@@ -152,6 +160,43 @@ impl GraphQueryParser {
         }
     }
 
+    pub fn parse_bool_node(&self, query: BoolNodeQuery) -> (Box<dyn Query>, Box<dyn Query>) {
+        (
+            self.inner_parse_bool_node(query.0.clone(), NodePosition::Source),
+            self.inner_parse_bool_node(query.0, NodePosition::Destination),
+        )
+    }
+
+    fn inner_parse_bool_node(&self, query: BooleanExpression<Node>, position: NodePosition) -> Box<dyn Query> {
+        match query {
+            BooleanExpression::Literal(node) => match position {
+                NodePosition::Source => {
+                    self.parse(GraphQuery::NodeQuery(NodeQuery::SourceNode(Expression::Value(node))))
+                }
+                NodePosition::Destination => self.parse(GraphQuery::NodeQuery(NodeQuery::DestinationNode(
+                    Expression::Value(node),
+                ))),
+            },
+            BooleanExpression::Not(subquery) => {
+                let mut subqueries = vec![];
+                subqueries.push((Occur::Must, Box::new(AllQuery) as Box<dyn Query>));
+                subqueries.push((Occur::MustNot, self.inner_parse_bool_node(*subquery, position)));
+                Box::new(BooleanQuery::new(subqueries))
+            }
+            BooleanExpression::Operation(operation) => {
+                let mut subqueries = vec![];
+                let occur = match operation.operator {
+                    Operator::And => Occur::Must,
+                    Operator::Or => Occur::Should,
+                };
+                for expression in operation.operands {
+                    subqueries.push((occur, self.inner_parse_bool_node(expression, position)));
+                }
+                Box::new(BooleanQuery::new(subqueries))
+            }
+        }
+    }
+
     pub fn parse(&self, query: GraphQuery) -> Box<dyn Query> {
         // REVIEW: if at some point we only want to return what the query really asks (nodes,
         // relations or paths), we may want to return the query and some kind of response builder
@@ -200,7 +245,7 @@ impl GraphQueryParser {
                 let mut subqueries = vec![];
 
                 subqueries.extend(self.has_node_expression_as_source(&source_expression));
-                subqueries.extend(self.has_relation(relation_expression).into_iter());
+                subqueries.extend(self.has_relation(&relation_expression).into_iter());
                 subqueries.extend(self.has_node_expression_as_destination(&destination_expression));
 
                 // Due to implementation details on tantivy, a query containing only MustNot won't
@@ -344,29 +389,29 @@ impl GraphQueryParser {
     }
 
     // Return a list of queries needed to match a triplet containing such relation expression
-    fn has_relation(&self, expression: Expression<Relation>) -> Vec<(Occur, Box<dyn Query>)> {
+    fn has_relation(&self, expression: &Expression<Relation>) -> Vec<(Occur, Box<dyn Query>)> {
         let mut subqueries = vec![];
 
         match expression {
             Expression::Value(Relation { value: None }) => {}
             Expression::Value(Relation { value: Some(value) }) if value.is_empty() => {}
             Expression::Value(Relation { value: Some(value) }) => {
-                subqueries.push((Occur::Must, self.has_relation_label(&value)));
+                subqueries.push((Occur::Must, self.has_relation_label(value)));
             }
 
             Expression::Not(Relation { value: None }) => {}
             Expression::Not(Relation { value: Some(value) }) if value.is_empty() => {}
             Expression::Not(Relation { value: Some(value) }) => {
-                subqueries.push((Occur::MustNot, self.has_relation_label(&value)));
+                subqueries.push((Occur::MustNot, self.has_relation_label(value)));
             }
 
             Expression::Or(relations) => {
                 subqueries.extend(relations.into_iter().flat_map(|relation| {
-                    if let Some(label) = relation.value {
+                    if let Some(label) = &relation.value {
                         if label.is_empty() {
                             None
                         } else {
-                            Some((Occur::Should, self.has_relation_label(&label)))
+                            Some((Occur::Should, self.has_relation_label(label)))
                         }
                     } else {
                         None
@@ -456,6 +501,64 @@ impl From<&str> for Term {
     }
 }
 
+impl TryFrom<&nidx_protos::graph_query::PathQuery> for BoolNodeQuery {
+    type Error = anyhow::Error;
+
+    fn try_from(value: &nidx_protos::graph_query::PathQuery) -> Result<Self, Self::Error> {
+        let bool_expr = match &value.query {
+            Some(query) => match query {
+                nidx_protos::graph_query::path_query::Query::Path(path) => {
+                    if !(path.source.is_some()
+                        && path.relation.is_none()
+                        && path.destination.is_none()
+                        && path.undirected == true)
+                    {
+                        // We are doing something wrong between search API and nidx
+                        return Err(anyhow!(
+                            "Invalid node query, we only expect a source for an undirected path! {path:?}"
+                        ));
+                    }
+
+                    let pb_node = path.source.as_ref().unwrap();
+                    let node = Node::try_from(pb_node)?;
+                    BooleanExpression::Literal(node)
+                }
+
+                nidx_protos::graph_query::path_query::Query::BoolNot(bool_not) => {
+                    let subquery = BoolNodeQuery::try_from(bool_not.as_ref())?.0;
+                    BooleanExpression::Not(Box::new(subquery))
+                }
+
+                nidx_protos::graph_query::path_query::Query::BoolAnd(bool_and) => {
+                    BooleanExpression::Operation(BooleanOperation {
+                        operator: Operator::And,
+                        operands: bool_and
+                            .operands
+                            .iter()
+                            .map(|o| BoolNodeQuery::try_from(o).map(|x| x.0))
+                            .collect::<anyhow::Result<Vec<_>>>()?,
+                    })
+                }
+
+                nidx_protos::graph_query::path_query::Query::BoolOr(bool_or) => {
+                    BooleanExpression::Operation(BooleanOperation {
+                        operator: Operator::Or,
+                        operands: bool_or
+                            .operands
+                            .iter()
+                            .map(|o| BoolNodeQuery::try_from(o).map(|x| x.0))
+                            .collect::<anyhow::Result<Vec<_>>>()?,
+                    })
+                }
+            },
+
+            None => BooleanExpression::Literal(Node::default()),
+        };
+
+        Ok(BoolNodeQuery(bool_expr))
+    }
+}
+
 impl TryFrom<&nidx_protos::graph_query::PathQuery> for BoolGraphQuery {
     type Error = anyhow::Error;
 

{nidx_binding-6.3.3.post92 → nidx_binding-6.3.3.post96}/nidx_relation/src/lib.rs

@@ -18,6 +18,7 @@
 // along with this program. If not, see <http://www.gnu.org/licenses/>.
 //
 
+mod graph_collector;
 pub mod graph_query_parser;
 mod io_maps;
 mod reader;

{nidx_binding-6.3.3.post92 → nidx_binding-6.3.3.post96}/nidx_relation/src/reader.rs

@@ -21,6 +21,7 @@ use std::collections::HashSet;
 use std::fmt::Debug;
 use std::path::Path;
 
+use nidx_protos::graph_search_request::QueryKind;
 use nidx_protos::relation_prefix_search_request::Search;
 use nidx_protos::{
     EntitiesSubgraphResponse, GraphSearchRequest, GraphSearchResponse, RelationNode, RelationPrefixSearchResponse,
@@ -30,8 +31,9 @@ use tantivy::collector::TopDocs;
 use tantivy::query::{BooleanQuery, Occur, Query};
 use tantivy::{DocAddress, Index, IndexReader, Searcher};
 
+use crate::graph_collector::{NodeSelector, TopUniqueNodeCollector, TopUniqueRelationCollector};
 use crate::graph_query_parser::{
-    BoolGraphQuery, Expression, FuzzyTerm, GraphQuery, GraphQueryParser, Node, NodeQuery, Term,
+    BoolGraphQuery, BoolNodeQuery, Expression, FuzzyTerm, GraphQuery, GraphQueryParser, Node, NodeQuery, Term,
 };
 use crate::schema::Schema;
 use crate::{io_maps, schema};
@@ -75,13 +77,22 @@ impl RelationsReaderService {
             return Ok(GraphSearchResponse::default());
         };
 
-        // Convert proto to tantivy query
-        let query = BoolGraphQuery::try_from(query)?;
+        let top_k = request.top_k as usize;
+
+        let response = match request.kind() {
+            QueryKind::Path => self.paths_graph_search(query, top_k),
+            QueryKind::Nodes => self.nodes_graph_search(query, top_k),
+            QueryKind::Relations => self.relations_graph_search(query, top_k),
+        };
+        response
+    }
+
+    pub fn inner_graph_search(&self, query: GraphQuery) -> anyhow::Result<nidx_protos::GraphSearchResponse> {
         let parser = GraphQueryParser::new();
-        let index_query = parser.parse_bool(query);
+        let index_query: Box<dyn Query> = parser.parse(query);
 
-        // Tantivy searcher query
-        let collector = TopDocs::with_limit(request.top_k as usize);
+        // TODO: parametrize this magic constant
+        let collector = TopDocs::with_limit(1000);
 
         let searcher = self.reader.searcher();
         let matching_docs = searcher.search(&index_query, &collector)?;
@@ -92,6 +103,87 @@ impl RelationsReaderService {
         )
     }
 
+    fn paths_graph_search(
+        &self,
+        query: &nidx_protos::graph_query::PathQuery,
+        top_k: usize,
+    ) -> anyhow::Result<GraphSearchResponse> {
+        let query = BoolGraphQuery::try_from(query)?;
+        let parser = GraphQueryParser::new();
+        let index_query = parser.parse_bool(query);
+
+        let collector = TopDocs::with_limit(top_k);
+        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 nodes_graph_search(
+        &self,
+        query: &nidx_protos::graph_query::PathQuery,
+        top_k: usize,
+    ) -> anyhow::Result<GraphSearchResponse> {
+        let query = BoolNodeQuery::try_from(query)?;
+        let parser = GraphQueryParser::new();
+        let (source_query, destination_query) = parser.parse_bool_node(query);
+
+        let mut unique_nodes = HashSet::new();
+
+        let collector = TopUniqueNodeCollector::new(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(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)
+    }
+
+    fn relations_graph_search(
+        &self,
+        query: &nidx_protos::graph_query::PathQuery,
+        top_k: usize,
+    ) -> anyhow::Result<GraphSearchResponse> {
+        let query = BoolGraphQuery::try_from(query)?;
+        let parser = GraphQueryParser::new();
+        let index_query = parser.parse_bool(query);
+
+        let collector = TopUniqueRelationCollector::new(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();
+
+        let response = nidx_protos::GraphSearchResponse {
+            relations,
+            ..Default::default()
+        };
+        Ok(response)
+    }
+
     fn build_graph_response(
         &self,
         searcher: &Searcher,
@@ -133,22 +225,6 @@ impl RelationsReaderService {
         };
         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);
-
-        // TODO: parametrize this magic constant
-        let collector = TopDocs::with_limit(1000);
-
-        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),
-        )
-    }
 }
 
 impl RelationsReaderService {

{nidx_binding-6.3.3.post92 → nidx_binding-6.3.3.post96}/pyproject.toml

@@ -13,7 +13,7 @@ classifiers = [
     "Programming Language :: Python :: Implementation :: CPython",
     "Programming Language :: Python :: Implementation :: PyPy",
 ]
-version = "6.3.3.post92"
+version = "6.3.3.post96"
 
 [project.urls]
 Homepage = "https://nuclia.com"