From cfc27c9665f6db1971b4ca4063b058decff3c330 Mon Sep 17 00:00:00 2001 From: Evance Souamoro Date: Thu, 29 Apr 2021 11:49:27 +0000 Subject: [PATCH 1/7] add support for more like this query --- src/query/mlt/mlt.rs | 387 +++++++++++++++++++++++++++++++++++++++++ src/query/mlt/mod.rs | 5 + src/query/mlt/query.rs | 284 ++++++++++++++++++++++++++++++ src/query/mod.rs | 2 + 4 files changed, 678 insertions(+) create mode 100644 src/query/mlt/mlt.rs create mode 100644 src/query/mlt/mod.rs create mode 100644 src/query/mlt/query.rs diff --git a/src/query/mlt/mlt.rs b/src/query/mlt/mlt.rs new file mode 100644 index 000000000..9839f61b8 --- /dev/null +++ b/src/query/mlt/mlt.rs @@ -0,0 +1,387 @@ +use std::collections::{BinaryHeap, HashMap}; + +use crate::{ + query::{BooleanQuery, BoostQuery, Occur, Query, TermQuery}, + schema::{Field, FieldType, FieldValue, IndexRecordOption, Term, Value}, + tokenizer::{BoxTokenStream, FacetTokenizer, PreTokenizedStream, Tokenizer}, + DocAddress, Result, Searcher, TantivyError, +}; + +#[derive(Debug, PartialEq)] +struct ScoreTerm { + pub term: Term, + pub score: f32, +} + +impl ScoreTerm { + fn new(term: Term, score: f32) -> Self { + Self { term, score } + } +} + +impl Eq for ScoreTerm {} + +impl PartialOrd for ScoreTerm { + fn partial_cmp(&self, other: &Self) -> Option { + self.score.partial_cmp(&other.score) + } +} + +impl Ord for ScoreTerm { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.partial_cmp(other).unwrap_or(std::cmp::Ordering::Equal) + } +} + +/// A struct used as helper to build [`MoreLikeThisQuery`] +#[derive(Debug, Clone)] +pub struct MoreLikeThis { + /// Ignore words which do not occur in at least this many docs. + pub min_doc_frequency: Option, + /// Ignore words which occur in more than this many docs. + pub max_doc_frequency: Option, + /// Ignore words less frequent than this. + pub min_term_frequency: Option, + /// Don't return a query longer than this. + pub max_query_terms: Option, + /// Ignore words if less than this length. + pub min_word_length: Option, + /// Ignore words if greater than this length. + pub max_word_length: Option, + /// Boost factor to use when boosting the terms + pub boost_factor: Option, + /// Current set of stop words. + pub stop_words: Vec, +} + +impl Default for MoreLikeThis { + fn default() -> Self { + Self { + min_doc_frequency: Some(5), + max_doc_frequency: None, + min_term_frequency: Some(2), + max_query_terms: Some(25), + min_word_length: None, + max_word_length: None, + boost_factor: Some(1.0), + stop_words: vec![], + } + } +} + +impl MoreLikeThis { + /// Creates a [`BooleanQuery`] using a document address to collect + /// the top stored field values. + pub fn query_with_document( + &self, + searcher: &Searcher, + doc_address: DocAddress, + ) -> Result { + let score_terms = self.retrieve_terms_from_doc_address(searcher, doc_address)?; + let query = self.create_query(score_terms); + Ok(query) + } + + /// Creates a [`BooleanQuery`] using a set of field values. + pub fn query_with_document_fields( + &self, + searcher: &Searcher, + doc_fields: &[(Field, Vec)], + ) -> Result { + let score_terms = self.retrieve_terms_from_doc_fields(searcher, doc_fields)?; + let query = self.create_query(score_terms); + Ok(query) + } + + /// Creates a [`BooleanQuery`] from an ascendingly sorted list of ScoreTerm + /// This will map the list of ScoreTerm to a list of [`TermQuery`] and compose a + /// BooleanQuery using that list as sub queries. + fn create_query(&self, score_terms: Vec) -> BooleanQuery { + let best_score = score_terms.first().map_or(1f32, |x| x.score); + let mut queries = Vec::new(); + + for ScoreTerm { term, score } in score_terms { + let mut query: Box = + Box::new(TermQuery::new(term, IndexRecordOption::Basic)); + if let Some(factor) = self.boost_factor { + query = Box::new(BoostQuery::new(query, score * factor / best_score)); + } + queries.push((Occur::Should, query)); + } + BooleanQuery::from(queries) + } + + /// Finds terms for a more-like-this query. + /// doc_address is the address of document from which to find terms. + fn retrieve_terms_from_doc_address( + &self, + searcher: &Searcher, + doc_address: DocAddress, + ) -> Result> { + let doc = searcher.doc(doc_address)?; + let field_to_field_values = doc + .get_sorted_field_values() + .iter() + .map(|(field, values)| { + ( + *field, + values + .iter() + .map(|v| (**v).clone()) + .collect::>(), + ) + }) + .collect::>(); + self.retrieve_terms_from_doc_fields(searcher, &field_to_field_values) + } + + /// Finds terms for a more-like-this query. + /// field_to_field_values is a mapping from field to possible values of taht field. + fn retrieve_terms_from_doc_fields( + &self, + searcher: &Searcher, + field_to_field_values: &[(Field, Vec)], + ) -> Result> { + if field_to_field_values.is_empty() { + return Err(TantivyError::InvalidArgument("Cannot create more like this query on empty field values. The document may not have stored fields".to_string())); + } + + let mut field_to_term_freq_map = HashMap::new(); + for (field, field_values) in field_to_field_values { + self.add_term_frequencies(searcher, *field, field_values, &mut field_to_term_freq_map)?; + } + self.create_score_term(searcher, field_to_term_freq_map) + } + + /// Computes the frequency of values for a field while updating the term frequencies + /// Note: A FieldValue can be made up of multiple terms. + /// We are interested in extracting terms within FieldValue + fn add_term_frequencies( + &self, + searcher: &Searcher, + field: Field, + field_values: &[FieldValue], + term_frequencies: &mut HashMap, + ) -> Result<()> { + let schema = searcher.schema(); + let tokenizer_manager = searcher.index().tokenizers(); + + let field_entry = schema.get_field_entry(field); + if !field_entry.is_indexed() { + return Ok(()); + } + + // extract the raw value, possibly tokenizing & filtering to update the term frequency map + match field_entry.field_type() { + FieldType::HierarchicalFacet(_) => { + let facets: Vec<&str> = field_values + .iter() + .map(|field_value| match *field_value.value() { + Value::Facet(ref facet) => Ok(facet.encoded_str()), + _ => Err(TantivyError::InvalidArgument( + "invalid field value".to_string(), + )), + }) + .collect::>>()?; + for fake_str in facets { + FacetTokenizer.token_stream(fake_str).process(&mut |token| { + if self.is_noise_word(token.text.clone()) { + let term = Term::from_field_text(field, &token.text); + *term_frequencies.entry(term).or_insert(0) += 1; + } + }); + } + } + FieldType::Str(text_options) => { + let mut token_streams: Vec = vec![]; + let mut offsets = vec![]; + let mut total_offset = 0; + + for field_value in field_values { + match field_value.value() { + Value::PreTokStr(tok_str) => { + offsets.push(total_offset); + if let Some(last_token) = tok_str.tokens.last() { + total_offset += last_token.offset_to; + } + token_streams.push(PreTokenizedStream::from(tok_str.clone()).into()); + } + Value::Str(ref text) => { + if let Some(tokenizer) = text_options + .get_indexing_options() + .map(|text_indexing_options| { + text_indexing_options.tokenizer().to_string() + }) + .and_then(|tokenizer_name| tokenizer_manager.get(&tokenizer_name)) + { + offsets.push(total_offset); + total_offset += text.len(); + //let v = text.clone(); + token_streams.push(tokenizer.token_stream(text)); + } + } + _ => (), + } + } + + for mut token_stream in token_streams { + token_stream.process(&mut |token| { + if !self.is_noise_word(token.text.clone()) { + let term = Term::from_field_text(field, &token.text); + *term_frequencies.entry(term).or_insert(0) += 1; + } + }); + } + } + FieldType::U64(_) => { + for field_value in field_values { + let val = field_value + .value() + .u64_value() + .ok_or(TantivyError::InvalidArgument("invalid value".to_string()))?; + if !self.is_noise_word(val.to_string()) { + let term = Term::from_field_u64(field, val); + *term_frequencies.entry(term).or_insert(0) += 1; + } + } + } + FieldType::Date(_) => { + for field_value in field_values { + // TODO: Ask if this is the semantic (timestamp) we want + let val = field_value + .value() + .date_value() + .ok_or(TantivyError::InvalidArgument("invalid value".to_string()))? + .timestamp(); + if !self.is_noise_word(val.to_string()) { + let term = Term::from_field_i64(field, val); + *term_frequencies.entry(term).or_insert(0) += 1; + } + } + } + FieldType::I64(_) => { + for field_value in field_values { + let val = field_value + .value() + .i64_value() + .ok_or(TantivyError::InvalidArgument("invalid value".to_string()))?; + if !self.is_noise_word(val.to_string()) { + let term = Term::from_field_i64(field, val); + *term_frequencies.entry(term).or_insert(0) += 1; + } + } + } + FieldType::F64(_) => { + for field_value in field_values { + let val = field_value + .value() + .f64_value() + .ok_or(TantivyError::InvalidArgument("invalid value".to_string()))?; + if !self.is_noise_word(val.to_string()) { + let term = Term::from_field_f64(field, val); + *term_frequencies.entry(term).or_insert(0) += 1; + } + } + } + _ => {} + } + Ok(()) + } + + /// Determines if the term is likely to be of interest based on "more-like-this" settings + fn is_noise_word(&self, word: String) -> bool { + let word_length = word.len(); + if word_length == 0 { + return true; + } + if self + .min_word_length + .map(|min| word_length < min) + .unwrap_or(false) + { + return true; + } + if self + .max_word_length + .map(|max| word_length > max) + .unwrap_or(false) + { + return true; + } + return self.stop_words.contains(&word); + } + + /// Couputes the score for each term while ignoring not useful terms + fn create_score_term( + &self, + searcher: &Searcher, + per_field_term_frequencies: HashMap, + ) -> Result> { + let mut score_terms = BinaryHeap::new(); + let num_docs = searcher + .segment_readers() + .iter() + .map(|x| x.num_docs() as u64) + .sum::(); + + for (term, term_frequency) in per_field_term_frequencies.into_iter() { + // ignore terms with less than min_term_frequency + if self + .min_term_frequency + .map(|x| term_frequency < x) + .unwrap_or(false) + { + continue; + } + + let doc_freq = searcher.doc_freq(&term)?; + + // ignore terms with less than min_doc_frequency + if self + .min_doc_frequency + .map(|x| doc_freq < x) + .unwrap_or(false) + { + continue; + } + + // ignore terms with more than max_doc_frequency + if self + .max_doc_frequency + .map(|x| doc_freq > x) + .unwrap_or(false) + { + continue; + } + + // ignore terms with zero frequency + if doc_freq == 0 { + continue; + } + + // compute similarity & score + let idf = self.idf(doc_freq, num_docs); + let score = (term_frequency as f32) * idf; + score_terms.push(ScoreTerm::new(term, score)); + } + + // limit ourself to max_query terms. we need to sort so to avoid discarding important terms + let score_terms = if let Some(max_query_terms) = self.max_query_terms { + let max_num_terms = std::cmp::min(max_query_terms, score_terms.len()); + score_terms + .into_sorted_vec() + .into_iter() + .take(max_num_terms) + .collect() + } else { + score_terms.into_sorted_vec() + }; + Ok(score_terms) + } + + /// Computes the similarity + fn idf(&self, doc_freq: u64, doc_count: u64) -> f32 { + let x = ((doc_count - doc_freq) as f32 + 0.5) / (doc_freq as f32 + 0.5); + (1f32 + x).ln() + } +} diff --git a/src/query/mlt/mod.rs b/src/query/mlt/mod.rs new file mode 100644 index 000000000..97f541908 --- /dev/null +++ b/src/query/mlt/mod.rs @@ -0,0 +1,5 @@ +mod mlt; +mod query; + +pub use self::mlt::MoreLikeThis; +pub use self::query::{MoreLikeThisQuery}; diff --git a/src/query/mlt/query.rs b/src/query/mlt/query.rs new file mode 100644 index 000000000..3bc5cc833 --- /dev/null +++ b/src/query/mlt/query.rs @@ -0,0 +1,284 @@ +use super::MoreLikeThis; + +use crate::{ + query::{Query, Weight}, + schema::{Field, FieldValue}, + DocAddress, Result, Searcher, TantivyError, +}; + +/// A query that matches all of the documents similar to a document +/// or a set of field values provided. +/// +/// # Examples +/// +/// ``` +/// use tantivy::DocAddress; +/// use tantivy::query::MoreLikeThisQuery; +/// +/// let query = MoreLikeThisQuery::builder() +/// .with_min_doc_frequency(1) +/// .with_max_doc_frequency(10) +/// .with_min_term_frequency(1) +/// .with_min_word_length(2) +/// .with_max_word_length(5) +/// .with_boost_factor(1.0) +/// .with_stop_words(vec!["for".to_string()]) +/// .with_document(DocAddress::new(2, 1)); +/// +/// ``` +#[derive(Debug, Clone)] +pub struct MoreLikeThisQuery { + mlt: MoreLikeThis, + doc_address: Option, + doc_fields: Option)>>, +} + +impl MoreLikeThisQuery { + /// Creates a new builder. + pub fn builder() -> MoreLikeThisQueryBuilder { + MoreLikeThisQueryBuilder::default() + } +} + +impl Query for MoreLikeThisQuery { + fn weight(&self, searcher: &Searcher, scoring_enabled: bool) -> Result> { + if let Some(doc_address) = self.doc_address { + return self + .mlt + .query_with_document(searcher, doc_address)? + .weight(searcher, scoring_enabled); + } + + if let Some(ref doc_fields) = self.doc_fields { + return self + .mlt + .query_with_document_fields(searcher, doc_fields)? + .weight(searcher, scoring_enabled); + } + + Err(TantivyError::InvalidArgument("".to_string())) + } +} + +/// The builder for more-like-this query +#[derive(Debug, Clone)] +pub struct MoreLikeThisQueryBuilder { + mlt: MoreLikeThis, +} + +impl Default for MoreLikeThisQueryBuilder { + fn default() -> Self { + Self { + mlt: MoreLikeThis::default(), + } + } +} + +impl MoreLikeThisQueryBuilder { + /// Sets the minimum document frequency. + /// + /// The resulting query will ignore words which do not occur + /// in at least this many docs. + pub fn with_min_doc_frequency(mut self, value: u64) -> Self { + self.mlt.min_doc_frequency = Some(value); + self + } + + /// Sets the maximum document frequency. + /// + /// The resulting query will ignore words which occur + /// in more than this many docs. + pub fn with_max_doc_frequency(mut self, value: u64) -> Self { + self.mlt.max_doc_frequency = Some(value); + self + } + + /// Sets the minimum term frequency. + /// + /// The resulting query will ignore words less + /// frequent that this number. + pub fn with_min_term_frequency(mut self, value: usize) -> Self { + self.mlt.min_term_frequency = Some(value); + self + } + + /// Sets the maximum query terms. + /// + /// The resulting query will not return a query with more clause than this. + pub fn with_max_query_terms(mut self, value: usize) -> Self { + self.mlt.max_query_terms = Some(value); + self + } + + /// Sets the minimum word length. + /// + /// The resulting query will ignore words shorter than this length. + pub fn with_min_word_length(mut self, value: usize) -> Self { + self.mlt.min_word_length = Some(value); + self + } + + /// Sets the maximum word length. + /// + /// The resulting query will ignore words longer than this length. + pub fn with_max_word_length(mut self, value: usize) -> Self { + self.mlt.max_word_length = Some(value); + self + } + + /// Sets the boost factor + /// + /// The boost factor used by the resulting query for boosting terms. + pub fn with_boost_factor(mut self, value: f32) -> Self { + self.mlt.boost_factor = Some(value); + self + } + + /// Sets the set of stop words + /// + /// The resulting query will ignore these set of words. + pub fn with_stop_words(mut self, value: Vec) -> Self { + self.mlt.stop_words = value; + self + } + + /// Sets the document address + /// Returns the constructed [`MoreLikeThisQuery`] + /// + /// This document will be used to collect field values, extract frequent terms + /// needed for composing the query. + /// + /// Note that field values will only be collected from stored fields in the index. + /// You can construct your own field values from any source. + pub fn with_document(self, doc_address: DocAddress) -> MoreLikeThisQuery { + MoreLikeThisQuery { + mlt: self.mlt, + doc_address: Some(doc_address), + doc_fields: None, + } + } + + /// Sets the document fields + /// Returns the constructed [`MoreLikeThisQuery`] + /// + /// This represents the list field values possibly collected from multiple documents + /// that will be used to compose the resulting query. + /// This interface is meant to be used when you want to provide your own set of fields + /// not necessarily from a specific document. + pub fn with_document_fields( + self, + doc_fields: Vec<(Field, Vec)>, + ) -> MoreLikeThisQuery { + MoreLikeThisQuery { + mlt: self.mlt, + doc_address: None, + doc_fields: Some(doc_fields), + } + } +} + +#[cfg(test)] +mod tests { + use super::MoreLikeThisQuery; + use crate::collector::TopDocs; + use crate::schema::{Schema, STORED, TEXT}; + use crate::DocAddress; + use crate::Index; + + fn create_test_index() -> Index { + let mut schema_builder = Schema::builder(); + let title = schema_builder.add_text_field("title", TEXT); + let body = schema_builder.add_text_field("body", TEXT | STORED); + let schema = schema_builder.build(); + let index = Index::create_in_ram(schema); + let mut index_writer = index.writer_for_tests().unwrap(); + index_writer.add_document(doc!(title => "aaa", body => "the old man and the sea")); + index_writer.add_document(doc!(title => "bbb", body => "an old man sailing on the sea")); + index_writer.add_document(doc!(title => "ccc", body=> "send this message to alice")); + index_writer.add_document(doc!(title => "ddd", body=> "a lady was riding and old bike")); + index_writer.add_document(doc!(title => "eee", body=> "Yes, my lady.")); + index_writer.commit().unwrap(); + index + } + + #[test] + fn test_more_like_this_query_builder() { + // default settings + let query = MoreLikeThisQuery::builder().with_document_fields(vec![]); + + assert_eq!(query.mlt.min_doc_frequency, Some(5)); + assert_eq!(query.mlt.max_doc_frequency, None); + assert_eq!(query.mlt.min_term_frequency, Some(2)); + assert_eq!(query.mlt.max_query_terms, Some(25)); + assert_eq!(query.mlt.min_word_length, None); + assert_eq!(query.mlt.max_word_length, None); + assert_eq!(query.mlt.boost_factor, Some(1.0)); + assert_eq!(query.mlt.stop_words, Vec::::new()); + assert_eq!(query.doc_fields, Some(vec![])); + assert_eq!(query.doc_address, None); + + // custom settings + let query = MoreLikeThisQuery::builder() + .with_min_doc_frequency(2) + .with_max_doc_frequency(5) + .with_min_term_frequency(2) + .with_min_word_length(2) + .with_max_word_length(4) + .with_boost_factor(0.5) + .with_stop_words(vec!["all".to_string(), "for".to_string()]) + .with_document(DocAddress::new(1, 2)); + + assert_eq!(query.mlt.min_doc_frequency, Some(2)); + assert_eq!(query.mlt.max_doc_frequency, Some(5)); + assert_eq!(query.mlt.min_term_frequency, Some(2)); + assert_eq!(query.mlt.min_word_length, Some(2)); + assert_eq!(query.mlt.max_word_length, Some(4)); + assert_eq!(query.mlt.boost_factor, Some(0.5)); + assert_eq!( + query.mlt.stop_words, + vec!["all".to_string(), "for".to_string()] + ); + assert_eq!(query.doc_fields, None); + assert_eq!(query.doc_address, Some(DocAddress::new(1, 2))); + } + + #[test] + fn test_more_like_this_query() { + let index = create_test_index(); + let reader = index.reader().unwrap(); + let searcher = reader.searcher(); + + // search base 1st doc with words [sea, and] skipping [old] + let query = MoreLikeThisQuery::builder() + .with_min_doc_frequency(1) + .with_max_doc_frequency(10) + .with_min_term_frequency(1) + .with_min_word_length(2) + .with_max_word_length(5) + .with_boost_factor(1.0) + .with_stop_words(vec!["old".to_string()]) + .with_document(DocAddress::new(0, 0)); + let top_docs = searcher.search(&query, &TopDocs::with_limit(5)).unwrap(); + let mut doc_ids: Vec<_> = top_docs.iter().map(|item| item.1.doc_id).collect(); + doc_ids.sort(); + + assert_eq!(doc_ids.len(), 3); + assert_eq!(doc_ids, vec![0, 1, 3]); + + // search base 5th doc with words [lady] + let query = MoreLikeThisQuery::builder() + .with_min_doc_frequency(1) + .with_max_doc_frequency(10) + .with_min_term_frequency(1) + .with_min_word_length(2) + .with_max_word_length(5) + .with_boost_factor(1.0) + .with_document(DocAddress::new(0, 4)); + let top_docs = searcher.search(&query, &TopDocs::with_limit(5)).unwrap(); + let mut doc_ids: Vec<_> = top_docs.iter().map(|item| item.1.doc_id).collect(); + doc_ids.sort(); + + assert_eq!(doc_ids.len(), 2); + assert_eq!(doc_ids, vec![3, 4]); + } +} diff --git a/src/query/mod.rs b/src/query/mod.rs index bc8e517bf..9d02540e8 100644 --- a/src/query/mod.rs +++ b/src/query/mod.rs @@ -11,6 +11,7 @@ mod exclude; mod explanation; mod fuzzy_query; mod intersection; +mod mlt; mod phrase_query; mod query; mod query_parser; @@ -45,6 +46,7 @@ pub use self::explanation::Explanation; pub(crate) use self::fuzzy_query::DfaWrapper; pub use self::fuzzy_query::FuzzyTermQuery; pub use self::intersection::intersect_scorers; +pub use self::mlt::MoreLikeThisQuery; pub use self::phrase_query::PhraseQuery; pub use self::query::{Query, QueryClone}; pub use self::query_parser::QueryParser; From 27f587aa1325a9d44406f54ccd7d0f5c3cb564ad Mon Sep 17 00:00:00 2001 From: Evance Souamoro Date: Thu, 29 Apr 2021 12:15:34 +0000 Subject: [PATCH 2/7] applied cargo fmt --- src/query/mlt/mlt.rs | 16 ++++++++-------- src/query/mlt/mod.rs | 2 +- src/query/mlt/query.rs | 10 +++++----- 3 files changed, 14 insertions(+), 14 deletions(-) diff --git a/src/query/mlt/mlt.rs b/src/query/mlt/mlt.rs index 9839f61b8..2c9ffbc5d 100644 --- a/src/query/mlt/mlt.rs +++ b/src/query/mlt/mlt.rs @@ -94,8 +94,8 @@ impl MoreLikeThis { } /// Creates a [`BooleanQuery`] from an ascendingly sorted list of ScoreTerm - /// This will map the list of ScoreTerm to a list of [`TermQuery`] and compose a - /// BooleanQuery using that list as sub queries. + /// This will map the list of ScoreTerm to a list of [`TermQuery`] and compose a + /// BooleanQuery using that list as sub queries. fn create_query(&self, score_terms: Vec) -> BooleanQuery { let best_score = score_terms.first().map_or(1f32, |x| x.score); let mut queries = Vec::new(); @@ -153,9 +153,9 @@ impl MoreLikeThis { self.create_score_term(searcher, field_to_term_freq_map) } - /// Computes the frequency of values for a field while updating the term frequencies - /// Note: A FieldValue can be made up of multiple terms. - /// We are interested in extracting terms within FieldValue + /// Computes the frequency of values for a field while updating the term frequencies + /// Note: A FieldValue can be made up of multiple terms. + /// We are interested in extracting terms within FieldValue fn add_term_frequencies( &self, searcher: &Searcher, @@ -171,7 +171,7 @@ impl MoreLikeThis { return Ok(()); } - // extract the raw value, possibly tokenizing & filtering to update the term frequency map + // extract the raw value, possibly tokenizing & filtering to update the term frequency map match field_entry.field_type() { FieldType::HierarchicalFacet(_) => { let facets: Vec<&str> = field_values @@ -325,7 +325,7 @@ impl MoreLikeThis { .sum::(); for (term, term_frequency) in per_field_term_frequencies.into_iter() { - // ignore terms with less than min_term_frequency + // ignore terms with less than min_term_frequency if self .min_term_frequency .map(|x| term_frequency < x) @@ -359,7 +359,7 @@ impl MoreLikeThis { continue; } - // compute similarity & score + // compute similarity & score let idf = self.idf(doc_freq, num_docs); let score = (term_frequency as f32) * idf; score_terms.push(ScoreTerm::new(term, score)); diff --git a/src/query/mlt/mod.rs b/src/query/mlt/mod.rs index 97f541908..f7e4da25e 100644 --- a/src/query/mlt/mod.rs +++ b/src/query/mlt/mod.rs @@ -2,4 +2,4 @@ mod mlt; mod query; pub use self::mlt::MoreLikeThis; -pub use self::query::{MoreLikeThisQuery}; +pub use self::query::MoreLikeThisQuery; diff --git a/src/query/mlt/query.rs b/src/query/mlt/query.rs index 3bc5cc833..3728495f6 100644 --- a/src/query/mlt/query.rs +++ b/src/query/mlt/query.rs @@ -6,7 +6,7 @@ use crate::{ DocAddress, Result, Searcher, TantivyError, }; -/// A query that matches all of the documents similar to a document +/// A query that matches all of the documents similar to a document /// or a set of field values provided. /// /// # Examples @@ -60,7 +60,7 @@ impl Query for MoreLikeThisQuery { } } -/// The builder for more-like-this query +/// The builder for more-like-this query #[derive(Debug, Clone)] pub struct MoreLikeThisQueryBuilder { mlt: MoreLikeThis, @@ -149,7 +149,7 @@ impl MoreLikeThisQueryBuilder { /// needed for composing the query. /// /// Note that field values will only be collected from stored fields in the index. - /// You can construct your own field values from any source. + /// You can construct your own field values from any source. pub fn with_document(self, doc_address: DocAddress) -> MoreLikeThisQuery { MoreLikeThisQuery { mlt: self.mlt, @@ -247,8 +247,8 @@ mod tests { let index = create_test_index(); let reader = index.reader().unwrap(); let searcher = reader.searcher(); - - // search base 1st doc with words [sea, and] skipping [old] + + // search base 1st doc with words [sea, and] skipping [old] let query = MoreLikeThisQuery::builder() .with_min_doc_frequency(1) .with_max_doc_frequency(10) From 2c0f6e33194809e1f7ddce869859446b634a1c08 Mon Sep 17 00:00:00 2001 From: Evance Souamoro Date: Thu, 29 Apr 2021 12:38:16 +0000 Subject: [PATCH 3/7] add builder to the public for documentation --- src/query/mlt/mod.rs | 2 +- src/query/mod.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/query/mlt/mod.rs b/src/query/mlt/mod.rs index f7e4da25e..6ab7d9a3c 100644 --- a/src/query/mlt/mod.rs +++ b/src/query/mlt/mod.rs @@ -2,4 +2,4 @@ mod mlt; mod query; pub use self::mlt::MoreLikeThis; -pub use self::query::MoreLikeThisQuery; +pub use self::query::{MoreLikeThisQuery, MoreLikeThisQueryBuilder}; diff --git a/src/query/mod.rs b/src/query/mod.rs index 9d02540e8..537aada50 100644 --- a/src/query/mod.rs +++ b/src/query/mod.rs @@ -46,7 +46,7 @@ pub use self::explanation::Explanation; pub(crate) use self::fuzzy_query::DfaWrapper; pub use self::fuzzy_query::FuzzyTermQuery; pub use self::intersection::intersect_scorers; -pub use self::mlt::MoreLikeThisQuery; +pub use self::mlt::{MoreLikeThisQuery, MoreLikeThisQueryBuilder}; pub use self::phrase_query::PhraseQuery; pub use self::query::{Query, QueryClone}; pub use self::query_parser::QueryParser; From cde324d4b461c0dce24e3b68387732b241f7aac6 Mon Sep 17 00:00:00 2001 From: Evance Souamoro Date: Fri, 30 Apr 2021 21:14:19 +0000 Subject: [PATCH 4/7] fixed issues based on comment, still need to check BM25 suggestion --- src/query/mlt/mlt.rs | 62 ++++++++++++++++++++++-------------------- src/query/mlt/query.rs | 44 +++++++++++++++--------------- 2 files changed, 55 insertions(+), 51 deletions(-) diff --git a/src/query/mlt/mlt.rs b/src/query/mlt/mlt.rs index 2c9ffbc5d..78cdc1f85 100644 --- a/src/query/mlt/mlt.rs +++ b/src/query/mlt/mlt.rs @@ -1,3 +1,4 @@ +use std::cmp::Reverse; use std::collections::{BinaryHeap, HashMap}; use crate::{ @@ -34,6 +35,11 @@ impl Ord for ScoreTerm { } /// A struct used as helper to build [`MoreLikeThisQuery`] +/// This more-like-this implementation is inspired by the Appache Lucene +/// amd closely follows the same implementation with adaptabtion to Tantivy vocabulary and API. +/// +/// [MoreLikeThis](https://github.com/apache/lucene/blob/main/lucene/queries/src/java/org/apache/lucene/queries/mlt/MoreLikeThis.java#L147) +/// [MoreLikeThisQuery](https://github.com/apache/lucene/blob/main/lucene/queries/src/java/org/apache/lucene/queries/mlt/MoreLikeThisQuery.java#L36) #[derive(Debug, Clone)] pub struct MoreLikeThis { /// Ignore words which do not occur in at least this many docs. @@ -194,16 +200,10 @@ impl MoreLikeThis { } FieldType::Str(text_options) => { let mut token_streams: Vec = vec![]; - let mut offsets = vec![]; - let mut total_offset = 0; for field_value in field_values { match field_value.value() { Value::PreTokStr(tok_str) => { - offsets.push(total_offset); - if let Some(last_token) = tok_str.tokens.last() { - total_offset += last_token.offset_to; - } token_streams.push(PreTokenizedStream::from(tok_str.clone()).into()); } Value::Str(ref text) => { @@ -214,9 +214,6 @@ impl MoreLikeThis { }) .and_then(|tokenizer_name| tokenizer_manager.get(&tokenizer_name)) { - offsets.push(total_offset); - total_offset += text.len(); - //let v = text.clone(); token_streams.push(tokenizer.token_stream(text)); } } @@ -317,18 +314,18 @@ impl MoreLikeThis { searcher: &Searcher, per_field_term_frequencies: HashMap, ) -> Result> { - let mut score_terms = BinaryHeap::new(); + let mut score_terms: BinaryHeap> = BinaryHeap::new(); let num_docs = searcher .segment_readers() .iter() - .map(|x| x.num_docs() as u64) + .map(|segment_reader| segment_reader.num_docs() as u64) .sum::(); - for (term, term_frequency) in per_field_term_frequencies.into_iter() { + for (term, term_frequency) in per_field_term_frequencies.iter() { // ignore terms with less than min_term_frequency if self .min_term_frequency - .map(|x| term_frequency < x) + .map(|min_term_frequency| *term_frequency < min_term_frequency) .unwrap_or(false) { continue; @@ -339,7 +336,7 @@ impl MoreLikeThis { // ignore terms with less than min_doc_frequency if self .min_doc_frequency - .map(|x| doc_freq < x) + .map(|min_doc_frequency| doc_freq < min_doc_frequency) .unwrap_or(false) { continue; @@ -348,7 +345,7 @@ impl MoreLikeThis { // ignore terms with more than max_doc_frequency if self .max_doc_frequency - .map(|x| doc_freq > x) + .map(|max_doc_frequency| doc_freq > max_doc_frequency) .unwrap_or(false) { continue; @@ -361,22 +358,29 @@ impl MoreLikeThis { // compute similarity & score let idf = self.idf(doc_freq, num_docs); - let score = (term_frequency as f32) * idf; - score_terms.push(ScoreTerm::new(term, score)); + let score = (*term_frequency as f32) * idf; + if let Some(limit) = self.max_query_terms { + if score_terms.len() > limit { + // update the least significant term + let least_significant_term_score = score_terms.peek().unwrap().0.score; + if least_significant_term_score < score { + score_terms.peek_mut().unwrap().0 = ScoreTerm::new(term.clone(), score); + } + } else { + score_terms.push(Reverse(ScoreTerm::new(term.clone(), score))); + } + } else { + score_terms.push(Reverse(ScoreTerm::new(term.clone(), score))); + } } - // limit ourself to max_query terms. we need to sort so to avoid discarding important terms - let score_terms = if let Some(max_query_terms) = self.max_query_terms { - let max_num_terms = std::cmp::min(max_query_terms, score_terms.len()); - score_terms - .into_sorted_vec() - .into_iter() - .take(max_num_terms) - .collect() - } else { - score_terms.into_sorted_vec() - }; - Ok(score_terms) + let mut score_terms_vec: Vec = score_terms + .into_iter() + .map(|reverse_score_term| reverse_score_term.0) + .collect(); + score_terms_vec.sort_unstable(); + + Ok(score_terms_vec) } /// Computes the similarity diff --git a/src/query/mlt/query.rs b/src/query/mlt/query.rs index 3728495f6..a24d04a59 100644 --- a/src/query/mlt/query.rs +++ b/src/query/mlt/query.rs @@ -3,7 +3,7 @@ use super::MoreLikeThis; use crate::{ query::{Query, Weight}, schema::{Field, FieldValue}, - DocAddress, Result, Searcher, TantivyError, + DocAddress, Result, Searcher, }; /// A query that matches all of the documents similar to a document @@ -29,8 +29,13 @@ use crate::{ #[derive(Debug, Clone)] pub struct MoreLikeThisQuery { mlt: MoreLikeThis, - doc_address: Option, - doc_fields: Option)>>, + target: TargetDocument, +} + +#[derive(Debug, PartialEq, Clone)] +enum TargetDocument { + DocumentAdress(DocAddress), + DocumentFields(Vec<(Field, Vec)>), } impl MoreLikeThisQuery { @@ -42,21 +47,16 @@ impl MoreLikeThisQuery { impl Query for MoreLikeThisQuery { fn weight(&self, searcher: &Searcher, scoring_enabled: bool) -> Result> { - if let Some(doc_address) = self.doc_address { - return self + match &self.target { + TargetDocument::DocumentAdress(doc_address) => self .mlt - .query_with_document(searcher, doc_address)? - .weight(searcher, scoring_enabled); - } - - if let Some(ref doc_fields) = self.doc_fields { - return self + .query_with_document(searcher, *doc_address)? + .weight(searcher, scoring_enabled), + TargetDocument::DocumentFields(doc_fields) => self .mlt .query_with_document_fields(searcher, doc_fields)? - .weight(searcher, scoring_enabled); + .weight(searcher, scoring_enabled), } - - Err(TantivyError::InvalidArgument("".to_string())) } } @@ -153,8 +153,7 @@ impl MoreLikeThisQueryBuilder { pub fn with_document(self, doc_address: DocAddress) -> MoreLikeThisQuery { MoreLikeThisQuery { mlt: self.mlt, - doc_address: Some(doc_address), - doc_fields: None, + target: TargetDocument::DocumentAdress(doc_address), } } @@ -171,8 +170,7 @@ impl MoreLikeThisQueryBuilder { ) -> MoreLikeThisQuery { MoreLikeThisQuery { mlt: self.mlt, - doc_address: None, - doc_fields: Some(doc_fields), + target: TargetDocument::DocumentFields(doc_fields), } } } @@ -180,6 +178,7 @@ impl MoreLikeThisQueryBuilder { #[cfg(test)] mod tests { use super::MoreLikeThisQuery; + use super::TargetDocument; use crate::collector::TopDocs; use crate::schema::{Schema, STORED, TEXT}; use crate::DocAddress; @@ -214,8 +213,7 @@ mod tests { assert_eq!(query.mlt.max_word_length, None); assert_eq!(query.mlt.boost_factor, Some(1.0)); assert_eq!(query.mlt.stop_words, Vec::::new()); - assert_eq!(query.doc_fields, Some(vec![])); - assert_eq!(query.doc_address, None); + assert_eq!(query.target, TargetDocument::DocumentFields(vec![])); // custom settings let query = MoreLikeThisQuery::builder() @@ -238,8 +236,10 @@ mod tests { query.mlt.stop_words, vec!["all".to_string(), "for".to_string()] ); - assert_eq!(query.doc_fields, None); - assert_eq!(query.doc_address, Some(DocAddress::new(1, 2))); + assert_eq!( + query.target, + TargetDocument::DocumentAdress(DocAddress::new(1, 2)) + ); } #[test] From 712c01aa93adf92df417182b5a4df291d68afc35 Mon Sep 17 00:00:00 2001 From: Evance Souamoro Date: Sat, 1 May 2021 05:40:59 +0000 Subject: [PATCH 5/7] fixed term sorting & moved it to a better place --- src/query/mlt/mlt.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/query/mlt/mlt.rs b/src/query/mlt/mlt.rs index 78cdc1f85..19f50d11d 100644 --- a/src/query/mlt/mlt.rs +++ b/src/query/mlt/mlt.rs @@ -102,7 +102,8 @@ impl MoreLikeThis { /// Creates a [`BooleanQuery`] from an ascendingly sorted list of ScoreTerm /// This will map the list of ScoreTerm to a list of [`TermQuery`] and compose a /// BooleanQuery using that list as sub queries. - fn create_query(&self, score_terms: Vec) -> BooleanQuery { + fn create_query(&self, mut score_terms: Vec) -> BooleanQuery { + score_terms.sort_by(|left_ts, right_ts| right_ts.cmp(left_ts)); let best_score = score_terms.first().map_or(1f32, |x| x.score); let mut queries = Vec::new(); @@ -374,11 +375,10 @@ impl MoreLikeThis { } } - let mut score_terms_vec: Vec = score_terms + let score_terms_vec: Vec = score_terms .into_iter() .map(|reverse_score_term| reverse_score_term.0) .collect(); - score_terms_vec.sort_unstable(); Ok(score_terms_vec) } From d71aa57077479aefba73b9d4ccf39b9fac2f48df Mon Sep 17 00:00:00 2001 From: Evance Souamoro Date: Mon, 3 May 2021 10:05:40 +0000 Subject: [PATCH 6/7] reusing idf from bm25 module as it was the same logic --- CHANGELOG.md | 4 +++- src/query/bm25.rs | 2 +- src/query/mlt/mlt.rs | 9 ++------- 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 164da9cb8..7f7b1ee7e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,8 +8,9 @@ Tantivy 0.15.0 - Bugfix consistent tie break handling in facet's topk (@hardikpnsp) #357 - Date field support for range queries (@rihardsk) #516 - Added lz4-flex as the default compression scheme in tantivy (@PSeitz) #1009 -- Renamed a lot of symbols to avoid all uppercasing on acronyms, as per new clippy recommendation. For instance, RAMDireotory -> RamDirectory. (@pmasurel) +- Renamed a lot of symbols to avoid all uppercasing on acronyms, as per new clippy recommendation. For instance, RAMDirectory -> RamDirectory. (@pmasurel) - Simplified positions index format (@fulmicoton) #1022 +- Added support for more-like-this query in tantivy (@evanxg852000) #1011 Tantivy 0.14.0 ========================= @@ -25,6 +26,7 @@ Tantivy 0.14.0 - Simplified the encoding of the skip reader struct. BlockWAND max tf is now encoded over a single byte. (@fulmicoton) - `FilterCollector` now supports all Fast Field value types (@barrotsteindev) - FastField are not all loaded when opening the segment reader. (@fulmicoton) +- Added an API to merge segments, see `tantivy::merge_segments` #1005. (@evanxg852000) This version breaks compatibility and requires users to reindex everything. diff --git a/src/query/bm25.rs b/src/query/bm25.rs index f017e3d4b..c056ba971 100644 --- a/src/query/bm25.rs +++ b/src/query/bm25.rs @@ -9,7 +9,7 @@ use serde::Serialize; const K1: Score = 1.2; const B: Score = 0.75; -fn idf(doc_freq: u64, doc_count: u64) -> Score { +pub(crate) fn idf(doc_freq: u64, doc_count: u64) -> Score { assert!(doc_count >= doc_freq, "{} >= {}", doc_count, doc_freq); let x = ((doc_count - doc_freq) as Score + 0.5) / (doc_freq as Score + 0.5); (1.0 + x).ln() diff --git a/src/query/mlt/mlt.rs b/src/query/mlt/mlt.rs index 19f50d11d..176fcc3fd 100644 --- a/src/query/mlt/mlt.rs +++ b/src/query/mlt/mlt.rs @@ -2,7 +2,7 @@ use std::cmp::Reverse; use std::collections::{BinaryHeap, HashMap}; use crate::{ - query::{BooleanQuery, BoostQuery, Occur, Query, TermQuery}, + query::{BooleanQuery, BoostQuery, Occur, Query, TermQuery, bm25::idf}, schema::{Field, FieldType, FieldValue, IndexRecordOption, Term, Value}, tokenizer::{BoxTokenStream, FacetTokenizer, PreTokenizedStream, Tokenizer}, DocAddress, Result, Searcher, TantivyError, @@ -358,7 +358,7 @@ impl MoreLikeThis { } // compute similarity & score - let idf = self.idf(doc_freq, num_docs); + let idf = idf(doc_freq, num_docs); let score = (*term_frequency as f32) * idf; if let Some(limit) = self.max_query_terms { if score_terms.len() > limit { @@ -383,9 +383,4 @@ impl MoreLikeThis { Ok(score_terms_vec) } - /// Computes the similarity - fn idf(&self, doc_freq: u64, doc_count: u64) -> f32 { - let x = ((doc_count - doc_freq) as f32 + 0.5) / (doc_freq as f32 + 0.5); - (1f32 + x).ln() - } } From 372d12766a4eff7a1ba36251e0c987dca42c6464 Mon Sep 17 00:00:00 2001 From: Evance Souamoro Date: Mon, 3 May 2021 10:26:56 +0000 Subject: [PATCH 7/7] fix cargo fmt --- src/query/mlt/mlt.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/query/mlt/mlt.rs b/src/query/mlt/mlt.rs index 176fcc3fd..477f4da0b 100644 --- a/src/query/mlt/mlt.rs +++ b/src/query/mlt/mlt.rs @@ -2,7 +2,7 @@ use std::cmp::Reverse; use std::collections::{BinaryHeap, HashMap}; use crate::{ - query::{BooleanQuery, BoostQuery, Occur, Query, TermQuery, bm25::idf}, + query::{bm25::idf, BooleanQuery, BoostQuery, Occur, Query, TermQuery}, schema::{Field, FieldType, FieldValue, IndexRecordOption, Term, Value}, tokenizer::{BoxTokenStream, FacetTokenizer, PreTokenizedStream, Tokenizer}, DocAddress, Result, Searcher, TantivyError, @@ -382,5 +382,4 @@ impl MoreLikeThis { Ok(score_terms_vec) } - }