From d47f9c90912bc8fe5f70d6742f14b69320e0fc87 Mon Sep 17 00:00:00 2001 From: Paul Masurel Date: Thu, 15 Oct 2026 12:10:32 +0200 Subject: [PATCH] (Calculated fields) Allow aggregations on calculated fields Add `calculated_fields` to `SearchRequest`: named jitexpr expressions computed from fast fields. Aggregations can use their names like fast field names. - Bump tantivy to get the jitexpr `ValueSource` for aggregations. - At the leaf, the calculated fields are registered in the aggregation context as `JitExprValueSourceProvider`s. - Warmup replaces calculated fields by the fast fields their expressions read. - Root validation rejects invalid expressions and checks that the expression inputs are fast fields. --- quickwit/Cargo.lock | 46 +++-- quickwit/Cargo.toml | 4 +- .../protos/quickwit/search.proto | 13 ++ .../src/codegen/quickwit/quickwit.search.rs | 16 ++ .../quickwit-search/src/calculated_fields.rs | 194 ++++++++++++++++++ quickwit/quickwit-search/src/collector.rs | 104 +++++++++- quickwit/quickwit-search/src/lib.rs | 1 + quickwit/quickwit-search/src/root.rs | 9 +- quickwit/quickwit-search/src/tests.rs | 134 +++++++++++- 9 files changed, 491 insertions(+), 30 deletions(-) create mode 100644 quickwit/quickwit-search/src/calculated_fields.rs diff --git a/quickwit/Cargo.lock b/quickwit/Cargo.lock index 601594bdcda..7dcb18c6394 100644 --- a/quickwit/Cargo.lock +++ b/quickwit/Cargo.lock @@ -1377,7 +1377,7 @@ version = "3.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6dee98b0db6a962de883bf5d20362dee4d7ca0d12fe39a7c6c73c844e1cd7c1f" dependencies = [ - "darling 0.23.0", + "darling 0.20.11", "ident_case", "prettyplease", "proc-macro2", @@ -2481,6 +2481,7 @@ dependencies = [ "ident_case", "proc-macro2", "quote", + "strsim", "syn 2.0.118", ] @@ -3032,7 +3033,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] @@ -4434,7 +4435,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] @@ -4532,7 +4533,7 @@ dependencies = [ [[package]] name = "jitexpr" version = "0.1.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=8a4cfc40f5b7446c0700471aa50290dbb7f8ea97#8a4cfc40f5b7446c0700471aa50290dbb7f8ea97" dependencies = [ "cranelift", "cranelift-jit", @@ -5995,7 +5996,7 @@ checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" [[package]] name = "ownedbytes" version = "0.9.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=8a4cfc40f5b7446c0700471aa50290dbb7f8ea97#8a4cfc40f5b7446c0700471aa50290dbb7f8ea97" dependencies = [ "stable_deref_trait", ] @@ -6844,7 +6845,7 @@ version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "itertools 0.10.5", "log", "multimap", @@ -6864,7 +6865,7 @@ version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03da047801ff44bb6a4d407d4860c05fd70bb81714e6b2f3812603d5b145b042" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "itertools 0.10.5", "log", "multimap", @@ -8222,7 +8223,7 @@ dependencies = [ "once_cell", "socket2 0.5.10", "tracing", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] @@ -9007,7 +9008,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.12.1", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] @@ -9087,7 +9088,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] @@ -9328,7 +9329,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5b55fb86dfd3a2f5f76ea78310a88f96c4ea21a3031f8d212443d56123fd0521" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] @@ -9829,7 +9830,7 @@ version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c1c97747dbf44bb1ca44a561ece23508e99cb592e862f22222dcf42f51d1e451" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "proc-macro2", "quote", "syn 2.0.118", @@ -10306,7 +10307,7 @@ checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" [[package]] name = "tantivy" version = "0.27.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=8a4cfc40f5b7446c0700471aa50290dbb7f8ea97#8a4cfc40f5b7446c0700471aa50290dbb7f8ea97" dependencies = [ "aho-corasick", "arc-swap", @@ -10324,6 +10325,7 @@ dependencies = [ "fs4", "futures-channel", "futures-util", + "hashbrown 0.17.1", "htmlescape", "itertools 0.14.0", "jitexpr", @@ -10347,7 +10349,7 @@ dependencies = [ "tantivy-columnar", "tantivy-common", "tantivy-fst", - "tantivy-query-grammar 0.26.0 (git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe)", + "tantivy-query-grammar 0.26.0 (git+https://github.com/quickwit-oss/tantivy/?rev=8a4cfc40f5b7446c0700471aa50290dbb7f8ea97)", "tantivy-sstable", "tantivy-stacker", "tantivy-tokenizer-api", @@ -10364,7 +10366,7 @@ dependencies = [ [[package]] name = "tantivy-bitpacker" version = "0.10.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=8a4cfc40f5b7446c0700471aa50290dbb7f8ea97#8a4cfc40f5b7446c0700471aa50290dbb7f8ea97" dependencies = [ "bitpacking", ] @@ -10372,7 +10374,7 @@ dependencies = [ [[package]] name = "tantivy-columnar" version = "0.7.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=8a4cfc40f5b7446c0700471aa50290dbb7f8ea97#8a4cfc40f5b7446c0700471aa50290dbb7f8ea97" dependencies = [ "downcast-rs", "fastdivide", @@ -10388,7 +10390,7 @@ dependencies = [ [[package]] name = "tantivy-common" version = "0.11.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=8a4cfc40f5b7446c0700471aa50290dbb7f8ea97#8a4cfc40f5b7446c0700471aa50290dbb7f8ea97" dependencies = [ "async-trait", "byteorder", @@ -10424,7 +10426,7 @@ dependencies = [ [[package]] name = "tantivy-query-grammar" version = "0.26.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=8a4cfc40f5b7446c0700471aa50290dbb7f8ea97#8a4cfc40f5b7446c0700471aa50290dbb7f8ea97" dependencies = [ "fnv", "nom 7.1.3", @@ -10436,7 +10438,7 @@ dependencies = [ [[package]] name = "tantivy-sstable" version = "0.7.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=8a4cfc40f5b7446c0700471aa50290dbb7f8ea97#8a4cfc40f5b7446c0700471aa50290dbb7f8ea97" dependencies = [ "futures-util", "itertools 0.14.0", @@ -10449,7 +10451,7 @@ dependencies = [ [[package]] name = "tantivy-stacker" version = "0.7.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=8a4cfc40f5b7446c0700471aa50290dbb7f8ea97#8a4cfc40f5b7446c0700471aa50290dbb7f8ea97" dependencies = [ "murmurhash32", "tantivy-common", @@ -10458,7 +10460,7 @@ dependencies = [ [[package]] name = "tantivy-tokenizer-api" version = "0.7.0" -source = "git+https://github.com/quickwit-oss/tantivy/?rev=2e87c2eb0c8577c96c82a777f345bbf236a985fe#2e87c2eb0c8577c96c82a777f345bbf236a985fe" +source = "git+https://github.com/quickwit-oss/tantivy/?rev=8a4cfc40f5b7446c0700471aa50290dbb7f8ea97#8a4cfc40f5b7446c0700471aa50290dbb7f8ea97" dependencies = [ "serde", ] @@ -10479,7 +10481,7 @@ dependencies = [ "getrandom 0.4.3", "once_cell", "rustix 1.1.4", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] diff --git a/quickwit/Cargo.toml b/quickwit/Cargo.toml index a58d5f3797f..bbba90e5d4b 100644 --- a/quickwit/Cargo.toml +++ b/quickwit/Cargo.toml @@ -379,7 +379,7 @@ quickwit-telemetry-exporters = { path = "quickwit-telemetry-exporters" } quickwit-transport = { path = "quickwit-transport" } # This is the tip of tantivy main. -tantivy = { git = "https://github.com/quickwit-oss/tantivy/", rev = "2e87c2eb0c8577c96c82a777f345bbf236a985fe", default-features = false, features = [ +tantivy = { git = "https://github.com/quickwit-oss/tantivy/", rev = "8a4cfc40f5b7446c0700471aa50290dbb7f8ea97", default-features = false, features = [ "jitexpr", "lz4-compression", "mmap", @@ -388,7 +388,7 @@ tantivy = { git = "https://github.com/quickwit-oss/tantivy/", rev = "2e87c2eb0c8 "columnar-zstd-compression", ] } # Must stay on the same rev as `tantivy`. -tantivy-common = { git = "https://github.com/quickwit-oss/tantivy/", rev = "2e87c2eb0c8577c96c82a777f345bbf236a985fe" } +tantivy-common = { git = "https://github.com/quickwit-oss/tantivy/", rev = "8a4cfc40f5b7446c0700471aa50290dbb7f8ea97" } tantivy-fst = "0.5" # This is actually not used directly the goal is to fix the version diff --git a/quickwit/quickwit-proto/protos/quickwit/search.proto b/quickwit/quickwit-proto/protos/quickwit/search.proto index a1a10efbe09..b0a7fca1e2e 100644 --- a/quickwit/quickwit-proto/protos/quickwit/search.proto +++ b/quickwit/quickwit-proto/protos/quickwit/search.proto @@ -285,6 +285,19 @@ message SearchRequest { // Scheduling priority for leaf search execution. Negative values are allowed, // and lower values have higher priority. Callers that omit it get priority 0. int32 priority = 21; + + // Calculated fields: fields computed from fast fields with a jitexpr expression. + // Aggregations can use them as field names. + repeated CalculatedField calculated_fields = 22; +} + +message CalculatedField { + // Name of the calculated field, as used in the aggregation request. + // A calculated field shadows a fast field with the same name. + string name = 1; + // Jitexpr expression, serialized as a string. For example: `(MULTIPLY price stock)`. + // Variables refer to fast fields. + string expression = 2; } enum CountHits { diff --git a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs index ad0c1099cf7..683a96771f6 100644 --- a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs +++ b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs @@ -222,6 +222,22 @@ pub struct SearchRequest { #[prost(int32, tag = "21")] #[serde(default)] pub priority: i32, + /// Calculated fields: fields computed from fast fields with a jitexpr expression. + /// Aggregations can use them as field names. + #[prost(message, repeated, tag = "22")] + pub calculated_fields: ::prost::alloc::vec::Vec, +} +#[derive(serde::Serialize, serde::Deserialize, utoipa::ToSchema)] +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct CalculatedField { + /// Name of the calculated field, as used in the aggregation request. + /// A calculated field shadows a fast field with the same name. + #[prost(string, tag = "1")] + pub name: ::prost::alloc::string::String, + /// Jitexpr expression, serialized as a string. For example: `(MULTIPLY price stock)`. + /// Variables refer to fast fields. + #[prost(string, tag = "2")] + pub expression: ::prost::alloc::string::String, } #[derive(serde::Serialize, serde::Deserialize, utoipa::ToSchema)] #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] diff --git a/quickwit/quickwit-search/src/calculated_fields.rs b/quickwit/quickwit-search/src/calculated_fields.rs new file mode 100644 index 00000000000..78b25f26852 --- /dev/null +++ b/quickwit/quickwit-search/src/calculated_fields.rs @@ -0,0 +1,194 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Calculated fields: fields computed from fast fields with a jitexpr expression. +//! +//! They are defined in [`SearchRequest::calculated_fields`] and can be used as field names in +//! the aggregation request. A calculated field shadows a fast field with the same name. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use quickwit_proto::search::{CalculatedField, SearchRequest}; +use tantivy::aggregation::{JitExprValueSourceProvider, ValueSourceRegistry}; +use tantivy::jitexpr::ast::{UntypedExpr, infer_types}; + +use crate::SearchError; + +/// The calculated fields of a search request, parsed and type checked. +#[derive(Default)] +pub(crate) struct CalculatedFields { + expressions: HashMap, +} + +impl CalculatedFields { + /// Parses and type checks the calculated fields of `search_request`. + pub fn from_search_request(search_request: &SearchRequest) -> crate::Result { + Self::parse(&search_request.calculated_fields) + } + + fn parse(calculated_fields: &[CalculatedField]) -> crate::Result { + let mut expressions: HashMap = + HashMap::with_capacity(calculated_fields.len()); + for calculated_field in calculated_fields { + if calculated_field.name.is_empty() { + return Err(SearchError::InvalidArgument( + "calculated field name must not be empty".to_string(), + )); + } + let expression: UntypedExpr = tantivy::jitexpr::ast::deserialize( + &calculated_field.expression, + ) + .map_err(|deserialize_error| { + SearchError::InvalidArgument(format!( + "invalid expression for calculated field `{}`: {deserialize_error}", + calculated_field.name + )) + })?; + if let Err(type_error) = infer_types(&expression) { + return Err(SearchError::InvalidArgument(format!( + "invalid expression for calculated field `{}`: {type_error}", + calculated_field.name + ))); + } + if expressions + .insert(calculated_field.name.clone(), expression) + .is_some() + { + return Err(SearchError::InvalidArgument(format!( + "calculated field `{}` is defined several times", + calculated_field.name + ))); + } + } + Ok(CalculatedFields { expressions }) + } + + pub fn is_empty(&self) -> bool { + self.expressions.is_empty() + } + + /// Builds the registry through which tantivy aggregations resolve calculated fields. + pub fn value_source_registry(&self) -> crate::Result> { + let mut registry = ValueSourceRegistry::default(); + for (name, expression) in &self.expressions { + // Cannot fail in practice: the expression was type checked in `parse`. + let provider = + JitExprValueSourceProvider::new(expression.clone()).map_err(|type_error| { + SearchError::InvalidArgument(format!( + "invalid expression for calculated field `{name}`: {type_error}" + )) + })?; + registry.register(name, Arc::new(provider)); + } + Ok(Arc::new(registry)) + } + + /// Replaces the calculated fields of `field_names` by the fast fields their expressions read. + /// + /// Used to warm up the columns actually read by aggregations. + pub fn resolve_fast_field_names(&self, field_names: HashSet) -> HashSet { + let mut fast_field_names: HashSet = HashSet::with_capacity(field_names.len()); + for field_name in field_names { + let Some(expression) = self.expressions.get(&field_name) else { + fast_field_names.insert(field_name); + continue; + }; + // Cannot fail: the expression was type checked in `parse`. + let Ok(inferred_types) = infer_types(expression) else { + continue; + }; + fast_field_names.extend( + inferred_types + .into_keys() + .map(|variable_name| variable_name.to_string()), + ); + } + fast_field_names + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn calculated_field(name: &str, expression: &str) -> CalculatedField { + CalculatedField { + name: name.to_string(), + expression: expression.to_string(), + } + } + + #[test] + fn test_parse_calculated_fields() { + let calculated_fields = CalculatedFields::parse(&[ + calculated_field("doubled", "(MULTIPLY latency 2i64)"), + calculated_field("upper", "(UPPER service)"), + ]) + .unwrap(); + assert!(calculated_fields.expressions.contains_key("doubled")); + assert!(calculated_fields.expressions.contains_key("upper")); + assert!(!calculated_fields.expressions.contains_key("latency")); + calculated_fields.value_source_registry().unwrap(); + } + + #[test] + fn test_parse_calculated_fields_errors() { + for (calculated_fields, expected_message) in [ + ( + vec![calculated_field("", "(UPPER service)")], + "must not be empty", + ), + ( + vec![calculated_field("bad", "(UPPER service")], + "invalid expression for calculated field `bad`", + ), + ( + vec![calculated_field("bad", "(UPPER 1i64)")], + "invalid expression for calculated field `bad`", + ), + ( + vec![ + calculated_field("twice", "(UPPER service)"), + calculated_field("twice", "(LOWER service)"), + ], + "defined several times", + ), + ] { + let Err(error) = CalculatedFields::parse(&calculated_fields) else { + panic!("expected an error for {calculated_fields:?}"); + }; + assert!( + error.to_string().contains(expected_message), + "{calculated_fields:?}: {error}" + ); + } + } + + #[test] + fn test_resolve_fast_field_names() { + let calculated_fields = + CalculatedFields::parse(&[calculated_field("ratio", "(DIVIDE latency duration)")]) + .unwrap(); + let field_names: HashSet = ["ratio".to_string(), "service".to_string()].into(); + let fast_field_names = calculated_fields.resolve_fast_field_names(field_names); + let expected: HashSet = [ + "latency".to_string(), + "duration".to_string(), + "service".to_string(), + ] + .into(); + assert_eq!(fast_field_names, expected); + } +} diff --git a/quickwit/quickwit-search/src/collector.rs b/quickwit/quickwit-search/src/collector.rs index 08ab06e6a8f..ab21241ae0f 100644 --- a/quickwit/quickwit-search/src/collector.rs +++ b/quickwit/quickwit-search/src/collector.rs @@ -15,6 +15,7 @@ use std::borrow::Cow; use std::cmp::Ordering; use std::collections::HashSet; +use std::sync::Arc; use itertools::Itertools; use quickwit_common::binary_heap::{SortKeyMapper, TopK}; @@ -34,6 +35,7 @@ use tantivy::fastfield::Column; use tantivy::tokenizer::TokenizerManager; use tantivy::{DateTime, DocId, Score, SegmentOrdinal, SegmentReader, TantivyError}; +use crate::calculated_fields::CalculatedFields; use crate::find_trace_ids_collector::{FindTraceIdsCollector, FindTraceIdsSegmentCollector, Span}; use crate::top_k_collector::{QuickwitSegmentTopKCollector, specialized_top_k_segment_collector}; use crate::{GlobalDocAddress, add_leaf_stats, merge_leaf_stats_it}; @@ -716,6 +718,9 @@ pub(crate) struct QuickwitCollector { pub sort_by: SortByPair, pub aggregation: Option, pub agg_context_params: AggContextParams, + /// Calculated fields the aggregations can refer to. Their value sources are registered in + /// `agg_context_params`. + calculated_fields: Arc, search_after: Option, } @@ -739,7 +744,12 @@ impl QuickwitCollector { sort_by_second.add_fast_field(&mut fast_field_names); } if let Some(aggregations) = &self.aggregation { - fast_field_names.extend(aggregations.fast_field_names()); + // Calculated fields are not columns: we need the fast fields their expressions read. + let agg_field_names = aggregations.fast_field_names(); + fast_field_names.extend( + self.calculated_fields + .resolve_fast_field_names(agg_field_names), + ); } fast_field_names } @@ -1042,6 +1052,12 @@ pub(crate) fn make_collector_for_split( Some(aggregation) => Some(serde_json::from_str(aggregation)?), None => None, }; + let calculated_fields = CalculatedFields::from_search_request(search_request)?; + let agg_context_params = if calculated_fields.is_empty() { + agg_context_params + } else { + agg_context_params.with_value_sources(calculated_fields.value_source_registry()?) + }; let sort_by = sort_by_from_request(search_request); Ok(QuickwitCollector { split_id, @@ -1050,6 +1066,7 @@ pub(crate) fn make_collector_for_split( sort_by, aggregation, agg_context_params, + calculated_fields: Arc::new(calculated_fields), search_after: search_request.search_after.clone(), }) } @@ -1075,6 +1092,8 @@ pub(crate) fn make_merge_collector( sort_by, aggregation, agg_context_params, + // Merging intermediate results does not read any column. + calculated_fields: Arc::default(), search_after: search_request.search_after.clone(), }) } @@ -1317,8 +1336,8 @@ mod tests { use std::cmp::Ordering; use quickwit_proto::search::{ - LeafResourceStats, LeafSearchResponse, PartialHit, SearchRequest, SortByValue, SortField, - SortOrder, SortValue, SplitResourceStats, SplitSearchError, + CalculatedField, LeafResourceStats, LeafSearchResponse, PartialHit, SearchRequest, + SortByValue, SortField, SortOrder, SortValue, SplitResourceStats, SplitSearchError, }; use quickwit_proto::types::SplitId; use tantivy::TantivyDocument; @@ -2144,4 +2163,83 @@ mod tests { // merged response to segment_size (4). assert_eq!(buckets.entries().len(), 4); } + + #[test] + fn test_warmup_info_resolves_calculated_fields() { + let request = SearchRequest { + max_hits: 0, + aggregation_request: Some( + r##"{ + "by_level": { + "terms": { "field": "#upper_level" }, + "aggs": { "total": { "sum": { "field": "#ratio" } } } + }, + "by_host": { "terms": { "field": "host" } } + }"## + .to_string(), + ), + calculated_fields: vec![ + CalculatedField { + name: "#upper_level".to_string(), + expression: "(UPPER level)".to_string(), + }, + CalculatedField { + name: "#ratio".to_string(), + expression: "(DIVIDE latency duration)".to_string(), + }, + // Not used by the aggregations: its input must not be warmed up. + CalculatedField { + name: "#unused".to_string(), + expression: "(UPPER unused_input)".to_string(), + }, + ], + ..Default::default() + }; + let collector = super::make_collector_for_split( + SplitId::from("fake_split_id"), + &request, + Default::default(), + ) + .unwrap(); + let mut warmed_up_fields: Vec<(String, bool)> = collector + .warmup_info() + .fast_fields + .into_iter() + .map(|fast_field| (fast_field.name, fast_field.with_subfields)) + .collect(); + warmed_up_fields.sort(); + assert_eq!( + warmed_up_fields, + vec![ + ("duration".to_string(), false), + ("host".to_string(), false), + ("latency".to_string(), false), + ("level".to_string(), false), + ] + ); + } + + #[test] + fn test_make_collector_rejects_invalid_calculated_field() { + let request = SearchRequest { + calculated_fields: vec![CalculatedField { + name: "#bad".to_string(), + expression: "(UPPER level".to_string(), + }], + ..Default::default() + }; + let Err(error) = super::make_collector_for_split( + SplitId::from("fake_split_id"), + &request, + Default::default(), + ) else { + panic!("expected an error"); + }; + assert!( + error + .to_string() + .contains("invalid expression for calculated field `#bad`"), + "{error}" + ); + } } diff --git a/quickwit/quickwit-search/src/lib.rs b/quickwit/quickwit-search/src/lib.rs index 64b28776bd5..75c835aa81f 100644 --- a/quickwit/quickwit-search/src/lib.rs +++ b/quickwit/quickwit-search/src/lib.rs @@ -17,6 +17,7 @@ #![allow(clippy::bool_assert_comparison)] #![deny(clippy::disallowed_methods)] +mod calculated_fields; mod client; mod cluster_client; mod collector; diff --git a/quickwit/quickwit-search/src/root.rs b/quickwit/quickwit-search/src/root.rs index e316c60bf85..4e02e8422cb 100644 --- a/quickwit/quickwit-search/src/root.rs +++ b/quickwit/quickwit-search/src/root.rs @@ -50,6 +50,7 @@ use tantivy::schema::{Field, FieldEntry, FieldType, Schema}; use tracing::{Span, debug, error, info, info_span, instrument}; use tracing_opentelemetry::OpenTelemetrySpanExt; +use crate::calculated_fields::CalculatedFields; use crate::cluster_client::ClusterClient; use crate::collector::{QuickwitAggregations, make_merge_collector}; use crate::cost::{compute_query_complexity_factor, compute_split_query_cost}; @@ -366,6 +367,8 @@ fn simplify_search_request_for_scroll_api(req: &SearchRequest) -> crate::Result< // We remove all aggregation request. // The aggregation will not be computed for each scroll request. aggregation_request: None, + // Calculated fields are only used by aggregations. + calculated_fields: Vec::new(), // We remove the snippet fields. This feature is not supported for scroll requests. snippet_fields: Vec::new(), // We remove the scroll ttl parameter. It is irrelevant to process later request @@ -519,6 +522,9 @@ fn validate_request( validate_requested_snippet_fields(schema, &search_request.snippet_fields)?; + // Rejects invalid calculated fields before the request reaches the leaves. + let calculated_fields = CalculatedFields::from_search_request(search_request)?; + if let Some(agg) = search_request.aggregation_request.as_ref() { let aggs: QuickwitAggregations = serde_json::from_str(agg).map_err(|_err| { let err = serde_json::from_str::(agg) @@ -527,7 +533,8 @@ fn validate_request( })?; // ensure that the required fast fields are indeed configured as fast fields. - let fast_field_names = aggs.fast_field_names(); + // Calculated fields are replaced by the fast fields their expressions read. + let fast_field_names = calculated_fields.resolve_fast_field_names(aggs.fast_field_names()); let dynamic_field = schema.get_field(DYNAMIC_FIELD_NAME).ok(); for fast_field_name in &fast_field_names { check_is_fast_field(schema, fast_field_name, dynamic_field)?; diff --git a/quickwit/quickwit-search/src/tests.rs b/quickwit/quickwit-search/src/tests.rs index d15c14a3cec..c9599f9efd4 100644 --- a/quickwit/quickwit-search/src/tests.rs +++ b/quickwit/quickwit-search/src/tests.rs @@ -22,8 +22,8 @@ use quickwit_doc_mapper::DocMapper; use quickwit_doc_mapper::tag_pruning::extract_tags_from_query; use quickwit_indexing::TestSandbox; use quickwit_proto::search::{ - CountHits, LeafListTermsResponse, ListTermsRequest, SearchRequest, SortByValue, SortField, - SortOrder, SortValue, TraceId, + CalculatedField, CountHits, LeafListTermsResponse, ListTermsRequest, SearchRequest, + SortByValue, SortField, SortOrder, SortValue, TraceId, }; use quickwit_query::query_ast::{ BoolQuery, HitSet, PredicateCache, QueryAst, RangeQuery, qast_helper, qast_json_helper, @@ -1578,6 +1578,136 @@ async fn test_single_node_aggregation() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +async fn test_single_node_aggregation_on_calculated_fields() { + let index_id = "single-node-agg-calculated-fields"; + let doc_mapping_yaml = r#" + field_mappings: + - name: color + type: text + tokenizer: raw + fast: true + - name: price + type: f64 + fast: true + - name: quantity + type: u64 + fast: true + "#; + let test_sandbox = TestSandbox::create(index_id, doc_mapping_yaml, "{}", &["color"]) + .await + .unwrap(); + let docs = vec![ + json!({"color": "blue", "price": 10.0, "quantity": 2}), + json!({"color": "Blue", "price": 15.0, "quantity": 1}), + json!({"color": "green", "price": 10.0, "quantity": 3}), + // No quantity: the calculated `#total` has no value for this document. + json!({"color": "green", "price": 100.0}), + ]; + test_sandbox.add_documents(docs).await.unwrap(); + let agg_req = r##" + { + "colors": { + "terms": { "field": "#upper_color", "order": { "_count": "desc" } }, + "aggs": { + "total": { "sum": { "field": "#total" } } + } + } + }"##; + let search_request = SearchRequest { + index_id_patterns: vec![index_id.to_string()], + query_ast: qast_json_helper("*", &[]), + max_hits: 0, + aggregation_request: Some(agg_req.to_string()), + calculated_fields: vec![ + CalculatedField { + name: "#upper_color".to_string(), + expression: "(UPPER color)".to_string(), + }, + CalculatedField { + name: "#total".to_string(), + expression: "(MULTIPLY price quantity)".to_string(), + }, + ], + ..Default::default() + }; + let single_node_result = single_node_search( + search_request, + test_sandbox.metastore(), + test_sandbox.storage_resolver(), + ) + .await + .unwrap(); + let agg_res_struct = + AggregationResults::from_postcard(&single_node_result.aggregation_postcard.unwrap()) + .unwrap(); + let agg_res_json: JsonValue = serde_json::to_value(&agg_res_struct).unwrap(); + let mut buckets: Vec<(String, u64, f64)> = agg_res_json["colors"]["buckets"] + .as_array() + .unwrap() + .iter() + .map(|bucket| { + ( + bucket["key"].as_str().unwrap().to_string(), + bucket["doc_count"].as_u64().unwrap(), + bucket["total"]["value"].as_f64().unwrap(), + ) + }) + .collect(); + buckets.sort_by(|left, right| left.0.cmp(&right.0)); + assert_eq!( + buckets, + vec![ + ("BLUE".to_string(), 2, 35.0), + ("GREEN".to_string(), 2, 30.0), + ] + ); + test_sandbox.assert_quit().await; +} + +#[tokio::test] +async fn test_single_node_aggregation_calculated_field_requires_fast_inputs() { + let index_id = "single-node-agg-calculated-fields-not-fast"; + let doc_mapping_yaml = r#" + field_mappings: + - name: color + type: text + "#; + let test_sandbox = TestSandbox::create(index_id, doc_mapping_yaml, "{}", &["color"]) + .await + .unwrap(); + test_sandbox + .add_documents(vec![json!({"color": "blue"})]) + .await + .unwrap(); + let search_request = SearchRequest { + index_id_patterns: vec![index_id.to_string()], + query_ast: qast_json_helper("*", &[]), + aggregation_request: Some( + r##"{ "colors": { "terms": { "field": "#upper_color" } } }"##.to_string(), + ), + calculated_fields: vec![CalculatedField { + name: "#upper_color".to_string(), + expression: "(UPPER color)".to_string(), + }], + ..Default::default() + }; + let single_node_error = single_node_search( + search_request, + test_sandbox.metastore(), + test_sandbox.storage_resolver(), + ) + .await + .unwrap_err(); + assert!( + single_node_error + .to_string() + .contains("Field \"color\" is not configured as a fast field"), + "{single_node_error}" + ); + test_sandbox.assert_quit().await; +} + #[tokio::test] async fn test_single_node_aggregation_missing_fast_field() { let index_id = "single-node-agg-2";