Share prefiltering struct (#2858)

nuclia · Feb 12, 2025 · 42c4a36 · 42c4a36
1 parent 42bed9c
commit 42c4a36
Show file tree

Hide file tree

Showing 28 changed files with 334 additions and 356 deletions.
diff --git a/nidx/Cargo.lock b/nidx/Cargo.lock
diff --git a/nidx/nidx_paragraph/src/lib.rs b/nidx/nidx_paragraph/src/lib.rs
@@ -34,7 +34,7 @@ use nidx_tantivy::{
     index_reader::{open_index_with_deletions, DeletionQueryBuilder},
     TantivyIndexer, TantivyMeta, TantivySegmentMetadata,
 };
-use nidx_types::OpenIndexMetadata;
+use nidx_types::{prefilter::PrefilterResult, OpenIndexMetadata};
 use reader::ParagraphReaderService;
 use resource_indexer::index_paragraphs;
 use schema::ParagraphSchema;
@@ -140,13 +140,21 @@ impl ParagraphSearcher {
     }
 
     #[instrument(name = "paragraph::search", skip_all)]
-    pub fn search(&self, request: &ParagraphSearchRequest) -> anyhow::Result<ParagraphSearchResponse> {
-        self.reader.search(request)
+    pub fn search(
+        &self,
+        request: &ParagraphSearchRequest,
+        prefilter: &PrefilterResult,
+    ) -> anyhow::Result<ParagraphSearchResponse> {
+        self.reader.search(request, prefilter)
     }
 
     #[instrument(name = "paragraph::suggest", skip_all)]
-    pub fn suggest(&self, request: &SuggestRequest) -> anyhow::Result<ParagraphSearchResponse> {
-        self.reader.suggest(request)
+    pub fn suggest(
+        &self,
+        request: &SuggestRequest,
+        prefilter: &PrefilterResult,
+    ) -> anyhow::Result<ParagraphSearchResponse> {
+        self.reader.suggest(request, prefilter)
     }
 
     pub fn iterator(&self, request: &StreamRequest) -> anyhow::Result<impl Iterator<Item = ParagraphItem>> {

diff --git a/nidx/nidx_paragraph/src/reader.rs b/nidx/nidx_paragraph/src/reader.rs
@@ -23,6 +23,7 @@ use std::time::Instant;
 
 use nidx_protos::order_by::{OrderField, OrderType};
 use nidx_protos::{OrderBy, ParagraphItem, ParagraphSearchResponse, StreamRequest, SuggestRequest};
+use nidx_types::prefilter::PrefilterResult;
 use tantivy::collector::{Collector, Count, FacetCollector, TopDocs};
 use tantivy::query::{AllQuery, Query, QueryParser};
 use tantivy::{schema::*, DateTime, Order};
@@ -56,7 +57,11 @@ impl ParagraphReaderService {
         Ok(count)
     }
 
-    pub fn suggest(&self, request: &SuggestRequest) -> anyhow::Result<ParagraphSearchResponse> {
+    pub fn suggest(
+        &self,
+        request: &SuggestRequest,
+        prefilter: &PrefilterResult,
+    ) -> anyhow::Result<ParagraphSearchResponse> {
         let time = Instant::now();
         let id = Some(&request.shard);
 
@@ -65,7 +70,8 @@ impl ParagraphReaderService {
 
         let parser = QueryParser::for_index(&self.index, vec![self.schema.text]);
         let text = self.adapt_text(&parser, &request.body);
-        let (original, termc, fuzzied) = suggest_query(&parser, &text, request, &self.schema, FUZZY_DISTANCE);
+        let (original, termc, fuzzied) =
+            suggest_query(&parser, &text, request, prefilter, &self.schema, FUZZY_DISTANCE);
         let v = time.elapsed().as_millis();
         debug!("{id:?} - Creating query: ends at {v} ms");
 
@@ -120,7 +126,11 @@ impl ParagraphReaderService {
         Ok(producer.flatten())
     }
 
-    pub fn search(&self, request: &ParagraphSearchRequest) -> anyhow::Result<ParagraphSearchResponse> {
+    pub fn search(
+        &self,
+        request: &ParagraphSearchRequest,
+        prefilter: &PrefilterResult,
+    ) -> anyhow::Result<ParagraphSearchResponse> {
         let time = Instant::now();
         let id = Some(&request.id);
 
@@ -151,6 +161,7 @@ impl ParagraphReaderService {
             &parser,
             &text,
             request,
+            prefilter,
             &self.schema,
             FUZZY_DISTANCE,
             advanced

diff --git a/nidx/nidx_paragraph/src/request_types.rs b/nidx/nidx_paragraph/src/request_types.rs
@@ -24,7 +24,6 @@ use nidx_types::query_language::BooleanExpression;
 pub struct ParagraphSearchRequest {
     pub id: String,
     pub uuid: String,
-    pub fields: Vec<String>,
     /// query this text in all the paragraphs
     pub body: String,
     pub filter: Option<nidx_protos::Filter>,
@@ -37,7 +36,6 @@ pub struct ParagraphSearchRequest {
     pub with_duplicates: bool,
     pub only_faceted: bool,
     pub advanced_query: Option<String>,
-    pub key_filters: Vec<String>,
     pub min_score: f32,
     pub security: Option<nidx_protos::utils::Security>,
 

diff --git a/nidx/nidx_paragraph/src/search_query.rs b/nidx/nidx_paragraph/src/search_query.rs
@@ -23,6 +23,7 @@ use crate::set_query::SetQuery;
 use itertools::Itertools;
 use nidx_protos::prost_types::Timestamp as ProstTimestamp;
 use nidx_protos::{StreamRequest, SuggestRequest};
+use nidx_types::prefilter::PrefilterResult;
 use std::borrow::Cow;
 use std::collections::{HashMap, HashSet};
 use std::ops::Bound;
@@ -280,10 +281,27 @@ pub fn produce_date_range_query(
     Some(query)
 }
 
+fn apply_prefilter(
+    queries: &mut [&mut Vec<(Occur, Box<dyn Query>)>],
+    schema: &ParagraphSchema,
+    prefilter: &PrefilterResult,
+) {
+    if let PrefilterResult::Some(field_keys) = prefilter {
+        let set_query = Box::new(SetQuery::new(
+            schema.field_uuid,
+            field_keys.iter().map(|x| format!("{}{}", x.resource_id.simple(), x.field_id)),
+        ));
+        for q in queries {
+            q.push((Occur::Must, set_query.clone()));
+        }
+    }
+}
+
 pub fn suggest_query(
     parser: &QueryParser,
     text: &str,
     request: &SuggestRequest,
+    prefilter: &PrefilterResult,
     schema: &ParagraphSchema,
     distance: u8,
 ) -> (Box<dyn Query>, SharedTermC, Box<dyn Query>) {
@@ -326,19 +344,7 @@ pub fn suggest_query(
             originals.push((Occur::Must, Box::new(facet_term_query)));
         });
 
-    if !request.key_filters.is_empty() {
-        let (field_ids, resource_ids) = request.key_filters.iter().cloned().partition::<Vec<_>, _>(|k| k.contains('/'));
-        if !field_ids.is_empty() {
-            let set_query = Box::new(SetQuery::new(schema.field_uuid, field_ids));
-            fuzzies.push((Occur::Must, set_query.clone()));
-            originals.push((Occur::Must, set_query.clone()));
-        }
-        if !resource_ids.is_empty() {
-            let set_query = Box::new(SetQuery::new(schema.uuid, resource_ids));
-            fuzzies.push((Occur::Must, set_query.clone()));
-            originals.push((Occur::Must, set_query.clone()));
-        }
-    }
+    apply_prefilter(&mut [&mut fuzzies, &mut originals], schema, prefilter);
 
     if originals.len() == 1 && originals[0].1.is::<AllQuery>() {
         let original = originals.pop().unwrap().1;
@@ -358,6 +364,7 @@ pub fn search_query(
     parser: &QueryParser,
     text: &str,
     search: &ParagraphSearchRequest,
+    prefilter: &PrefilterResult,
     schema: &ParagraphSchema,
     distance: u8,
     with_advance: Option<Box<dyn Query>>,
@@ -401,23 +408,6 @@ pub fn search_query(
             originals.push((Occur::Must, Box::new(created)));
         }
     }
-    // Fields
-    let mut field_filter: Vec<(Occur, Box<dyn Query>)> = vec![];
-    search
-        .fields
-        .iter()
-        .map(|value| format!("/{value}"))
-        .flat_map(|facet_key| Facet::from_text(&facet_key).ok().into_iter())
-        .for_each(|facet| {
-            let facet_term = Term::from_facet(schema.field, &facet);
-            let facet_term_query = TermQuery::new(facet_term, IndexRecordOption::Basic);
-            field_filter.push((Occur::Should, Box::new(facet_term_query)));
-        });
-    if !field_filter.is_empty() {
-        let field_filter = Box::new(BooleanQuery::new(field_filter));
-        fuzzies.push((Occur::Must, field_filter.clone()));
-        originals.push((Occur::Must, field_filter));
-    }
 
     // Label filters
     if let Some(formula) = &search.filtering_formula {
@@ -426,19 +416,7 @@ pub fn search_query(
         originals.push((Occur::Must, query));
     }
 
-    if !search.key_filters.is_empty() {
-        let (field_ids, resource_ids) = search.key_filters.iter().cloned().partition::<Vec<_>, _>(|k| k.contains('/'));
-        if !field_ids.is_empty() {
-            let set_query = Box::new(SetQuery::new(schema.field_uuid, field_ids));
-            fuzzies.push((Occur::Must, set_query.clone()));
-            originals.push((Occur::Must, set_query.clone()));
-        }
-        if !resource_ids.is_empty() {
-            let set_query = Box::new(SetQuery::new(schema.uuid, resource_ids));
-            fuzzies.push((Occur::Must, set_query.clone()));
-            originals.push((Occur::Must, set_query.clone()));
-        }
-    }
+    apply_prefilter(&mut [&mut fuzzies, &mut originals], schema, prefilter);
 
     if originals.len() == 1 && originals[0].1.is::<AllQuery>() {
         let original = originals.pop().unwrap().1;

diff --git a/nidx/nidx_paragraph/src/set_query.rs b/nidx/nidx_paragraph/src/set_query.rs
@@ -36,8 +36,8 @@ pub struct SetQuery {
 
 impl SetQuery {
     /// Create a Term Set Query
-    pub fn new(field: Field, values: Vec<String>) -> Self {
-        let values = values.into_iter().collect();
+    pub fn new(field: Field, values: impl Iterator<Item = String>) -> Self {
+        let values = values.collect();
         let set = SetWeightWrapper::new(SetWeight {
             field,
             values,