lb_tantivy/query/disjunction_max_query.rs
1use crate::query::{BooleanWeight, DisjunctionMaxCombiner, EnableScoring, Occur, Query, Weight};
2use crate::{Score, Term};
3
4/// The disjunction max query returns documents matching one or more wrapped queries,
5/// called query clauses or clauses.
6///
7/// If a returned document matches multiple query clauses,
8/// the `DisjunctionMaxQuery` assigns the document the highest relevance score from any matching
9/// clause, plus a tie breaking increment for any additional matching subqueries.
10///
11/// ```rust
12/// use tantivy::collector::TopDocs;
13/// use tantivy::doc;
14/// use tantivy::query::{DisjunctionMaxQuery, Query, QueryClone, TermQuery};
15/// use tantivy::schema::{IndexRecordOption, Schema, TEXT};
16/// use tantivy::Term;
17/// use tantivy::Index;
18/// use tantivy::IndexWriter;
19///
20/// fn main() -> tantivy::Result<()> {
21/// let mut schema_builder = Schema::builder();
22/// let title = schema_builder.add_text_field("title", TEXT);
23/// let body = schema_builder.add_text_field("body", TEXT);
24/// let schema = schema_builder.build();
25/// let index = Index::create_in_ram(schema);
26/// {
27/// let mut index_writer: IndexWriter = index.writer(15_000_000)?;
28/// index_writer.add_document(doc!(
29/// title => "The Name of Girl",
30/// ))?;
31/// index_writer.add_document(doc!(
32/// title => "The Diary of Muadib",
33/// ))?;
34/// index_writer.add_document(doc!(
35/// title => "The Diary of Girl",
36/// ))?;
37/// index_writer.commit()?;
38/// }
39///
40/// let reader = index.reader()?;
41/// let searcher = reader.searcher();
42///
43/// // Make TermQuery's for "girl" and "diary" in the title
44/// let girl_term_query: Box<dyn Query> = Box::new(TermQuery::new(
45/// Term::from_field_text(title, "girl"),
46/// IndexRecordOption::Basic,
47/// ));
48/// let diary_term_query: Box<dyn Query> = Box::new(TermQuery::new(
49/// Term::from_field_text(title, "diary"),
50/// IndexRecordOption::Basic,
51/// ));
52///
53/// // TermQuery "diary" and "girl" should be present and only one should be accounted in score
54/// let queries1 = vec![diary_term_query.box_clone(), girl_term_query.box_clone()];
55/// let diary_and_girl = DisjunctionMaxQuery::new(queries1);
56/// let documents = searcher.search(&diary_and_girl, &TopDocs::with_limit(3))?;
57/// assert_eq!(documents[0].0, documents[1].0);
58/// assert_eq!(documents[1].0, documents[2].0);
59///
60/// // TermQuery "diary" and "girl" should be present
61/// // and one should be accounted with multiplier 0.7
62/// let queries2 = vec![diary_term_query.box_clone(), girl_term_query.box_clone()];
63/// let tie_breaker = 0.7;
64/// let diary_and_girl_with_tie_breaker = DisjunctionMaxQuery::with_tie_breaker(queries2, tie_breaker);
65/// let documents = searcher.search(&diary_and_girl_with_tie_breaker, &TopDocs::with_limit(3))?;
66/// assert_eq!(documents[1].0, documents[2].0);
67/// // For this test all terms brings the same score. So we can do easy math and assume that
68/// // `DisjunctionMaxQuery` with tie breakers score should be equal
69/// // to term1 score + `tie_breaker` * term2 score or (1.0 + tie_breaker) * term score
70/// assert!(f32::abs(documents[0].0 - documents[1].0 * (1.0 + tie_breaker)) < 0.001);
71/// Ok(())
72/// }
73/// ```
74#[derive(Debug)]
75pub struct DisjunctionMaxQuery {
76 disjuncts: Vec<Box<dyn Query>>,
77 tie_breaker: Score,
78}
79
80impl Clone for DisjunctionMaxQuery {
81 fn clone(&self) -> Self {
82 DisjunctionMaxQuery::with_tie_breaker(
83 self.disjuncts
84 .iter()
85 .map(|disjunct| disjunct.box_clone())
86 .collect::<Vec<_>>(),
87 self.tie_breaker,
88 )
89 }
90}
91
92impl Query for DisjunctionMaxQuery {
93 fn weight(&self, enable_scoring: EnableScoring<'_>) -> crate::Result<Box<dyn Weight>> {
94 let disjuncts = self
95 .disjuncts
96 .iter()
97 .map(|disjunct| Ok((Occur::Should, disjunct.weight(enable_scoring)?)))
98 .collect::<crate::Result<_>>()?;
99 let tie_breaker = self.tie_breaker;
100 Ok(Box::new(BooleanWeight::new(
101 disjuncts,
102 enable_scoring.is_scoring_enabled(),
103 Box::new(move || DisjunctionMaxCombiner::with_tie_breaker(tie_breaker)),
104 )))
105 }
106
107 fn query_terms<'a>(&'a self, visitor: &mut dyn FnMut(&'a Term, bool)) {
108 for disjunct in &self.disjuncts {
109 disjunct.query_terms(visitor);
110 }
111 }
112}
113
114impl DisjunctionMaxQuery {
115 /// Creates a new `DisjunctionMaxQuery` with tie breaker.
116 pub fn with_tie_breaker(
117 disjuncts: Vec<Box<dyn Query>>,
118 tie_breaker: Score,
119 ) -> DisjunctionMaxQuery {
120 DisjunctionMaxQuery {
121 disjuncts,
122 tie_breaker,
123 }
124 }
125
126 /// Creates a new `DisjunctionMaxQuery` with no tie breaker.
127 pub fn new(disjuncts: Vec<Box<dyn Query>>) -> DisjunctionMaxQuery {
128 DisjunctionMaxQuery::with_tie_breaker(disjuncts, 0.0)
129 }
130}