From ef5dbdeae6353cad128bb3eebd6abaede37ded97 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20Coss=C3=ADo?= Date: Tue, 22 Apr 2025 13:16:43 -0400 Subject: [PATCH] [full-text index] return earlier in text query parsing (#6264) --- .../src/index/field_index/field_index_base.rs | 4 +- .../immutable_inverted_index.rs | 25 ++++------ .../full_text_index/inverted_index.rs | 47 +++++++++++++------ .../mmap_inverted_index/mod.rs | 27 +++++------ .../full_text_index/mutable_inverted_index.rs | 9 ++-- .../field_index/full_text_index/tests/mod.rs | 11 ++--- .../field_index/full_text_index/text_index.rs | 25 ++++++---- .../condition_converter/match_converter.rs | 4 +- 8 files changed, 80 insertions(+), 72 deletions(-) diff --git a/lib/segment/src/index/field_index/field_index_base.rs b/lib/segment/src/index/field_index/field_index_base.rs index 002cf25f91..8569825c46 100644 --- a/lib/segment/src/index/field_index/field_index_base.rs +++ b/lib/segment/src/index/field_index/field_index_base.rs @@ -178,7 +178,9 @@ impl FieldIndex { FieldIndex::BoolIndex(_) => None, FieldIndex::FullTextIndex(full_text_index) => match &condition.r#match { Some(Match::Text(MatchText { text })) => { - let query = full_text_index.parse_query(text, hw_counter); + let Some(query) = full_text_index.parse_query(text, hw_counter) else { + return Some(false); + }; for value in FullTextIndex::get_values(payload_value) { let document = full_text_index.parse_document(&value, hw_counter); if query.check_match(&document) { diff --git a/lib/segment/src/index/field_index/full_text_index/immutable_inverted_index.rs b/lib/segment/src/index/field_index/full_text_index/immutable_inverted_index.rs index ac348297cb..1a32783460 100644 --- a/lib/segment/src/index/field_index/full_text_index/immutable_inverted_index.rs +++ b/lib/segment/src/index/field_index/full_text_index/immutable_inverted_index.rs @@ -52,13 +52,9 @@ impl InvertedIndex for ImmutableInvertedIndex { let postings_opt: Option> = query .tokens .iter() - .map(|&token_id| match token_id { - None => None, - // if a ParsedQuery token was given an index, then it must exist in the vocabulary - Some(idx) => { - let postings = self.postings.get(idx as usize); - postings - } + .map(|&token_id| { + let postings = self.postings.get(token_id as usize); + postings }) .collect(); @@ -100,23 +96,20 @@ impl InvertedIndex for ImmutableInvertedIndex { point_id: PointOffsetType, _: &HardwareCounterCell, ) -> bool { - if parsed_query.tokens.contains(&None) { + if parsed_query.tokens.is_empty() { return false; } + // check presence of the document if self.values_is_empty(point_id) { return false; } // Check that all tokens are in document - parsed_query - .tokens - .iter() - // unwrap crash safety: all tokens exist in the vocabulary if it passes the above check - .all(|query_token| { - let postings = &self.postings[query_token.unwrap() as usize]; - postings.reader().contains(point_id) - }) + parsed_query.tokens.iter().all(|token_id| { + let postings = &self.postings[*token_id as usize]; + postings.reader().contains(point_id) + }) } fn values_is_empty(&self, point_id: PointOffsetType) -> bool { diff --git a/lib/segment/src/index/field_index/full_text_index/inverted_index.rs b/lib/segment/src/index/field_index/full_text_index/inverted_index.rs index 53b29bb219..e3b04d491e 100644 --- a/lib/segment/src/index/field_index/full_text_index/inverted_index.rs +++ b/lib/segment/src/index/field_index/full_text_index/inverted_index.rs @@ -33,27 +33,26 @@ impl Document { &self.tokens } - pub fn check(&self, token: TokenId) -> bool { - self.tokens.binary_search(&token).is_ok() + pub fn check(&self, token: &TokenId) -> bool { + self.tokens.binary_search(token).is_ok() } } #[derive(Debug, Clone)] pub struct ParsedQuery { - pub tokens: Vec>, + pub tokens: Vec, } impl ParsedQuery { pub fn check_match(&self, document: &Document) -> bool { - if self.tokens.contains(&None) { + if self.tokens.is_empty() { return false; } // Check that all tokens are in document self.tokens .iter() - // unwrap crash safety: all tokens exist in the vocabulary if it passes the above check - .all(|query_token| document.check(query_token.unwrap())) + .all(|query_token| document.check(query_token)) } } @@ -108,10 +107,7 @@ pub trait InvertedIndex { let posting_lengths: Option> = query .tokens .iter() - .map(|&vocab_idx| match vocab_idx { - None => None, - Some(idx) => self.get_posting_len(idx, hw_counter), - }) + .map(|&vocab_idx| self.get_posting_len(vocab_idx, hw_counter)) .collect(); if posting_lengths.is_none() || points_count == 0 { // There are unseen tokens -> no matches @@ -228,12 +224,16 @@ mod tests { (0..len).map(|_| generate_word()).collect() } + /// Tries to parse a query. If there is an unknown id to a token, returns `None` fn to_parsed_query( query: Vec, token_to_id: impl Fn(String) -> Option, - ) -> ParsedQuery { - let tokens: Vec<_> = query.into_iter().map(token_to_id).collect(); - ParsedQuery { tokens } + ) -> Option { + let tokens = query + .into_iter() + .map(token_to_id) + .collect::>>()?; + Some(ParsedQuery { tokens }) } fn mutable_inverted_index(indexed_count: u32, deleted_count: u32) -> MutableInvertedIndex { @@ -384,11 +384,20 @@ mod tests { }) .collect(); - for (mut_query, imm_query) in mut_parsed_queries + for queries in mut_parsed_queries .iter() .cloned() .zip(imm_parsed_queries.iter().cloned()) { + let (Some(mut_query), Some(imm_query)) = queries else { + // Immutable index can have a smaller vocabulary, since it only contains tokens that have + // non-empty posting lists. + // Since we removed some documents from the mutable index, it can happen that the immutable + // index returns None when parsing the query, even if the mutable index returns Some. + // + // In this case both queries would filter to an empty set of documents. + continue; + }; let mut_filtered = mutable.filter(mut_query, &hw_counter).collect::>(); let imm_filtered = mmap_index .filter(imm_query, &hw_counter) @@ -410,11 +419,19 @@ mod tests { // Check congruence after deletion - for (mut_query, imm_query) in mut_parsed_queries + for queries in mut_parsed_queries .iter() .cloned() .zip(imm_parsed_queries.iter().cloned()) { + let (Some(mut_query), Some(imm_query)) = queries else { + // Both queries must be None + assert!( + queries.0.is_none() && queries.1.is_none(), + "Both queries must be parsed or not parsed entirely" + ); + continue; + }; let mut_filtered = mutable.filter(mut_query, &hw_counter).collect::>(); let imm_filtered = mmap_index .filter(imm_query, &hw_counter) diff --git a/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/mod.rs b/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/mod.rs index d8950fd87c..d2b4b44afc 100644 --- a/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/mod.rs +++ b/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/mod.rs @@ -206,11 +206,7 @@ impl InvertedIndex for MmapInvertedIndex { let postings_opt: Option> = query .tokens .iter() - .map(|&token_id| match token_id { - None => None, - // if a ParsedQuery token was given an index, then it must exist in the vocabulary - Some(idx) => self.postings.get(idx, hw_counter), - }) + .map(|&token_id| self.postings.get(token_id, hw_counter)) .collect(); let Some(posting_readers) = postings_opt else { // There are unseen tokens -> no matches @@ -252,24 +248,23 @@ impl InvertedIndex for MmapInvertedIndex { point_id: PointOffsetType, hw_counter: &HardwareCounterCell, ) -> bool { - if parsed_query.tokens.contains(&None) { + // check non-empty query + if parsed_query.tokens.is_empty() { return false; } + // check presence of the document if self.values_is_empty(point_id) { return false; } // Check that all tokens are in document - parsed_query - .tokens - .iter() - // unwrap safety: all tokens exist in the vocabulary if it passes the above check - .all(|query_token| { - self.postings - .get(query_token.unwrap(), hw_counter) - .unwrap() - .contains(point_id) - }) + parsed_query.tokens.iter().all(|query_token| { + self.postings + .get(*query_token, hw_counter) + // unwrap safety: all tokens exist in the vocabulary, otherwise there'd be no query tokens + .unwrap() + .contains(point_id) + }) } fn values_is_empty(&self, point_id: PointOffsetType) -> bool { diff --git a/lib/segment/src/index/field_index/full_text_index/mutable_inverted_index.rs b/lib/segment/src/index/field_index/full_text_index/mutable_inverted_index.rs index b0a8a026db..44c9a33767 100644 --- a/lib/segment/src/index/field_index/full_text_index/mutable_inverted_index.rs +++ b/lib/segment/src/index/field_index/full_text_index/mutable_inverted_index.rs @@ -150,14 +150,11 @@ impl InvertedIndex for MutableInvertedIndex { let postings_opt: Option> = query .tokens .iter() - .map(|&vocab_idx| match vocab_idx { - None => None, + .map(|&vocab_idx| { // if a ParsedQuery token was given an index, then it must exist in the vocabulary // dictionary. Posting list entry can be None but it exists. - Some(idx) => { - let postings = self.postings.get(idx as usize).unwrap().as_ref(); - postings - } + let postings = self.postings.get(vocab_idx as usize).unwrap().as_ref(); + postings }) .collect(); if postings_opt.is_none() { diff --git a/lib/segment/src/index/field_index/full_text_index/tests/mod.rs b/lib/segment/src/index/field_index/full_text_index/tests/mod.rs index 4b30ec069c..280f08b766 100644 --- a/lib/segment/src/index/field_index/full_text_index/tests/mod.rs +++ b/lib/segment/src/index/field_index/full_text_index/tests/mod.rs @@ -188,7 +188,7 @@ fn test_prefix_search(#[case] immutable: bool) { let res: Vec<_> = index.query("ROBO", &hw_counter).collect(); - let query = index.parse_query("ROBO", &hw_counter); + let query = index.parse_query("ROBO", &hw_counter).unwrap(); for idx in res.iter().copied() { assert!(index.check_match(&query, idx, &hw_counter)); @@ -197,12 +197,7 @@ fn test_prefix_search(#[case] immutable: bool) { assert_eq!(res.len(), 3); let res: Vec<_> = index.query("q231", &hw_counter).collect(); + assert!(res.is_empty()); - let query = index.parse_query("q231", &hw_counter); - - for idx in [1, 2, 3] { - assert!(!index.check_match(&query, idx, &hw_counter)); - } - - assert_eq!(res.len(), 0); + assert!(index.parse_query("q231", &hw_counter).is_none()); } diff --git a/lib/segment/src/index/field_index/full_text_index/text_index.rs b/lib/segment/src/index/field_index/full_text_index/text_index.rs index 78a26b5a82..890a6e0325 100644 --- a/lib/segment/src/index/field_index/full_text_index/text_index.rs +++ b/lib/segment/src/index/field_index/full_text_index/text_index.rs @@ -1,7 +1,8 @@ -use std::collections::{BTreeSet, HashSet}; +use std::collections::BTreeSet; use std::path::PathBuf; use std::sync::Arc; +use ahash::AHashSet; use common::counter::hardware_counter::HardwareCounterCell; use common::types::PointOffsetType; use parking_lot::RwLock; @@ -236,14 +237,14 @@ impl FullTextIndex { } } - pub fn parse_query(&self, text: &str, hw_counter: &HardwareCounterCell) -> ParsedQuery { - let mut tokens = HashSet::new(); + /// Tries to parse a query. If there are any unseen tokens, returns `None` + pub fn parse_query(&self, text: &str, hw_counter: &HardwareCounterCell) -> Option { + let mut tokens = AHashSet::new(); Tokenizer::tokenize_query(text, self.config(), |token| { tokens.insert(self.get_token(token, hw_counter)); }); - ParsedQuery { - tokens: tokens.into_iter().collect(), - } + let tokens = tokens.into_iter().collect::>>()?; + Some(ParsedQuery { tokens }) } pub fn parse_document(&self, text: &str, hw_counter: &HardwareCounterCell) -> Document { @@ -262,7 +263,9 @@ impl FullTextIndex { query: &'a str, hw_counter: &'a HardwareCounterCell, ) -> Box + 'a> { - let parsed_query = self.parse_query(query, hw_counter); + let Some(parsed_query) = self.parse_query(query, hw_counter) else { + return Box::new(std::iter::empty()); + }; self.filter(parsed_query, hw_counter) } @@ -398,7 +401,9 @@ impl PayloadFieldIndex for FullTextIndex { hw_counter: &'a HardwareCounterCell, ) -> Option + 'a>> { if let Some(Match::Text(text_match)) = &condition.r#match { - let parsed_query = self.parse_query(&text_match.text, hw_counter); + let Some(parsed_query) = self.parse_query(&text_match.text, hw_counter) else { + return Some(Box::new(std::iter::empty())); + }; return Some(self.filter(parsed_query, hw_counter)); } None @@ -410,7 +415,9 @@ impl PayloadFieldIndex for FullTextIndex { hw_counter: &HardwareCounterCell, ) -> Option { if let Some(Match::Text(text_match)) = &condition.r#match { - let parsed_query = self.parse_query(&text_match.text, hw_counter); + let Some(parsed_query) = self.parse_query(&text_match.text, hw_counter) else { + return Some(CardinalityEstimation::exact(0)); + }; return Some(self.estimate_cardinality(&parsed_query, condition, hw_counter)); } None diff --git a/lib/segment/src/index/query_optimization/condition_converter/match_converter.rs b/lib/segment/src/index/query_optimization/condition_converter/match_converter.rs index 9199a9f6e2..55b45d8379 100644 --- a/lib/segment/src/index/query_optimization/condition_converter/match_converter.rs +++ b/lib/segment/src/index/query_optimization/condition_converter/match_converter.rs @@ -254,7 +254,9 @@ fn get_match_text_checker( let hw_counter = hw_acc.get_counter_cell(); match index { FieldIndex::FullTextIndex(full_text_index) => { - let parsed_query = full_text_index.parse_query(&text, &hw_counter); + let Some(parsed_query) = full_text_index.parse_query(&text, &hw_counter) else { + return Some(Box::new(|_| false)); + }; Some(Box::new(move |point_id: PointOffsetType| { full_text_index.check_match(&parsed_query, point_id, &hw_counter) }))