Adding a Jitexpr ValueSource to allow jitexpr in aggregations.

Cargo.toml

@@ -78,6 +78,8 @@ bon = "3.3.1"
 
 # EXPERIMENTAL. The API is likely to change in the near future.
 jitexpr = { version = "0.1", path = "./jitexpr", optional = true }
+# Used by the jitexpr value source to intern the computed terms.
+hashbrown = { version = "0.17", default-features = false, optional = true }
 
 columnar = { version = "0.7", path = "./columnar", package = "tantivy-columnar" }
 sstable = { version = "0.7", path = "./sstable", package = "tantivy-sstable", optional = true }
@@ -152,7 +154,7 @@ failpoints = ["fail", "fail/failpoints"]
 unstable = []                            # useful for benches.
 
 quickwit = ["sstable", "futures-util", "futures-channel"]
-jitexpr = ["dep:jitexpr"]
+jitexpr = ["dep:jitexpr", "dep:hashbrown"]
 
 # Compares only the hash of a string when indexing data.
 # Increases indexing speed, but may lead to extremely rare missing terms, when there's a hash collision.
@@ -184,6 +186,10 @@ name = "failpoints"
 path = "tests/failpoints/mod.rs"
 required-features = ["failpoints"]
 
+[[example]]
+name = "aggregation_jitexpr"
+required-features = ["jitexpr"]
+
 [[bench]]
 name = "analyzer"
 harness = false

columnar/src/column/dictionary_encoded.rs

@@ -100,6 +100,11 @@ impl StrColumn {
         self.0.dictionary.as_ref()
     }
 
+    /// Splits the column into its dictionary and its column of term ords.
+    pub fn into_parts(self) -> (Arc<Dictionary<VoidSSTable>>, Column<u64>) {
+        (self.0.dictionary, self.0.term_ord_column)
+    }
+
     /// Fills the buffer
     pub fn ord_to_str(&self, term_ord: u64, output: &mut String) -> io::Result<bool> {
         unsafe {

examples/aggregation_jitexpr.rs

@@ -0,0 +1,86 @@
+// # Jitexpr fields in aggregations
+//
+// This example shows how aggregations can use a jitexpr field: a field that is not in the
+// index, but computed from the fast fields of each document with a jitexpr expression.
+//
+// - A jitexpr field is registered under a name in a `ValueSourceRegistry`
+// - Aggregations use that name like the name of a fast field.
+//
+// Requires the `jitexpr` feature.
+// ---
+
+use std::sync::Arc;
+
+use serde_json::{json, Value};
+use tantivy::aggregation::agg_req::Aggregations;
+use tantivy::aggregation::{
+    AggContextParams, AggregationCollector, JitExprValueSourceProvider, ValueSourceRegistry,
+};
+use tantivy::query::AllQuery;
+use tantivy::schema::{Schema, FAST, STRING};
+use tantivy::{doc, Index, IndexWriter};
+
+fn main() -> tantivy::Result<()> {
+    // # Create Schema
+    //
+    // Aggregations read fast fields: the inputs of the jitexpr field must be fast.
+    let mut schema_builder = Schema::builder();
+    let category = schema_builder.add_text_field("category", STRING | FAST);
+    let price = schema_builder.add_f64_field("price", FAST);
+    let stock = schema_builder.add_f64_field("stock", FAST);
+    let schema = schema_builder.build();
+
+    // # Indexing documents
+    let index = Index::create_in_ram(schema);
+    let mut index_writer: IndexWriter = index.writer(50_000_000)?;
+    index_writer.add_document(doc!(category => "Womens Footwear", price => 99.0, stock => 5.0))?;
+    index_writer.add_document(doc!(category => "Womens Footwear", price => 59.0, stock => 8.0))?;
+    index_writer.add_document(doc!(category => "Mens Footwear", price => 120.0, stock => 2.0))?;
+    index_writer.add_document(doc!(category => "Mens Footwear", price => 75.0, stock => 10.0))?;
+    // A document without a price: the jitexpr field has no value for it.
+    index_writer.add_document(doc!(category => "Mens Footwear", stock => 3.0))?;
+    index_writer.commit()?;
+
+    // # Registering the jitexpr field
+    //
+    // The expression is written in the jitexpr s-expression syntax. Variables (e.g. `price`) are
+    // resolved as fast field names, for each segment.
+    //
+    // A jitexpr field shadows a fast field with the same name. Here, the name starts with `#`
+    // so that it cannot collide with a fast field.
+    let expression = tantivy::jitexpr::ast::deserialize("(MULTIPLY price stock)")
+        .map_err(|err| tantivy::TantivyError::InvalidArgument(format!("{err:?}")))?;
+    let provider = JitExprValueSourceProvider::new(expression)
+        .map_err(|err| tantivy::TantivyError::InvalidArgument(format!("{err:?}")))?;
+    let mut value_sources = ValueSourceRegistry::default();
+    value_sources.register("#stock_value", Arc::new(provider));
+    let context = AggContextParams::default().with_value_sources(Arc::new(value_sources));
+
+    // # Aggregating over the jitexpr field
+    //
+    // The total stock value of each category.
+    let aggs: Aggregations = serde_json::from_value(json!({
+        "categories": {
+            "terms": { "field": "category" },
+            "aggs": {
+                "total_stock_value": { "sum": { "field": "#stock_value" } }
+            }
+        }
+    }))?;
+    let collector = AggregationCollector::from_aggs(aggs, context);
+
+    let searcher = index.reader()?.searcher();
+    let agg_res = searcher.search(&AllQuery, &collector)?;
+    let res: Value = serde_json::to_value(agg_res)?;
+    println!("{}", serde_json::to_string_pretty(&res)?);
+
+    let buckets = &res["categories"]["buckets"];
+    assert_eq!(buckets[0]["key"], "Mens Footwear");
+    // 120 * 2 + 75 * 10. The document without a price has no stock value.
+    assert_eq!(buckets[0]["total_stock_value"]["value"], 990.0);
+    assert_eq!(buckets[1]["key"], "Womens Footwear");
+    // 99 * 5 + 59 * 8.
+    assert_eq!(buckets[1]["total_stock_value"]["value"], 967.0);
+
+    Ok(())
+}

jitexpr/src/compile/compile_fn_builder.rs

@@ -202,6 +202,24 @@ impl<'types, 'names> CompileFnBuilder<'types, 'names> {
         self.lower_typed_expr(expression)?.into_compiled_fn()
     }
 
+    /// Compiles an expression whose variable ids are already assigned.
+    ///
+    /// Precondition: `input_vars` are the inputs of `expression`, as assigned by
+    /// `assign_variable_ids`.
+    pub(super) fn compile_with_inputs(
+        input_vars: Vec<TypedVariable>,
+        expression: TypedExpr,
+    ) -> Result<CompiledFn, CompileError> {
+        // The variable types are only used to build a typed expression: `expression` is already
+        // typed.
+        let variable_types: HashMap<&str, VarType> = HashMap::new();
+        let builder = CompileFnBuilder {
+            variable_types: &variable_types,
+            input_vars,
+        };
+        builder.compile_typed_expr(expression)
+    }
+
     pub(super) fn compile_typed_expr_to_assembly(
         self,
         expression: TypedExpr,

jitexpr/src/compile/compiled_fn.rs

@@ -1,9 +1,9 @@
 use std::ops::Deref;
-use std::sync::Arc;
+use std::sync::{Arc, OnceLock};
 
 use cranelift_jit::JITModule;
 
-use super::{StringArena, TypedExpr, TypedVariable};
+use super::{CompileError, CompileFnBuilder, StringArena, TypedExpr, TypedVariable};
 use crate::types::{VarType, VariableValue};
 
 #[cfg(not(any(
@@ -36,7 +36,9 @@ pub struct CompiledFn {
     module: Option<JITModule>,
     inputs: Vec<TypedVariable>,
     // This AST owns the Arc-backed literals and regexes embedded in generated code.
-    _typed_expr: Box<TypedExpr>,
+    typed_expr: Box<TypedExpr>,
+    // Memoizes `coerce_result_to_f64`.
+    coerced_to_f64: OnceLock<Result<Arc<CompiledFn>, CompileError>>,
 }
 
 // `JITModule` is not `Sync` because it supports lazily looking up symbols through
@@ -69,7 +71,8 @@ impl CompiledFn {
             entry,
             module: Some(module),
             inputs,
-            _typed_expr: typed_expr,
+            typed_expr,
+            coerced_to_f64: OnceLock::new(),
         }
     }
 
@@ -80,7 +83,32 @@ impl CompiledFn {
 
     /// Returns the concrete result type selected during compilation.
     pub fn result_type(&self) -> VarType {
-        self._typed_expr.return_type
+        self.typed_expr.return_type
+    }
+
+    /// Returns the same function, with its result coerced to `f64`.
+    ///
+    /// Returns an error if the result type is not `I64`, `U64` or `F64`. For `F64`, returns `self`.
+    ///
+    /// The coerced function is compiled on the first call, and memoized.
+    pub fn coerce_result_to_f64(self: &Arc<Self>) -> Result<Arc<CompiledFn>, CompileError> {
+        let result_type = self.result_type();
+        match result_type {
+            VarType::F64 => return Ok(Arc::clone(self)),
+            VarType::I64 | VarType::U64 => {}
+            VarType::Bool | VarType::Str | VarType::None => {
+                return Err(CompileError::UnsupportedCoercion {
+                    from_type: result_type,
+                    target: VarType::F64,
+                });
+            }
+        }
+        self.coerced_to_f64
+            .get_or_init(|| {
+                let typed_expr: TypedExpr = (*self.typed_expr).clone().coerce(VarType::F64);
+                CompileFnBuilder::compile_with_inputs(self.inputs.clone(), typed_expr).map(Arc::new)
+            })
+            .clone()
     }
 
     /// Creates an evaluation context with a private string arena.

jitexpr/src/compile/mod.rs

@@ -253,6 +253,42 @@ mod tests {
         assert_eq!(string_arena.used_bytes(), 1);
     }
 
+    #[test]
+    fn test_coerce_result_to_f64() {
+        let untyped_expr = crate::ast::deserialize("(ADD number 1i64)").unwrap();
+        let variable_types = HashMap::from([("number", VarType::I64)]);
+        let compiled_fn = compile(&untyped_expr, &variable_types).unwrap();
+        assert_eq!(compiled_fn.result_type(), VarType::I64);
+
+        let coerced_fn = compiled_fn.coerce_result_to_f64().unwrap();
+        assert_eq!(coerced_fn.result_type(), VarType::F64);
+        assert!(coerced_fn.inputs() == compiled_fn.inputs());
+        let mut string_arena = StringArena::new();
+        let output = unsafe { coerced_fn.call(&[VariableValue::some(-3i64)], &mut string_arena) };
+        assert_eq!(unsafe { output.as_f64() }, Some(-2.0));
+        // A null result stays null.
+        let output = unsafe { coerced_fn.call(&[VariableValue::none()], &mut string_arena) };
+        assert_eq!(unsafe { output.as_f64() }, None);
+
+        // Memoized.
+        assert!(Arc::ptr_eq(
+            &coerced_fn,
+            &compiled_fn.coerce_result_to_f64().unwrap()
+        ));
+        // Already `f64`.
+        assert!(Arc::ptr_eq(
+            &coerced_fn,
+            &coerced_fn.coerce_result_to_f64().unwrap()
+        ));
+        // Not numerical.
+        let str_fn = compile(
+            &UntypedExpr::variable("label"),
+            &HashMap::from([("label", VarType::Str)]),
+        )
+        .unwrap();
+        assert!(str_fn.coerce_result_to_f64().is_err());
+    }
+
     #[test]
     fn test_compiled_fn_is_send_and_sync() {
         fn assert_send_and_sync<T: Send + Sync>() {}

src/aggregation/accessor_helpers.rs

@@ -144,15 +144,20 @@ pub(crate) fn get_all_value_sources(
     Ok(sources)
 }
 
-/// Builds the value source of a field without any column, with the given type.
+/// Builds a value source without any value, with the given type.
+///
+/// `num_rows` is the number of rows of the column (the `max_doc` of the segment).
 ///
 /// A `Str` shim is an empty `StrColumn`, so that it has an (empty) term dictionary, as required by
 /// `ValueSource::term_dictionary`.
-fn build_empty_value_source(num_docs: u32, column_type: ColumnType) -> Box<dyn ValueSource> {
+pub(crate) fn build_empty_value_source(
+    num_rows: u32,
+    column_type: ColumnType,
+) -> Box<dyn ValueSource> {
     if column_type == ColumnType::Str {
-        Box::new(StrColumn::wrap(BytesColumn::empty(num_docs)))
+        Box::new(StrColumn::wrap(BytesColumn::empty(num_rows)))
     } else {
-        Box::new((Column::build_empty_column(num_docs), column_type))
+        Box::new((Column::build_empty_column(num_rows), column_type))
     }
 }
 

src/aggregation/agg_data.rs

@@ -16,8 +16,8 @@ use crate::aggregation::bucket::{
     build_segment_multi_terms_collector, build_segment_range_collector, CompositeAggReqData,
     CompositeAggregation, CompositeSourceAccessors, FilterAggReqData, HistogramAggReqData,
     IncludeExcludeParam, MissingTermAggReqData, MultiTermsAggReqData, MultiTermsAggregation,
-    MultiTermsFieldAccessor, MultiTermsMissingAccessor, RangeAggReqData, TermMissingAgg,
-    TermsAggReqData, TermsAggregation, TermsAggregationInternal,
+    MultiTermsFieldAccessor, MultiTermsMissingAccessor, OrderTarget, RangeAggReqData,
+    TermMissingAgg, TermsAggReqData, TermsAggregation, TermsAggregationInternal,
 };
 use crate::aggregation::metric::{
     build_segment_stats_collector, AverageAggregation, CardinalityAggReqData,
@@ -448,7 +448,14 @@ fn build_nodes(
                     ))
                 }
             };
-            let accessor = get_value_source(reader, value_sources, field, allowed_column_types)?;
+            let mut accessor =
+                get_value_source(reader, value_sources, field, allowed_column_types)?;
+            if missing.is_some() {
+                // The collectors encode `missing` with the column type of the source: a `U64`
+                // source would turn `-1` into `0`, an `I64` source `2.5` into `2`. They collect
+                // `f64` values anyway, so coercing the source to `F64` loses no precision.
+                accessor = accessor.to_f64();
+            }
             let field_type = accessor.column_type();
             let node_data = AggNodeData::Metric(MetricAggReqData {
                 accessor,
@@ -467,12 +474,16 @@ fn build_nodes(
         // Percentiles handled as Metric as well
         AggregationVariants::Percentiles(percentiles_req) => {
             percentiles_req.validate()?;
-            let accessor = get_value_source(
+            let mut accessor = get_value_source(
                 reader,
                 value_sources,
                 percentiles_req.field_name(),
                 Some(get_numeric_or_date_column_types()),
             )?;
+            if percentiles_req.missing.is_some() {
+                // See the metric aggregations above.
+                accessor = accessor.to_f64();
+            }
             let field_type = accessor.column_type();
             let node_data = AggNodeData::Metric(MetricAggReqData {
                 accessor,
@@ -1002,11 +1013,35 @@ fn build_terms_or_cardinality_nodes(
                         missing.is_some(),
                     )?;
                 };
+                // Some terms aggregation options are not supported on computed value sources. We
+                // return an explicit error for them.
+                let req_internal = TermsAggregationInternal::from_req(req);
+                if accessor.as_physical_column().is_none() {
+                    // `_key` ordering orders the term ords. The term ords of a computed text
+                    // source are assigned in first-seen order, not sorted with the terms. Numeric
+                    // and bool values use the monotonic `u64` mapping, which is sorted with the
+                    // values.
+                    if accessor.column_type() == ColumnType::Str
+                        && req_internal.order.target == OrderTarget::Key
+                    {
+                        return Err(crate::TantivyError::InvalidArgument(format!(
+                            "terms aggregation ordered by `_key` is not supported on the computed \
+                             value source `{field_name}`"
+                        )));
+                    }
+                    // `min_doc_count: 0` requires listing all of the possible values.
+                    if req_internal.min_doc_count == 0 {
+                        return Err(crate::TantivyError::InvalidArgument(format!(
+                            "terms aggregation with `min_doc_count: 0` is not supported on the \
+                             computed value source `{field_name}`"
+                        )));
+                    }
+                }
                 AggNodeData::Terms(TermsAggReqData {
                     accessor,
                     missing_value_for_accessor,
                     name: agg_name.to_string(),
-                    req: TermsAggregationInternal::from_req(req),
+                    req: req_internal,
                     sub_aggregations: sub_aggs.clone(),
                     allowed_term_ids,
                     is_top_level,

src/aggregation/metric/percentiles.rs

@@ -703,4 +703,26 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_percentiles_negative_missing_on_u64_column() {
+        let mut schema_builder = Schema::builder();
+        let score = schema_builder.add_u64_field("score", FAST);
+        let index = Index::create_in_ram(schema_builder.build());
+        let mut index_writer = index.writer_for_tests().unwrap();
+        index_writer.add_document(doc!(score => 10u64)).unwrap();
+        index_writer.add_document(doc!()).unwrap();
+        index_writer.commit().unwrap();
+
+        let agg_req: Aggregations = serde_json::from_value(json!({
+            "percentiles": { "percentiles": { "field": "score", "missing": -1.0 } }
+        }))
+        .unwrap();
+        let res = exec_request_with_query(agg_req, &index, None).unwrap();
+
+        // The doc without a value counts as -1, which a `u64` column cannot represent.
+        let first_percentile = res["percentiles"]["values"]["1.0"].as_f64().unwrap();
+        assert_le!(first_percentile, -0.95);
+        assert_ge!(first_percentile, -1.05);
+    }
 }