nidx-binding 6.3.4.post154__tar.gz → 6.3.4.post161__tar.gz

{nidx_binding-6.3.4.post154 → nidx_binding-6.3.4.post161}/PKG-INFO

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

{nidx_binding-6.3.4.post154 → nidx_binding-6.3.4.post161}/nidx_protos/pyproject.toml

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

nidx_binding-6.3.4.post161/nidx_relation/src/graph_collector.rs

@@ -0,0 +1,97 @@
+// 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 tantivy::{
+    DocId, Score, SegmentOrdinal, SegmentReader,
+    collector::{Collector, SegmentCollector},
+    columnar::Column,
+};
+
+use crate::top_unique_n::TopUniqueN;
+
+#[derive(Clone, Copy)]
+pub enum Selector {
+    SourceNodes,
+    DestinationNodes,
+    Relations,
+}
+
+pub struct TopUniqueCollector {
+    limit: usize,
+    selector: Selector,
+}
+
+pub struct TopUniqueSegmentCollector {
+    unique: TopUniqueN<Vec<u64>>,
+    encoded_field: Column<u64>,
+}
+
+impl TopUniqueCollector {
+    pub fn new(selector: Selector, top_k: usize) -> Self {
+        Self { limit: top_k, selector }
+    }
+}
+
+impl Collector for TopUniqueCollector {
+    type Fruit = TopUniqueN<Vec<u64>>;
+    type Child = TopUniqueSegmentCollector;
+
+    fn requires_scoring(&self) -> bool {
+        true
+    }
+
+    fn for_segment(&self, _segment_local_id: SegmentOrdinal, segment: &SegmentReader) -> tantivy::Result<Self::Child> {
+        let fast_field = match self.selector {
+            Selector::SourceNodes => segment.fast_fields().u64("encoded_source_id")?,
+            Selector::DestinationNodes => segment.fast_fields().u64("encoded_target_id")?,
+            Selector::Relations => segment.fast_fields().u64("encoded_relation_id")?,
+        };
+        let segment_collector = TopUniqueSegmentCollector {
+            unique: TopUniqueN::new(self.limit),
+            encoded_field: fast_field,
+        };
+        Ok(segment_collector)
+    }
+
+    fn merge_fruits(
+        &self,
+        segment_fruits: Vec<<Self::Child as SegmentCollector>::Fruit>,
+    ) -> tantivy::Result<Self::Fruit> {
+        let fruits = segment_fruits.into_iter().flat_map(|top| top.into_sorted_vec());
+        let mut unique = TopUniqueN::new(self.limit);
+        for (key, score) in fruits {
+            unique.insert(key, score);
+        }
+        Ok(unique)
+    }
+}
+
+impl SegmentCollector for TopUniqueSegmentCollector {
+    type Fruit = TopUniqueN<Vec<u64>>;
+
+    fn collect(&mut self, doc_id: DocId, score: Score) {
+        let value = self.encoded_field.values_for_doc(doc_id).collect::<Vec<u64>>();
+        self.unique.insert(value, score);
+    }
+
+    fn harvest(self) -> Self::Fruit {
+        self.unique
+    }
+}

{nidx_binding-6.3.4.post154 → nidx_binding-6.3.4.post161}/nidx_relation/src/graph_query_parser.rs

@@ -18,11 +18,12 @@
 // along with this program. If not, see <http://www.gnu.org/licenses/>.
 //
 use anyhow::anyhow;
+use nidx_protos::graph_query::FacetFilter;
 use nidx_protos::relation::RelationType;
 use nidx_protos::relation_node::NodeType;
 use nidx_types::query_language::{BooleanExpression, BooleanOperation, Operator};
 use tantivy::query::{AllQuery, BooleanQuery, FuzzyTermQuery, Occur, Query, TermQuery};
-use tantivy::schema::{Field, IndexRecordOption};
+use tantivy::schema::{Facet, Field, IndexRecordOption};
 use tantivy::tokenizer::TokenizerManager;
 
 use crate::io_maps;
@@ -109,6 +110,8 @@ pub enum GraphQuery {
     RelationQuery(RelationQuery),
     // (:A)-[:R]->(:B)
     PathQuery(PathQuery),
+
+    Facet(String),
 }
 
 #[derive(Clone, Copy)]
@@ -117,7 +120,13 @@ enum NodePosition {
     Destination,
 }
 
-pub struct BoolNodeQuery(BooleanExpression<Node>);
+#[derive(Clone)]
+enum NodeExpression {
+    Node(Node),
+    Facet(String),
+}
+
+pub struct BoolNodeQuery(BooleanExpression<NodeExpression>);
 pub struct BoolGraphQuery(BooleanExpression<GraphQuery>);
 
 #[derive(Clone, Copy)]
@@ -172,9 +181,13 @@ impl<'a> GraphQueryParser<'a> {
         )
     }
 
-    fn inner_parse_bool_node(&self, query: BooleanExpression<Node>, position: NodePosition) -> Box<dyn Query> {
+    fn inner_parse_bool_node(
+        &self,
+        query: BooleanExpression<NodeExpression>,
+        position: NodePosition,
+    ) -> Box<dyn Query> {
         match query {
-            BooleanExpression::Literal(node) => match position {
+            BooleanExpression::Literal(NodeExpression::Node(node)) => match position {
                 NodePosition::Source => {
                     self.parse(GraphQuery::NodeQuery(NodeQuery::SourceNode(Expression::Value(node))))
                 }
@@ -182,6 +195,7 @@ impl<'a> GraphQueryParser<'a> {
                     Expression::Value(node),
                 ))),
             },
+            BooleanExpression::Literal(NodeExpression::Facet(facet)) => self.parse_facet(&facet),
             BooleanExpression::Not(subquery) => {
                 let subqueries = vec![
                     (Occur::Must, Box::new(AllQuery) as Box<dyn Query>),
@@ -210,6 +224,7 @@ impl<'a> GraphQueryParser<'a> {
             GraphQuery::NodeQuery(query) => self.parse_node_query(query),
             GraphQuery::RelationQuery(query) => self.parse_relation_query(query),
             GraphQuery::PathQuery(query) => self.parse_path_query(query),
+            GraphQuery::Facet(facet) => self.parse_facet(&facet),
         }
     }
 
@@ -245,6 +260,13 @@ impl<'a> GraphQueryParser<'a> {
         self.parse_path_query(equivalent_path_query)
     }
 
+    fn parse_facet(&self, facet: &str) -> Box<dyn Query> {
+        Box::new(TermQuery::new(
+            tantivy::Term::from_facet(self.schema.facets, &Facet::from_text(facet).expect("Invalid facet")),
+            IndexRecordOption::Basic,
+        ))
+    }
+
     fn parse_path_query(&self, query: PathQuery) -> Box<dyn Query> {
         match query {
             PathQuery::DirectedPath((source_expression, relation_expression, destination_expression)) => {
@@ -564,14 +586,12 @@ impl TryFrom<&nidx_protos::graph_query::PathQuery> for BoolNodeQuery {
 
                     let pb_node = path.source.as_ref().unwrap();
                     let node = Node::try_from(pb_node)?;
-                    BooleanExpression::Literal(node)
+                    BooleanExpression::Literal(NodeExpression::Node(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,
@@ -582,7 +602,6 @@ impl TryFrom<&nidx_protos::graph_query::PathQuery> for BoolNodeQuery {
                             .collect::<anyhow::Result<Vec<_>>>()?,
                     })
                 }
-
                 nidx_protos::graph_query::path_query::Query::BoolOr(bool_or) => {
                     BooleanExpression::Operation(BooleanOperation {
                         operator: Operator::Or,
@@ -593,9 +612,12 @@ impl TryFrom<&nidx_protos::graph_query::PathQuery> for BoolNodeQuery {
                             .collect::<anyhow::Result<Vec<_>>>()?,
                     })
                 }
+                nidx_protos::graph_query::path_query::Query::Facet(FacetFilter { facet }) => {
+                    BooleanExpression::Literal(NodeExpression::Facet(facet.clone()))
+                }
             },
 
-            None => BooleanExpression::Literal(Node::default()),
+            None => BooleanExpression::Literal(NodeExpression::Node(Node::default())),
         };
 
         Ok(BoolNodeQuery(bool_expr))
@@ -638,12 +660,10 @@ impl TryFrom<&nidx_protos::graph_query::PathQuery> for BoolGraphQuery {
 
                     BooleanExpression::Literal(path_query)
                 }
-
                 nidx_protos::graph_query::path_query::Query::BoolNot(bool_not) => {
                     let subquery = BoolGraphQuery::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,
@@ -654,7 +674,6 @@ impl TryFrom<&nidx_protos::graph_query::PathQuery> for BoolGraphQuery {
                             .collect::<anyhow::Result<Vec<_>>>()?,
                     })
                 }
-
                 nidx_protos::graph_query::path_query::Query::BoolOr(bool_or) => {
                     BooleanExpression::Operation(BooleanOperation {
                         operator: Operator::Or,
@@ -665,6 +684,9 @@ impl TryFrom<&nidx_protos::graph_query::PathQuery> for BoolGraphQuery {
                             .collect::<anyhow::Result<Vec<_>>>()?,
                     })
                 }
+                nidx_protos::graph_query::path_query::Query::Facet(FacetFilter { facet }) => {
+                    BooleanExpression::Literal(GraphQuery::Facet(facet.clone()))
+                }
             },
 
             None => BooleanExpression::Literal(GraphQuery::PathQuery(PathQuery::UndirectedPath((

{nidx_binding-6.3.4.post154 → nidx_binding-6.3.4.post161}/nidx_relation/src/lib.rs

@@ -24,6 +24,7 @@ mod io_maps;
 mod reader;
 mod resource_indexer;
 mod schema;
+mod top_unique_n;
 
 use anyhow::anyhow;
 use nidx_protos::{

{nidx_binding-6.3.4.post154 → nidx_binding-6.3.4.post161}/nidx_relation/src/reader.rs

@@ -31,14 +31,15 @@ use nidx_types::prefilter::{FieldId, PrefilterResult};
 use tantivy::collector::TopDocs;
 use tantivy::query::{BooleanQuery, EmptyQuery, Occur, Query, TermSetQuery};
 use tantivy::schema::Field;
-use tantivy::{DocAddress, Index, IndexReader, Searcher};
+use tantivy::{Index, IndexReader};
 use uuid::Uuid;
 
-use crate::graph_collector::{NodeSelector, TopUniqueNodeCollector2, TopUniqueRelationCollector2};
+use crate::graph_collector::{Selector, TopUniqueCollector};
 use crate::graph_query_parser::{
     BoolGraphQuery, BoolNodeQuery, Expression, FuzzyTerm, GraphQuery, GraphQueryParser, Node, NodeQuery, Term,
 };
 use crate::schema::{Schema, decode_node, decode_relation, encode_field_id};
+use crate::top_unique_n::TopUniqueN;
 use crate::{RelationConfig, io_maps};
 
 const FUZZY_DISTANCE: u8 = 1;
@@ -146,10 +147,38 @@ impl RelationsReaderService {
         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),
-        )
+
+        let mut nodes = Vec::new();
+        let mut relations = Vec::new();
+        let mut graph = Vec::new();
+
+        for (_score, doc_address) in matching_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)
     }
 
     fn nodes_graph_search(
@@ -166,25 +195,27 @@ impl RelationsReaderService {
 
         let searcher = self.reader.searcher();
 
-        let mut unique_nodes = HashSet::new();
+        let mut unique_nodes = TopUniqueN::new(top_k);
 
-        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 = TopUniqueCollector::new(Selector::SourceNodes, top_k);
+        let source_nodes = searcher.search(&source_query, &collector)?;
+        unique_nodes.merge(source_nodes);
 
-        let collector = TopUniqueNodeCollector2::new(NodeSelector::DestinationNodes, top_k);
-        let mut destination_nodes = searcher.search(&destination_query, &collector)?;
-        unique_nodes.extend(destination_nodes.drain());
+        let collector = TopUniqueCollector::new(Selector::DestinationNodes, top_k);
+        let destination_nodes = searcher.search(&destination_query, &collector)?;
+        unique_nodes.merge(destination_nodes);
 
         let nodes = unique_nodes
+            .into_sorted_vec()
             .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,
+            .map(|(encoded_node, _score)| {
+                let (value, node_type, node_subtype) = decode_node(&encoded_node);
+                RelationNode {
+                    value,
+                    ntype: io_maps::u64_to_node_type(node_type),
+                    subtype: node_subtype,
+                }
             })
-            .take(top_k)
             .collect();
 
         let response = nidx_protos::GraphSearchResponse {
@@ -207,12 +238,13 @@ impl RelationsReaderService {
 
         let searcher = self.reader.searcher();
 
-        let collector = TopUniqueRelationCollector2::new(top_k);
-        let matching_docs = searcher.search(&index_query, &collector)?;
+        let collector = TopUniqueCollector::new(Selector::Relations, top_k);
+        let top_relations = searcher.search(&index_query, &collector)?;
 
-        let relations = matching_docs
+        let relations = top_relations
+            .into_sorted_vec()
             .into_iter()
-            .map(|encoded_relation| {
+            .map(|(encoded_relation, _score)| {
                 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),
@@ -240,48 +272,6 @@ impl RelationsReaderService {
             }
         }
     }
-
-    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)
-    }
 }
 
 impl RelationsReaderService {

nidx_binding-6.3.4.post161/nidx_relation/src/top_unique_n.rs

@@ -0,0 +1,164 @@
+// 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;
+
+/// Top unique N
+///
+/// Maintain the top set of unique keys with greatest scores.
+pub struct TopUniqueN<K> {
+    elements: HashMap<K, f32>,
+    top_n: usize,
+    threshold: f32,
+}
+
+impl<K> TopUniqueN<K>
+where
+    K: Eq + std::hash::Hash + std::fmt::Debug,
+{
+    pub fn new(top_n: usize) -> Self {
+        Self {
+            top_n,
+            elements: HashMap::with_capacity(2 * top_n),
+            threshold: f32::NEG_INFINITY,
+        }
+    }
+
+    pub fn insert(&mut self, key: K, score: f32) {
+        if score < self.threshold {
+            return;
+        }
+
+        if self.elements.len() == self.elements.capacity() {
+            let lowest_score = self.truncate_top_n();
+            self.threshold = lowest_score;
+        }
+
+        self.elements
+            .entry(key)
+            .and_modify(|s| {
+                if score > *s {
+                    *s = score
+                }
+            })
+            .or_insert(score);
+    }
+
+    // Truncate the current set of element to N leaving only the top-scoring
+    // elements. Return the smallest score across the top.
+    fn truncate_top_n(&mut self) -> f32 {
+        let mut vec = Vec::from_iter(self.elements.drain());
+        vec.sort_unstable_by(|a, b| a.1.total_cmp(&b.1).reverse());
+        vec.truncate(self.top_n);
+        let lowest_score = vec.last().map(|(_, score)| *score).unwrap_or(f32::NEG_INFINITY);
+
+        self.elements.extend(vec);
+
+        lowest_score
+    }
+
+    pub fn into_sorted_vec(self) -> Vec<(K, f32)> {
+        let mut vec = Vec::from_iter(self.elements);
+        vec.sort_by(|a, b| a.1.total_cmp(&b.1).reverse());
+        vec.truncate(self.top_n);
+        vec
+    }
+
+    pub fn merge(&mut self, other: Self) {
+        for (key, score) in other.elements.into_iter() {
+            self.insert(key, score);
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_top_n() {
+        let mut top = TopUniqueN::new(2);
+        top.insert("A", 1.0);
+        top.insert("B", 3.0);
+        top.insert("C", 2.0);
+        top.insert("D", 4.0);
+        top.insert("E", -1.0);
+
+        let r = top.into_sorted_vec();
+        assert_eq!(r.len(), 2);
+        let r: HashMap<_, _> = HashMap::from_iter(r);
+        let expected = HashMap::from_iter([("B", 3.0), ("D", 4.0)]);
+        assert_eq!(r, expected);
+    }
+
+    /// Validate inserting more than it's capacity, values are truncated to N.
+    #[test]
+    fn test_internal_truncate() {
+        const N: usize = 2;
+        let mut top = TopUniqueN::new(N);
+
+        // capacity is at least 2 * N + 1, but in reality, it's usually more
+        let actual_capacity = top.elements.capacity();
+        assert!(actual_capacity >= 2 * N);
+
+        let mut key_id = 0;
+        let mut key_generator = std::iter::repeat_with(|| {
+            let key = key_id.to_string();
+            key_id += 1;
+            key
+        });
+
+        while top.elements.len() < top.elements.capacity() {
+            let key = key_generator.next().unwrap();
+            top.insert(key, 1.0);
+        }
+        assert_eq!(top.elements.len(), top.elements.capacity());
+        assert!(top.threshold < 0.0);
+
+        // this insert would overflow the capacity, but it truncates the internal values and don't
+        // increase it
+        top.insert("A".to_string(), 1.0);
+        assert_eq!(top.elements.capacity(), actual_capacity);
+        assert_eq!(top.elements.len(), N + 1);
+        assert_eq!(top.threshold, 1.0);
+    }
+
+    #[test]
+    fn test_merge() {
+        let mut top_a = TopUniqueN::new(4);
+        top_a.insert("A1", 1.0);
+        top_a.insert("A2", 3.0);
+
+        let mut top_b = TopUniqueN::new(3);
+        top_b.insert("B1", 1.0);
+        top_b.insert("B2", 3.0);
+        top_b.insert("B3", 4.0);
+        top_b.insert("B4", 2.0);
+
+        top_a.merge(top_b);
+
+        let r = top_a.into_sorted_vec();
+        assert_eq!(r.len(), 4);
+
+        let r: HashMap<_, _> = HashMap::from_iter(r);
+        let expected = HashMap::from_iter([("A2", 3.0), ("B2", 3.0), ("B3", 4.0), ("B4", 2.0)]);
+        assert_eq!(r, expected);
+    }
+}