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

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

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

{nidx_binding-6.3.4.post148 → nidx_binding-6.3.4.post154}/nidx_paragraph/tests/reader.rs

@@ -156,7 +156,6 @@ fn create_resource(shard_id: String, timestamp: Timestamp) -> Resource {
         labels: vec!["/l/mylabel_resource".to_string()],
         paragraphs,
         paragraphs_to_delete: vec![],
-        relations: vec![],
         vectors: HashMap::default(),
         vectors_to_delete: HashMap::default(),
         shard_id,

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

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

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

@@ -0,0 +1,177 @@
+// 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::HashSet;
+
+use tantivy::{
+    DocId, Score, SegmentOrdinal, SegmentReader,
+    collector::{Collector, SegmentCollector},
+    columnar::Column,
+};
+
+#[derive(Clone, Copy)]
+pub enum NodeSelector {
+    SourceNodes,
+    DestinationNodes,
+}
+
+// Node collector for schema v2
+//
+// We can now use fast fields to uniquely identify nodes.
+
+pub struct TopUniqueNodeCollector2 {
+    limit: usize,
+    selector: NodeSelector,
+}
+
+pub struct TopUniqueNodeSegmentCollector2 {
+    limit: usize,
+    unique: HashSet<Vec<u64>>,
+    encoded_node_reader: Column<u64>,
+}
+
+// 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 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 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
+        if self.unique.len() >= self.limit {
+            return;
+        }
+
+        let relation = self
+            .encoded_relation_reader
+            .values_for_doc(doc_id)
+            .collect::<Vec<u64>>();
+        self.unique.insert(relation);
+    }
+
+    fn harvest(self) -> Self::Fruit {
+        self.unique
+    }
+}

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

@@ -120,21 +120,10 @@ enum NodePosition {
 pub struct BoolNodeQuery(BooleanExpression<Node>);
 pub struct BoolGraphQuery(BooleanExpression<GraphQuery>);
 
-#[derive(Clone, Copy)]
-struct TokenizedNodeFields {
-    exact: Field,
-    tokenized: Field,
-}
-
-#[derive(Clone, Copy)]
-enum NodeValueField {
-    Normalized(Field),
-    Tokenized(TokenizedNodeFields),
-}
-
 #[derive(Clone, Copy)]
 struct NodeSchemaFields {
-    value: NodeValueField,
+    exact_value: Field,
+    tokenized_value: Field,
     node_type: Field,
     node_subtype: Field,
 }
@@ -286,18 +275,11 @@ impl<'a> GraphQueryParser<'a> {
 
     #[inline]
     fn has_node_expression_as_source(&self, expression: &Expression<Node>) -> Vec<(Occur, Box<dyn Query>)> {
-        let value = if self.schema.version == 1 {
-            NodeValueField::Normalized(self.schema.normalized_source_value)
-        } else {
-            NodeValueField::Tokenized(TokenizedNodeFields {
-                exact: self.schema.normalized_source_value,
-                tokenized: self.schema.source_value,
-            })
-        };
         self.has_node_expression(
             expression,
             NodeSchemaFields {
-                value,
+                exact_value: self.schema.normalized_source_value,
+                tokenized_value: self.schema.source_value,
                 node_type: self.schema.source_type,
                 node_subtype: self.schema.source_subtype,
             },
@@ -306,18 +288,11 @@ impl<'a> GraphQueryParser<'a> {
 
     #[inline]
     fn has_node_expression_as_destination(&self, expression: &Expression<Node>) -> Vec<(Occur, Box<dyn Query>)> {
-        let value = if self.schema.version == 1 {
-            NodeValueField::Normalized(self.schema.normalized_target_value)
-        } else {
-            NodeValueField::Tokenized(TokenizedNodeFields {
-                exact: self.schema.normalized_target_value,
-                tokenized: self.schema.target_value,
-            })
-        };
         self.has_node_expression(
             expression,
             NodeSchemaFields {
-                value,
+                exact_value: self.schema.normalized_target_value,
+                tokenized_value: self.schema.target_value,
                 node_type: self.schema.target_type,
                 node_subtype: self.schema.target_subtype,
             },
@@ -403,7 +378,7 @@ impl<'a> GraphQueryParser<'a> {
         let value_query = node
             .value
             .as_ref()
-            .and_then(|value| self.has_node_value(value, fields.value));
+            .and_then(|value| self.has_node_value(value, fields.exact_value, fields.tokenized_value));
         if let Some(query) = value_query {
             subqueries.push(query);
         }
@@ -468,7 +443,7 @@ impl<'a> GraphQueryParser<'a> {
         subqueries
     }
 
-    fn has_node_value(&self, value: &Term, field: NodeValueField) -> Option<Box<dyn Query>> {
+    fn has_node_value(&self, value: &Term, exact_field: Field, tokenized_field: Field) -> Option<Box<dyn Query>> {
         let text_value = match value {
             Term::Exact(value) => value,
             Term::Fuzzy(fuzzy) => &fuzzy.value,
@@ -476,31 +451,14 @@ impl<'a> GraphQueryParser<'a> {
         if text_value.is_empty() {
             return None;
         }
-        let exact_term = match field {
-            NodeValueField::Normalized(field) => {
-                tantivy::Term::from_field_text(field, &self.schema.normalize(text_value))
-            }
-            NodeValueField::Tokenized(TokenizedNodeFields { exact, .. }) => {
-                tantivy::Term::from_field_text(exact, &self.schema.normalize(text_value))
-            }
-        };
-        let tokenized_terms = match field {
-            NodeValueField::Normalized(field) => {
-                vec![tantivy::Term::from_field_text(
-                    field,
-                    &self.schema.normalize(text_value),
-                )]
-            }
-            NodeValueField::Tokenized(TokenizedNodeFields { tokenized, .. }) => {
-                let mut tokenizer = TokenizerManager::default().get("default").unwrap();
-                let mut token_stream = tokenizer.token_stream(text_value);
-                let mut terms = Vec::new();
-                while let Some(token) = token_stream.next() {
-                    terms.push(tantivy::Term::from_field_text(tokenized, &token.text));
-                }
-                terms
-            }
-        };
+        let exact_term = tantivy::Term::from_field_text(exact_field, &self.schema.normalize(text_value));
+
+        let mut tokenizer = TokenizerManager::default().get("default").unwrap();
+        let mut token_stream = tokenizer.token_stream(text_value);
+        let mut tokenized_terms = Vec::new();
+        while let Some(token) = token_stream.next() {
+            tokenized_terms.push(tantivy::Term::from_field_text(tokenized_field, &token.text));
+        }
 
         // TODO: Rethink this
         // Current logic:

{nidx_binding-6.3.4.post148 → nidx_binding-6.3.4.post154}/nidx_relation/src/io_maps.rs

@@ -105,25 +105,18 @@ pub fn doc_to_relation(schema: &Schema, doc: &TantivyDocument) -> ProtosRelation
         relation_label: schema.relationship_label(doc).to_string(),
         source: Some(source_to_relation_node(schema, doc)),
         to: Some(target_to_relation_node(schema, doc)),
-        resource_id: Some(schema.resource_id(doc)),
     }
 }
 
 pub fn doc_to_index_relation(schema: &Schema, doc: &TantivyDocument) -> IndexRelation {
-    let resource_field_id = schema.resource_field_id.and_then(|f| {
-        doc.get_first(f).map(|v| {
-            let (rid, fid) = decode_field_id(v.as_bytes().unwrap());
-            format!("{}/{}", rid.simple(), fid)
-        })
+    let resource_field_id = doc.get_first(schema.resource_field_id).map(|v| {
+        let (rid, fid) = decode_field_id(v.as_bytes().unwrap());
+        format!("{}/{}", rid.simple(), fid)
     });
-    let facets = schema
-        .facets
-        .map(|field| {
-            doc.get_all(field)
-                .map(|f| f.as_facet().unwrap().to_path_string())
-                .collect()
-        })
-        .unwrap_or_default();
+    let facets = doc
+        .get_all(schema.facets)
+        .map(|f| f.as_facet().unwrap().to_path_string())
+        .collect();
     IndexRelation {
         relation: Some(doc_to_relation(schema, doc)),
         resource_field_id,

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

@@ -25,6 +25,7 @@ mod reader;
 mod resource_indexer;
 mod schema;
 
+use anyhow::anyhow;
 use nidx_protos::{
     GraphSearchRequest, GraphSearchResponse, RelationNode, RelationNodeFilter, RelationPrefixSearchRequest,
     RelationSearchRequest, RelationSearchResponse, relation_node::NodeType, relation_prefix_search_request::Search,
@@ -82,35 +83,29 @@ pub struct RelationIndexer;
 
 pub struct RelationDeletionQueryBuilder {
     resource: Field,
-    field: Option<Field>,
+    field: Field,
 }
 impl DeletionQueryBuilder for RelationDeletionQueryBuilder {
     fn query<'a>(&self, keys: impl Iterator<Item = &'a String>) -> Box<dyn Query> {
-        if let Some(field) = self.field {
-            Box::new(TermSetQuery::new(keys.filter_map(|k| {
-                // Our keys can be resource or field ids, match the corresponding tantivy field
-                if k.len() < 32 {
-                    error!(?k, "Invalid deletion key for nidx_relation");
-                    return None;
-                }
-
-                let Ok(rid) = Uuid::parse_str(&k[..32]) else {
-                    error!(?k, "Invalid deletion key for nidx_relation");
-                    return None;
-                };
-
-                let is_field = k.len() > 32;
-                if is_field {
-                    Some(Term::from_field_bytes(field, &encode_field_id(rid, &k[33..])))
-                } else {
-                    Some(Term::from_field_bytes(self.resource, rid.as_bytes()))
-                }
-            })))
-        } else {
-            Box::new(TermSetQuery::new(
-                keys.map(|k| Term::from_field_bytes(self.resource, k.as_bytes())),
-            ))
-        }
+        Box::new(TermSetQuery::new(keys.filter_map(|k| {
+            // Our keys can be resource or field ids, match the corresponding tantivy field
+            if k.len() < 32 {
+                error!(?k, "Invalid deletion key for nidx_relation");
+                return None;
+            }
+
+            let Ok(rid) = Uuid::parse_str(&k[..32]) else {
+                error!(?k, "Invalid deletion key for nidx_relation");
+                return None;
+            };
+
+            let is_field = k.len() > 32;
+            if is_field {
+                Some(Term::from_field_bytes(self.field, &encode_field_id(rid, &k[33..])))
+            } else {
+                Some(Term::from_field_bytes(self.resource, rid.as_bytes()))
+            }
+        })))
     }
 }
 impl RelationDeletionQueryBuilder {
@@ -130,6 +125,9 @@ impl RelationIndexer {
         config: &RelationConfig,
         resource: &nidx_protos::Resource,
     ) -> anyhow::Result<Option<TantivySegmentMetadata>> {
+        if config.version != 2 {
+            return Err(anyhow!("Unsupported nidx_relation version"));
+        }
         let field_schema = RelationSchema::new(config.version);
         let mut indexer = TantivyIndexer::new(output_dir.to_path_buf(), field_schema.schema.clone())?;
 
@@ -141,17 +139,13 @@ impl RelationIndexer {
         indexer.finalize()
     }
 
-    pub fn deletions_for_resource(&self, config: &RelationConfig, resource: &nidx_protos::Resource) -> Vec<String> {
-        if config.version == 2 {
-            let rid = &resource.resource.as_ref().unwrap().uuid;
-            resource
-                .relation_fields_to_delete
-                .iter()
-                .map(|f| format!("{rid}/{f}"))
-                .collect()
-        } else {
-            vec![resource.resource.as_ref().unwrap().uuid.clone()]
-        }
+    pub fn deletions_for_resource(&self, _config: &RelationConfig, resource: &nidx_protos::Resource) -> Vec<String> {
+        let rid = &resource.resource.as_ref().unwrap().uuid;
+        resource
+            .relation_fields_to_delete
+            .iter()
+            .map(|f| format!("{rid}/{f}"))
+            .collect()
     }
 
     #[instrument(name = "relation::merge", skip_all)]
@@ -161,6 +155,9 @@ impl RelationIndexer {
         config: RelationConfig,
         open_index: impl OpenIndexMetadata<TantivyMeta>,
     ) -> anyhow::Result<TantivySegmentMetadata> {
+        if config.version != 2 {
+            return Err(anyhow!("Unsupported nidx_relation version"));
+        }
         let schema = RelationSchema::new(config.version);
         let deletions_query = RelationDeletionQueryBuilder::new(&schema);
         let index = open_index_with_deletions(schema.schema, open_index, deletions_query)?;
@@ -185,6 +182,9 @@ pub struct RelationSearcher {
 impl RelationSearcher {
     #[instrument(name = "relation::open", skip_all)]
     pub fn open(config: RelationConfig, open_index: impl OpenIndexMetadata<TantivyMeta>) -> anyhow::Result<Self> {
+        if config.version != 2 {
+            return Err(anyhow!("Unsupported nidx_relation version"));
+        }
         let schema = RelationSchema::new(config.version);
         let deletions_query = RelationDeletionQueryBuilder::new(&schema);
         let index = open_index_with_deletions(schema.schema, open_index, deletions_query)?;