perf: Implement a TermSet variant which uses fast fields (#69)

The `TermSet` `Query` currently produces one `Scorer`/`DocSet` per matched term by scanning the term dictionary and then consuming posting lists. For very large sets of terms and a fast field, it is faster to scan the fast field column while intersecting with a `HashSet` of (encoded) term values.

Following the pattern set by the two execution modes of `RangeQuery`, this PR introduces a variant of `TermSet` which uses fast fields, and then uses it when there are more than 1024 input terms (an arbitrary threshold!).

Performance is significantly improved for large `TermSet`s of primitives.

Author: stuhood

src/query/mod.rs

@@ -22,8 +22,8 @@ mod range_query;
 mod regex_query;
 mod reqopt_scorer;
 mod scorer;
-mod set_query;
 mod term_query;
+mod term_set_query;
 mod union;
 mod weight;
 
@@ -61,8 +61,8 @@ pub use self::regex_query::RegexQuery;
 pub use self::reqopt_scorer::RequiredOptionalScorer;
 pub use self::score_combiner::{DisjunctionMaxCombiner, ScoreCombiner, SumCombiner};
 pub use self::scorer::Scorer;
-pub use self::set_query::TermSetQuery;
 pub use self::term_query::TermQuery;
+pub use self::term_set_query::*;
 pub use self::union::{BufferedUnionScorer, SimpleUnion};
 #[cfg(test)]
 pub use self::vec_docset::VecDocSet;

src/query/term_set_query/mod.rs

@@ -0,0 +1,5 @@
+mod term_set_query;
+mod term_set_query_fastfield;
+
+pub use self::term_set_query::{InvertedIndexTermSetQuery, TermSetQuery};
+pub use self::term_set_query_fastfield::FastFieldTermSetQuery;

src/query/set_query.rs …c/query/term_set_query/term_set_query.rssrc/query/set_query.rs renamed to src/query/term_set_query/term_set_query.rs src/query/set_query.rs renamed to src/query/term_set_query/term_set_query.rs

@@ -3,9 +3,10 @@ use std::collections::HashMap;
 use tantivy_fst::raw::CompiledAddr;
 use tantivy_fst::{Automaton, Map};
 
+use super::term_set_query_fastfield::FastFieldTermSetWeight;
 use crate::query::score_combiner::DoNothingCombiner;
 use crate::query::{AutomatonWeight, BooleanWeight, EnableScoring, Occur, Query, Weight};
-use crate::schema::{Field, Schema};
+use crate::schema::{Field, Schema, Type};
 use crate::{SegmentReader, Term};
 
 /// A Term Set Query matches all of the documents containing any of the Term provided
@@ -44,20 +45,41 @@ impl TermSetQuery {
                 return Err(crate::TantivyError::SchemaError(error_msg));
             }
 
-            // In practice this won't fail because:
-            // - we are writing to memory, so no IoError
-            // - Terms are ordered
-            let map = Map::from_iter(
-                sorted_terms
-                    .iter()
-                    .map(|key| (key.serialized_value_bytes(), 0)),
-            )
-            .map_err(std::io::Error::other)?;
-
-            sub_queries.push((
-                Occur::Should,
-                Box::new(AutomatonWeight::new(field, SetDfaWrapper(map))),
-            ));
+            let supported_for_ff = sorted_terms
+                .get(0)
+                .map(|term| match term.typ() {
+                    Type::U64 | Type::I64 | Type::F64 | Type::Bool | Type::Date | Type::IpAddr => {
+                        true
+                    }
+                    Type::Json | Type::Str => {
+                        // Explicitly not supported yet: see `term_set_query_fastfield.rs`.
+                        false
+                    }
+                    _ => false,
+                })
+                .unwrap_or(false);
+
+            if field_type.is_fast() && supported_for_ff {
+                sub_queries.push((
+                    Occur::Should,
+                    Box::new(FastFieldTermSetWeight::new(field, sorted_terms.to_vec())),
+                ));
+            } else {
+                // In practice this won't fail because:
+                // - we are writing to memory, so no IoError
+                // - Terms are ordered
+                let map = Map::from_iter(
+                    sorted_terms
+                        .iter()
+                        .map(|key| (key.serialized_value_bytes(), 0)),
+                )
+                .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
+
+                sub_queries.push((
+                    Occur::Should,
+                    Box::new(AutomatonWeight::new(field, SetDfaWrapper(map))),
+                ));
+            }
         }
 
         Ok(BooleanWeight::new(
@@ -87,6 +109,59 @@ impl Query for TermSetQuery {
     }
 }
 
+/// `InvertedIndexTermSetQuery` is the same as [TermSetQuery] but only uses the inverted index.
+#[derive(Debug, Clone)]
+pub struct InvertedIndexTermSetQuery {
+    terms_map: HashMap<Field, Vec<Term>>,
+}
+
+impl InvertedIndexTermSetQuery {
+    /// Create a new `InvertedIndexTermSetQuery`.
+    pub fn new<T: IntoIterator<Item = Term>>(terms: T) -> Self {
+        let mut terms_map: HashMap<_, Vec<_>> = HashMap::new();
+        for term in terms {
+            terms_map.entry(term.field()).or_default().push(term);
+        }
+
+        for terms in terms_map.values_mut() {
+            terms.sort_unstable();
+            terms.dedup();
+        }
+
+        InvertedIndexTermSetQuery { terms_map }
+    }
+}
+
+impl Query for InvertedIndexTermSetQuery {
+    fn weight(&self, enable_scoring: EnableScoring<'_>) -> crate::Result<Box<dyn Weight>> {
+        let mut sub_queries: Vec<(_, Box<dyn Weight>)> = Vec::with_capacity(self.terms_map.len());
+        for (&field, sorted_terms) in &self.terms_map {
+            let schema = enable_scoring.schema();
+            let field_entry = schema.get_field_entry(field);
+            if !field_entry.field_type().is_indexed() {
+                let error_msg = format!("Field {:?} is not indexed.", field_entry.name());
+                return Err(crate::TantivyError::SchemaError(error_msg));
+            }
+            let map = Map::from_iter(
+                sorted_terms
+                    .iter()
+                    .map(|key| (key.serialized_value_bytes(), 0)),
+            )
+            .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
+
+            sub_queries.push((
+                Occur::Should,
+                Box::new(AutomatonWeight::new(field, SetDfaWrapper(map))),
+            ));
+        }
+        Ok(Box::new(BooleanWeight::new(
+            sub_queries,
+            false,
+            Box::new(DoNothingCombiner::default),
+        )))
+    }
+}
+
 struct SetDfaWrapper(Map<Vec<u8>>);
 
 impl Automaton for SetDfaWrapper {