Skip to main content

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}