nidx-binding 6.10.0.post725__tar.gz → 6.10.0.post732__tar.gz

{nidx_binding-6.10.0.post725 → nidx_binding-6.10.0.post732}/Cargo.lock

@@ -1023,6 +1023,12 @@ version = "0.1.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2"
 
+[[package]]
+name = "foldhash"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
+
 [[package]]
 name = "foreign-types"
 version = "0.3.2"
@@ -1263,7 +1269,7 @@ checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1"
 dependencies = [
  "allocator-api2",
  "equivalent",
- "foldhash",
+ "foldhash 0.1.5",
 ]
 
 [[package]]
@@ -1271,6 +1277,11 @@ name = "hashbrown"
 version = "0.16.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100"
+dependencies = [
+ "allocator-api2",
+ "equivalent",
+ "foldhash 0.2.0",
+]
 
 [[package]]
 name = "hashlink"
@@ -2002,6 +2013,15 @@ dependencies = [
  "hashbrown 0.15.5",
 ]
 
+[[package]]
+name = "lru"
+version = "0.16.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a1dc47f592c06f33f8e3aea9591776ec7c9f9e4124778ff8a3c3b87159f7e593"
+dependencies = [
+ "hashbrown 0.16.1",
+]
+
 [[package]]
 name = "lru-slab"
 version = "0.1.2"
@@ -2158,7 +2178,7 @@ dependencies = [
  "k8s-openapi",
  "kube",
  "lazy_static",
- "lru",
+ "lru 0.16.3",
  "nidx_paragraph",
  "nidx_protos",
  "nidx_relation",
@@ -4409,7 +4429,7 @@ dependencies = [
  "itertools 0.14.0",
  "levenshtein_automata",
  "log",
- "lru",
+ "lru 0.12.5",
  "lz4_flex",
  "measure_time",
  "memmap2 0.9.9",

{nidx_binding-6.10.0.post725 → nidx_binding-6.10.0.post732}/Cargo.toml

@@ -31,7 +31,7 @@ tracing = "0.1.40"
 tracing-subscriber = { version = "0.3.18", features = ["env-filter", "json"] }
 nidx_protos = { version = "0.1.0", path = "nidx_protos" }
 tonic = "0.12.3"
-lru = "0.12.5"
+lru = "0.16"
 nidx_text = { version = "0.1.0", path = "nidx_text" }
 nidx_paragraph = { version = "0.1.0", path = "nidx_paragraph" }
 nidx_relation = { version = "0.1.0", path = "nidx_relation" }

{nidx_binding-6.10.0.post725 → nidx_binding-6.10.0.post732}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nidx_binding
-Version: 6.10.0.post725
+Version: 6.10.0.post732
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy

nidx_binding-6.10.0.post732/migrations/20251218173017_vector_relation_index_kinds.sql

@@ -0,0 +1,2 @@
+ALTER TYPE index_kind ADD VALUE 'vector_relation_node';
+ALTER TYPE index_kind ADD VALUE 'vector_relation_edge';

{nidx_binding-6.10.0.post725 → nidx_binding-6.10.0.post732}/nidx_protos/nodereader.proto

@@ -2,8 +2,8 @@ syntax = "proto3";
 
 package nodereader;
 
-import public "nidx_protos/noderesources.proto";
 import "google/protobuf/timestamp.proto";
+import public "nidx_protos/noderesources.proto";
 import public "nucliadb_protos/utils.proto";
 
 message StreamFilter {
@@ -17,18 +17,16 @@ message StreamFilter {
     repeated string labels = 2;
 }
 
-
 message Faceted {
     repeated string labels = 1;
 }
 
-
 message OrderBy {
     reserved 1;
 
     enum OrderType {
-      DESC = 0;
-      ASC = 1;
+        DESC = 0;
+        ASC = 1;
     }
     enum OrderField {
         CREATED = 0;
@@ -179,15 +177,31 @@ message GraphQuery {
             uint32 distance = 2;
         }
 
+        message VectorMatch {
+            repeated float vector = 1;
+        }
+
         oneof match_kind {
             ExactMatch exact = 5;
             FuzzyMatch fuzzy = 6;
+            VectorMatch vector = 7;
         }
     }
 
     message Relation {
         optional string value = 1;
         optional utils.Relation.RelationType relation_type = 2;
+
+        message ExactMatch {}
+
+        message VectorMatch {
+            repeated float vector = 1;
+        }
+
+        oneof match_kind {
+            ExactMatch exact = 3;
+            VectorMatch vector = 4;
+        }
     }
 
     message Path {
@@ -237,6 +251,8 @@ message GraphSearchRequest {
 
     optional utils.Security security = 5;
     optional FilterExpression field_filter = 6;
+
+    optional string graph_vectorset = 7;
 }
 
 message GraphSearchResponse {
@@ -260,6 +276,7 @@ message GraphSearchResponse {
     repeated utils.RelationNode nodes = 1;
     repeated Relation relations = 2;
     repeated Path graph = 3;
+    repeated float scores = 4;
 }
 
 message FilterExpression {
@@ -324,6 +341,7 @@ message SearchRequest {
     repeated float vector = 10; // Embedded vector search.
 
     string vectorset = 15;
+    optional string graph_vectorset = 30;
 
     bool paragraph = 12;
     bool document = 13;
@@ -351,7 +369,6 @@ message SearchRequest {
     GraphSearch graph_search = 29;
 }
 
-
 enum SuggestFeatures {
     ENTITIES = 0;
     PARAGRAPHS = 1;
@@ -368,7 +385,6 @@ message SuggestRequest {
     FilterOperator filter_operator = 9;
 }
 
-
 message SuggestResponse {
     int32 total = 1;
     repeated ParagraphResult results = 2;

{nidx_binding-6.10.0.post725 → nidx_binding-6.10.0.post732}/nidx_protos/noderesources.proto

@@ -167,7 +167,7 @@ message Resource {
     bool skip_paragraphs = 19;
 
     repeated RelationNodeVector relation_node_vectors = 20;
-    repeated RelationVector relation_vectors = 21;
+    repeated RelationEdgeVector relation_edge_vectors = 21;
 }
 
 message RelationNodeVector {
@@ -175,7 +175,7 @@ message RelationNodeVector {
     repeated float vector = 2;
 }
 
-message RelationVector {
+message RelationEdgeVector {
     utils.Relation.RelationType relation_type = 1;
     string relation_label = 2;
     repeated float vector = 3;

{nidx_binding-6.10.0.post725 → nidx_binding-6.10.0.post732}/nidx_protos/pyproject.toml

@@ -10,7 +10,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "nidx_protos"
-version = "6.10.0.post725"
+version = "6.10.0.post732"
 license = "AGPL-3.0-or-later"
 description = "Protobuf definitions for nucliadb/nidx"
 authors = [{ name = "Nuclia", email = "nucliadb@nuclia.com" }]

{nidx_binding-6.10.0.post725 → nidx_binding-6.10.0.post732}/nidx_relation/src/graph_query_parser.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 // Copyright (C) 2021 Bosutech XXI S.L.
 //
 // nucliadb is offered under the AGPL v3.0 and as commercial software.
@@ -23,7 +25,9 @@ use nidx_protos::graph_query::node::MatchLocation;
 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, TermSetQuery};
+use tantivy::query::{
+    AllQuery, BooleanQuery, BoostQuery, EmptyQuery, FuzzyTermQuery, Occur, Query, TermQuery, TermSetQuery,
+};
 use tantivy::schema::{Facet, Field, IndexRecordOption};
 use tantivy::tokenizer::TokenizerManager;
 
@@ -43,6 +47,13 @@ pub enum Term {
     ExactWord(String),
     Fuzzy(FuzzyTerm),
     FuzzyWord(FuzzyTerm),
+    FromVectorQuery(String),
+}
+
+#[derive(Clone)]
+pub enum RelationTerm {
+    Exact(String),
+    FromVectorQuery(String),
 }
 
 #[derive(Default, Clone)]
@@ -54,7 +65,7 @@ pub struct Node {
 
 #[derive(Default, Clone)]
 pub struct Relation {
-    pub value: Option<String>,
+    pub label: Option<RelationTerm>,
     pub relation_type: Option<RelationType>,
 }
 
@@ -138,13 +149,23 @@ struct NodeSchemaFields {
     node_subtype: Field,
 }
 
+#[derive(Default)]
+pub struct VectorQueryResults {
+    pub nodes: HashMap<String, Vec<(String, f32)>>,
+    pub edges: HashMap<String, Vec<(String, f32)>>,
+}
+
 pub struct GraphQueryParser<'a> {
     schema: &'a Schema,
+    vector_results: VectorQueryResults,
 }
 
 impl<'a> GraphQueryParser<'a> {
-    pub fn new(schema: &'a Schema) -> Self {
-        Self { schema }
+    pub fn new(schema: &'a Schema, context: VectorQueryResults) -> Self {
+        Self {
+            schema,
+            vector_results: context,
+        }
     }
 
     pub fn parse_bool(&self, query: BoolGraphQuery) -> Box<dyn Query> {
@@ -422,11 +443,12 @@ impl<'a> GraphQueryParser<'a> {
         let mut subqueries = vec![];
 
         match expression {
-            Expression::Value(Relation { value, relation_type }) => {
-                if let Some(value) = value
-                    && !value.is_empty()
-                {
-                    subqueries.push((Occur::Must, self.has_relation_label(value)));
+            Expression::Value(Relation {
+                label: term,
+                relation_type,
+            }) => {
+                if let Some(term) = term {
+                    subqueries.push((Occur::Must, self.has_relation_label(term)));
                 };
 
                 relation_type.map(|relation_type| {
@@ -434,11 +456,12 @@ impl<'a> GraphQueryParser<'a> {
                 });
             }
 
-            Expression::Not(Relation { value, relation_type }) => {
-                if let Some(value) = value
-                    && !value.is_empty()
-                {
-                    subqueries.push((Occur::MustNot, self.has_relation_label(value)));
+            Expression::Not(Relation {
+                label: term,
+                relation_type,
+            }) => {
+                if let Some(term) = term {
+                    subqueries.push((Occur::MustNot, self.has_relation_label(term)));
                 };
 
                 relation_type.map(|relation_type| {
@@ -448,15 +471,10 @@ impl<'a> GraphQueryParser<'a> {
 
             Expression::Or(relations) => {
                 subqueries.extend(relations.iter().flat_map(|relation| {
-                    if let Some(label) = &relation.value {
-                        if label.is_empty() {
-                            None
-                        } else {
-                            Some((Occur::Should, self.has_relation_label(label)))
-                        }
-                    } else {
-                        None
-                    }
+                    relation
+                        .label
+                        .as_ref()
+                        .map(|term| (Occur::Should, self.has_relation_label(term)))
                 }));
             }
         };
@@ -465,9 +483,31 @@ impl<'a> GraphQueryParser<'a> {
     }
 
     fn has_node_value(&self, value: &Term, exact_field: Field, tokenized_field: Field) -> Option<Box<dyn Query>> {
+        if let Term::FromVectorQuery(key) = value {
+            let Some(keys) = self.vector_results.nodes.get(key) else {
+                return Some(Box::new(EmptyQuery));
+            };
+
+            return Some(Box::new(BooleanQuery::union(
+                keys.iter()
+                    .map(|(value, score)| {
+                        let q: Box<dyn Query> = Box::new(BoostQuery::new(
+                            Box::new(TermQuery::new(
+                                tantivy::Term::from_field_text(exact_field, &self.schema.normalize(value)),
+                                IndexRecordOption::Basic,
+                            )),
+                            *score,
+                        ));
+                        q
+                    })
+                    .collect(),
+            )));
+        }
+
         let text_value = match value {
             Term::Exact(value) | Term::ExactWord(value) => value,
             Term::Fuzzy(fuzzy) | Term::FuzzyWord(fuzzy) => &fuzzy.value,
+            Term::FromVectorQuery(_) => unreachable!(),
         };
         if text_value.is_empty() {
             return None;
@@ -513,6 +553,7 @@ impl<'a> GraphQueryParser<'a> {
                     ))
                 }
             }
+            Term::FromVectorQuery(_) => unreachable!(),
         };
 
         Some(query)
@@ -533,11 +574,33 @@ impl<'a> GraphQueryParser<'a> {
         ))
     }
 
-    fn has_relation_label(&self, label: &str) -> Box<dyn Query> {
-        Box::new(TermQuery::new(
-            tantivy::Term::from_field_text(self.schema.label, label),
-            IndexRecordOption::Basic,
-        ))
+    fn has_relation_label(&self, term: &RelationTerm) -> Box<dyn Query> {
+        match term {
+            RelationTerm::Exact(label) => Box::new(TermQuery::new(
+                tantivy::Term::from_field_text(self.schema.label, label),
+                IndexRecordOption::Basic,
+            )),
+            RelationTerm::FromVectorQuery(key) => {
+                let Some(keys) = self.vector_results.edges.get(key) else {
+                    return Box::new(EmptyQuery);
+                };
+
+                Box::new(BooleanQuery::union(
+                    keys.iter()
+                        .map(|(value, score)| {
+                            let q: Box<dyn Query> = Box::new(BoostQuery::new(
+                                Box::new(TermQuery::new(
+                                    tantivy::Term::from_field_text(self.schema.label, &self.schema.normalize(value)),
+                                    IndexRecordOption::Basic,
+                                )),
+                                *score,
+                            ));
+                            q
+                        })
+                        .collect(),
+                ))
+            }
+        }
     }
 
     fn has_relation_type(&self, relation_type: RelationType) -> Box<dyn Query> {
@@ -703,7 +766,8 @@ impl TryFrom<&nidx_protos::graph_query::Node> for Node {
             // Default to exact match to keep backwards compatibility
             let match_kind = node_pb
                 .match_kind
-                .unwrap_or(nidx_protos::graph_query::node::MatchKind::Exact(
+                .as_ref()
+                .unwrap_or(&nidx_protos::graph_query::node::MatchKind::Exact(
                     nidx_protos::graph_query::node::ExactMatch {
                         kind: MatchLocation::Full as i32,
                     },
@@ -745,6 +809,7 @@ impl TryFrom<&nidx_protos::graph_query::Node> for Node {
                         is_prefix: true,
                     }),
                 },
+                nidx_protos::graph_query::node::MatchKind::Vector(_) => Term::FromVectorQuery(value),
             }
         });
 
@@ -764,8 +829,22 @@ impl TryFrom<&nidx_protos::graph_query::Relation> for Relation {
 
     fn try_from(relation_pb: &nidx_protos::graph_query::Relation) -> Result<Self, Self::Error> {
         let value = relation_pb.value.clone();
+        let term = value.map(|value| {
+            match relation_pb
+                .match_kind
+                .as_ref()
+                .unwrap_or(&nidx_protos::graph_query::relation::MatchKind::Exact(
+                    nidx_protos::graph_query::relation::ExactMatch {},
+                )) {
+                nidx_protos::graph_query::relation::MatchKind::Exact(_) => RelationTerm::Exact(value),
+                nidx_protos::graph_query::relation::MatchKind::Vector(_) => RelationTerm::FromVectorQuery(value),
+            }
+        });
         let relation_type = relation_pb.relation_type.map(RelationType::try_from).transpose()?;
 
-        Ok(Relation { value, relation_type })
+        Ok(Relation {
+            label: term,
+            relation_type,
+        })
     }
 }

{nidx_binding-6.10.0.post725 → nidx_binding-6.10.0.post732}/nidx_relation/src/lib.rs

@@ -57,6 +57,8 @@ use tantivy::{
 use tracing::{error, instrument};
 use uuid::Uuid;
 
+use crate::graph_query_parser::VectorQueryResults;
+
 /// Minimum length for a word to be accepted as a entity to search for
 /// suggestions. Low values can provide too much noise and higher ones can
 /// remove important words from suggestion
@@ -211,8 +213,9 @@ impl RelationSearcher {
         &self,
         request: &GraphSearchRequest,
         prefilter: &PrefilterResult,
+        context: VectorQueryResults,
     ) -> anyhow::Result<GraphSearchResponse> {
-        self.reader.graph_search(request, prefilter)
+        self.reader.graph_search(request, prefilter, context)
     }
 
     #[instrument(name = "relation::suggest", skip_all)]
@@ -253,7 +256,7 @@ impl RelationSearcher {
             }),
             ..Default::default()
         };
-        let response = self.graph_search(&request, &PrefilterResult::All)?;
+        let response = self.graph_search(&request, &PrefilterResult::All, VectorQueryResults::default())?;
         Ok(response.nodes)
     }
 

{nidx_binding-6.10.0.post725 → nidx_binding-6.10.0.post732}/nidx_relation/src/reader.rs

@@ -30,7 +30,7 @@ use tantivy::{Index, IndexReader};
 use uuid::Uuid;
 
 use crate::graph_collector::{Selector, TopUniqueCollector};
-use crate::graph_query_parser::{BoolGraphQuery, BoolNodeQuery, GraphQueryParser};
+use crate::graph_query_parser::{BoolGraphQuery, BoolNodeQuery, GraphQueryParser, VectorQueryResults};
 use crate::schema::{Schema, decode_node, decode_relation, encode_field_id};
 use crate::top_unique_n::TopUniqueN;
 use crate::{RelationConfig, io_maps};
@@ -99,6 +99,7 @@ impl RelationsReaderService {
         &self,
         request: &GraphSearchRequest,
         prefilter: &PrefilterResult,
+        context: VectorQueryResults,
     ) -> anyhow::Result<GraphSearchResponse> {
         // No query? Empty graph
         let Some(query) = &request.query else {
@@ -110,21 +111,23 @@ impl RelationsReaderService {
 
         let top_k = request.top_k as usize;
 
+        let parser = GraphQueryParser::new(&self.schema, context);
+
         match request.kind() {
-            QueryKind::Path => self.paths_graph_search(query, prefilter, top_k),
-            QueryKind::Nodes => self.nodes_graph_search(query, prefilter, top_k),
-            QueryKind::Relations => self.relations_graph_search(query, prefilter, top_k),
+            QueryKind::Path => self.paths_graph_search(parser, query, prefilter, top_k),
+            QueryKind::Nodes => self.nodes_graph_search(parser, query, prefilter, top_k),
+            QueryKind::Relations => self.relations_graph_search(parser, query, prefilter, top_k),
         }
     }
 
     fn paths_graph_search(
         &self,
+        parser: GraphQueryParser,
         query: &nidx_protos::graph_query::PathQuery,
         prefilter: &PrefilterResult,
         top_k: usize,
     ) -> anyhow::Result<GraphSearchResponse> {
         let query = BoolGraphQuery::try_from(query)?;
-        let parser = GraphQueryParser::new(&self.schema);
         let index_query = parser.parse_bool(query);
         let index_query = self.apply_prefilter(index_query, prefilter);
 
@@ -135,8 +138,9 @@ impl RelationsReaderService {
         let mut nodes = Vec::new();
         let mut relations = Vec::new();
         let mut graph = Vec::new();
+        let mut scores = Vec::new();
 
-        for (_score, doc_address) in matching_docs {
+        for (score, doc_address) in matching_docs {
             let doc = searcher.doc(doc_address)?;
 
             let source = io_maps::source_to_relation_node(&self.schema, &doc);
@@ -157,25 +161,28 @@ impl RelationsReaderService {
                 metadata: io_maps::decode_metadata(&self.schema, &doc),
                 resource_field_id: io_maps::doc_to_resource_field_id(&self.schema, &doc),
                 facets: io_maps::doc_to_facets(&self.schema, &doc),
-            })
+            });
+
+            scores.push(score);
         }
 
         let response = nidx_protos::GraphSearchResponse {
             nodes,
             relations,
             graph,
+            scores,
         };
         Ok(response)
     }
 
     fn nodes_graph_search(
         &self,
+        parser: GraphQueryParser,
         query: &nidx_protos::graph_query::PathQuery,
         prefilter: &PrefilterResult,
         top_k: usize,
     ) -> anyhow::Result<GraphSearchResponse> {
         let query = BoolNodeQuery::try_from(query)?;
-        let parser = GraphQueryParser::new(&self.schema);
         let (source_query, destination_query) = parser.parse_bool_node(query);
         let source_query = self.apply_prefilter(source_query, prefilter);
         let destination_query = self.apply_prefilter(destination_query, prefilter);
@@ -192,21 +199,21 @@ impl RelationsReaderService {
         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, _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,
-                }
-            })
-            .collect();
+        let mut scores = vec![];
+        let mut nodes = vec![];
+        for (encoded_node, score) in unique_nodes.into_sorted_vec() {
+            let (value, node_type, node_subtype) = decode_node(&encoded_node);
+            scores.push(score);
+            nodes.push(RelationNode {
+                value,
+                ntype: io_maps::u64_to_node_type(node_type),
+                subtype: node_subtype,
+            });
+        }
 
         let response = nidx_protos::GraphSearchResponse {
             nodes,
+            scores,
             ..Default::default()
         };
         Ok(response)
@@ -214,12 +221,12 @@ impl RelationsReaderService {
 
     fn relations_graph_search(
         &self,
+        parser: GraphQueryParser,
         query: &nidx_protos::graph_query::PathQuery,
         prefilter: &PrefilterResult,
         top_k: usize,
     ) -> anyhow::Result<GraphSearchResponse> {
         let query = BoolGraphQuery::try_from(query)?;
-        let parser = GraphQueryParser::new(&self.schema);
         let index_query = parser.parse_bool(query);
         let index_query = self.apply_prefilter(index_query, prefilter);
 
@@ -228,20 +235,21 @@ impl RelationsReaderService {
         let collector = TopUniqueCollector::new(Selector::Relations, top_k);
         let top_relations = searcher.search(&index_query, &collector)?;
 
-        let relations = top_relations
-            .into_sorted_vec()
-            .into_iter()
-            .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),
-                    label: relation_label,
-                }
-            })
-            .collect();
+        let mut scores = vec![];
+        let mut relations = vec![];
+
+        for (encoded_relation, score) in top_relations.into_sorted_vec() {
+            let (relation_type, relation_label) = decode_relation(&encoded_relation);
+            scores.push(score);
+            relations.push(nidx_protos::graph_search_response::Relation {
+                relation_type: io_maps::u64_to_relation_type::<i32>(relation_type),
+                label: relation_label,
+            });
+        }
 
         let response = nidx_protos::GraphSearchResponse {
             relations,
+            scores,
             ..Default::default()
         };
         Ok(response)