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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nidx_binding
-Version: 6.3.4.post159
+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.post159 → 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.post159"
+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.post159 → 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.post159 → 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);
+    }
+}

{nidx_binding-6.3.4.post159 → nidx_binding-6.3.4.post161}/nidx_relation/tests/test_graph_search.rs

@@ -19,11 +19,12 @@
 //
 mod common;
 
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use nidx_protos::graph_query::node::MatchKind;
 use nidx_protos::graph_query::path_query::Query;
 use nidx_protos::graph_query::{BoolQuery, FacetFilter, Node, Path, PathQuery, Relation};
+use nidx_protos::graph_search_request::QueryKind;
 use nidx_protos::relation::RelationType;
 use nidx_protos::relation_node::NodeType;
 use nidx_protos::{
@@ -39,13 +40,30 @@ use common::TestOpener;
 use uuid::Uuid;
 
 fn search(reader: &RelationSearcher, query: Query) -> anyhow::Result<GraphSearchResponse> {
-    search_with_prefilter(reader, query, PrefilterResult::All)
+    _search(reader, query, QueryKind::Path, PrefilterResult::All)
+}
+
+fn search_nodes(reader: &RelationSearcher, query: Query) -> anyhow::Result<GraphSearchResponse> {
+    _search(reader, query, QueryKind::Nodes, PrefilterResult::All)
+}
+
+fn search_relations(reader: &RelationSearcher, query: Query) -> anyhow::Result<GraphSearchResponse> {
+    _search(reader, query, QueryKind::Relations, PrefilterResult::All)
 }
 
 fn search_with_prefilter(
     reader: &RelationSearcher,
     query: Query,
     prefilter: PrefilterResult,
+) -> anyhow::Result<GraphSearchResponse> {
+    _search(reader, query, QueryKind::Path, prefilter)
+}
+
+fn _search(
+    reader: &RelationSearcher,
+    query: Query,
+    kind: QueryKind,
+    prefilter: PrefilterResult,
 ) -> anyhow::Result<GraphSearchResponse> {
     reader.graph_search(
         &GraphSearchRequest {
@@ -53,18 +71,68 @@ fn search_with_prefilter(
                 path: Some(PathQuery { query: Some(query) }),
             }),
             top_k: 100,
+            kind: kind.into(),
             ..Default::default()
         },
         prefilter,
     )
 }
 
+#[test]
+fn test_node_search() -> anyhow::Result<()> {
+    let reader = create_reader()?;
+
+    // (:PLACE)
+    let result = search_nodes(
+        &reader,
+        Query::Path(Path {
+            source: Some(Node {
+                node_subtype: Some("PLACE".to_string()),
+                ..Default::default()
+            }),
+            undirected: true,
+            ..Default::default()
+        }),
+    )?;
+    assert_eq!(result.nodes.len(), 2);
+    let nodes: HashSet<&str> = result.nodes.iter().map(|node| node.value.as_str()).collect();
+    assert!(nodes.contains(&"New York"));
+    assert!(nodes.contains(&"UK"));
+
+    Ok(())
+}
+
+#[test]
+fn test_relation_search() -> anyhow::Result<()> {
+    let reader = create_reader()?;
+
+    // (:Synonym)
+    let result = search_relations(
+        &reader,
+        Query::Path(Path {
+            relation: Some(Relation {
+                relation_type: Some(RelationType::Synonym.into()),
+                ..Default::default()
+            }),
+            ..Default::default()
+        }),
+    )?;
+    assert_eq!(result.relations.len(), 1);
+    let relations: HashSet<&str> = result
+        .relations
+        .iter()
+        .map(|relation| relation.label.as_str())
+        .collect();
+    assert!(relations.contains(&"ALIAS"));
+
+    Ok(())
+}
+
 #[test]
 fn test_graph_node_query() -> anyhow::Result<()> {
     let reader = create_reader()?;
 
     // (s)-[]->()
-
     let result = search(&reader, Query::Path(Path::default()))?;
     assert_eq!(result.graph.len(), 17);
 

{nidx_binding-6.3.4.post159 → nidx_binding-6.3.4.post161}/pyproject.toml

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