Evaluate EQ(REGEXP_EXTRACT) once per distinct value with RegexExtractEqQuery

quickwit/Cargo.toml

@@ -386,6 +386,8 @@ tantivy = { git = "https://github.com/quickwit-oss/tantivy/", rev = "e229de6db74
   "zstd-compression",
   "columnar-zstd-compression",
 ] }
+# Must stay on the same rev as `tantivy`.
+tantivy-common = { git = "https://github.com/quickwit-oss/tantivy/", rev = "e229de6db748c6b21747a99a18220446e76104e1" }
 tantivy-fst = "0.5"
 
 # This is actually not used directly the goal is to fix the version

quickwit/quickwit-doc-mapper/src/query_builder.rs

@@ -380,6 +380,18 @@ impl<'a, 'b: 'a> QueryAstVisitor<'a> for ExtractPrefixTermRanges<'b> {
         );
         Ok(())
     }
+
+    fn visit_calc_field(&mut self, calc_field_query: &'a CalcFieldQuery) -> Result<(), Self::Err> {
+        let Some(regex_query) = calc_field_query.try_prefilter_regex_query(self.schema) else {
+            return Ok(());
+        };
+        let resolved = regex_query.to_resolved(self.schema, Some(self.tokenizer_manager))?;
+        self.add_automaton(
+            resolved.field,
+            Automaton::Regex(resolved.json_path, resolved.regex),
+        );
+        Ok(())
+    }
 }
 
 type TermRangeWarmupInfo = HashMap<Field, HashMap<TermRange, PositionNeeded>>;
@@ -400,7 +412,7 @@ fn extract_prefix_term_ranges_and_automaton(
 
 #[cfg(test)]
 mod test {
-    use std::collections::HashSet;
+    use std::collections::{HashMap, HashSet};
     use std::ops::Bound;
 
     use quickwit_common::shared_consts::FIELD_PRESENCE_FIELD_NAME;
@@ -413,11 +425,14 @@ mod test {
     };
     use tantivy::Term;
     use tantivy::jitexpr::ast::deserialize;
-    use tantivy::schema::{DateOptions, DateTimePrecision, FAST, INDEXED, STORED, Schema, TEXT};
+    use tantivy::schema::{
+        DateOptions, DateTimePrecision, FAST, Field, INDEXED, STORED, STRING, Schema, TEXT,
+    };
 
     use super::{ExtractPrefixTermRanges, build_query};
     use crate::{
-        DYNAMIC_FIELD_NAME, FastFieldWarmupInfo, SOURCE_FIELD_NAME, TermRange, WarmupInfo,
+        Automaton, DYNAMIC_FIELD_NAME, FastFieldWarmupInfo, SOURCE_FIELD_NAME, TermRange,
+        WarmupInfo,
     };
 
     fn calc_field(expression: &str) -> QueryAst {
@@ -488,6 +503,64 @@ mod test {
         );
     }
 
+    #[test]
+    fn test_calc_field_regex_extract_eq_warmup() {
+        let mut schema_builder = Schema::builder();
+        let service = schema_builder.add_text_field("service", STRING | FAST);
+        schema_builder.add_text_field("tokenized", TEXT | FAST);
+        schema_builder.add_text_field("fast_only", FAST);
+        let schema = schema_builder.build();
+        let context = BuildTantivyAstContext::for_test(&schema);
+
+        // The predicate always reads the fast field. It may also read postings of terms accepted
+        // by the prefilter when the field uses the raw tokenizer for both indexing and fast fields.
+        for (expression, fast_field, expected_prefilter) in [
+            (
+                r#"(EQ (REGEXP_EXTRACT service "^svc-([a-z]+)-prod$" 1u64) "api")"#,
+                "service",
+                Some("svc-api-prod"),
+            ),
+            (
+                r#"(EQ (REGEXP_EXTRACT service "svc-([a-z]+)-prod" 1u64) "api")"#,
+                "service",
+                Some("(?s:.*)svc-api-prod(?s:.*)"),
+            ),
+            (
+                r#"(EQ (REGEXP_EXTRACT tokenized "^svc-([a-z]+)-prod$" 1u64) "api")"#,
+                "tokenized",
+                None,
+            ),
+            (
+                r#"(EQ (REGEXP_EXTRACT fast_only "^svc-([a-z]+)-prod$" 1u64) "api")"#,
+                "fast_only",
+                None,
+            ),
+            (
+                r#"(EQ (REGEXP_EXTRACT service "(?m)^svc-([a-z]+)" 1u64) "api")"#,
+                "service",
+                None,
+            ),
+        ] {
+            let (_, warmup) = build_query(calc_field(expression), &context, None).unwrap();
+            assert_eq!(
+                warmup.fast_fields,
+                expected_fast_fields(&[fast_field]),
+                "{expression}"
+            );
+            let expected_automatons: HashMap<Field, HashSet<Automaton>> = expected_prefilter
+                .map(|regex| {
+                    let automaton = Automaton::Regex(None, regex.to_string());
+                    (service, HashSet::from([automaton]))
+                })
+                .into_iter()
+                .collect();
+            assert_eq!(
+                warmup.automatons_grouped_by_field, expected_automatons,
+                "{expression}"
+            );
+        }
+    }
+
     fn full_text_query_for_warmup(field: &str, tokenizer: Option<&str>) -> QueryAst {
         quickwit_query::query_ast::FullTextQuery {
             field: field.to_string(),

quickwit/quickwit-query/Cargo.toml

@@ -15,11 +15,14 @@ anyhow = { workspace = true }
 base64 = { workspace = true }
 bitpacking = { workspace = true }
 hex = { workspace = true }
+lru = { workspace = true }
 regex = { workspace = true }
+regex-syntax = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 serde_with = { workspace = true }
 tantivy = { workspace = true }
+tantivy-common = { workspace = true }
 unicode-segmentation = { workspace = true }
 tantivy-fst = { workspace = true }
 tracing = { workspace = true }

quickwit/quickwit-query/benches/tokenizers_bench.rs

@@ -12,7 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use criterion::{Criterion, Throughput, black_box, criterion_group, criterion_main};
+use std::hint::black_box;
+
+use criterion::{Criterion, Throughput, criterion_group, criterion_main};
 use quickwit_query::CodeTokenizer;
 use tantivy::tokenizer::{RegexTokenizer, TextAnalyzer, Token, TokenStream};