First stab

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.

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,
@@ -1002,11 +1002,17 @@ fn build_terms_or_cardinality_nodes(
                         missing.is_some(),
                     )?;
                 };
+                let req_internal = TermsAggregationInternal::from_req(req);
+                if let Some(term_dictionary) = accessor.term_dictionary() {
+                    if !term_dictionary.ords_sorted_with_terms() {
+                        check_terms_req_for_unsorted_ords(&req_internal, 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,
@@ -1027,6 +1033,30 @@ fn build_terms_or_cardinality_nodes(
     Ok(nodes)
 }
 
+/// Rejects the terms options that require term ords sorted with the terms.
+///
+/// Some computed text sources assign term ords in first-seen order:
+/// - with `_key` ordering, the segment cut-off selects the buckets by term ord.
+/// - `min_doc_count: 0` requires a dictionary listing the terms of all of the docs.
+fn check_terms_req_for_unsorted_ords(
+    req: &TermsAggregationInternal,
+    field_name: &str,
+) -> crate::Result<()> {
+    if req.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}`"
+        )));
+    }
+    if req.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}`"
+        )));
+    }
+    Ok(())
+}
+
 /// Builds a single BitSet of allowed term ordinals for a string dictionary column according to
 /// include/exclude parameters.
 ///

src/aggregation/mod.rs

@@ -147,6 +147,8 @@ use std::fmt::Display;
 use std::sync::Arc;
 
 pub(crate) use value_source::ColumnBlockAccessor;
+#[cfg(feature = "jitexpr")]
+pub use value_source::JitExprValueSourceProvider;
 pub use value_source::{
     ValueSource, ValueSourceDictionary, ValueSourceProvider, ValueSourceRegistry,
 };

src/aggregation/value_source/jitexpr_value_source/dynamic_term_dictionary.rs

@@ -0,0 +1,249 @@
+//! A term dictionary built on the fly, for computed text value sources.
+
+use std::hash::BuildHasher;
+use std::{fmt, io};
+
+use hashbrown::hash_table::Entry;
+use hashbrown::HashTable;
+use rustc_hash::FxBuildHasher;
+
+use crate::aggregation::value_source::ValueSourceDictionary;
+
+/// An append-only dictionary that interns terms and assigns them term ords.
+///
+/// Term ords are assigned in first-seen order. They are dense (`0..num_terms()`) and stable:
+/// interning new terms never changes the ord of a known term. They are NOT sorted with the terms.
+///
+/// All of the terms are concatenated in a single buffer. A separate offset index gives the
+/// boundaries of each term.
+///
+/// The hash table stores the term ord and a 32-bit hash of the term. This way, the table never
+/// reads the terms when it grows, and most of the comparisons with a different term stop at the
+/// hash. Equality always compares the term bytes: a hash collision costs an extra comparison, and
+/// never merges two terms.
+///
+/// Contract: at most `u32::MAX` terms. `intern` panics beyond that.
+pub(crate) struct DynamicTermDictionary {
+    /// Concatenation of all of the terms, in term ord order.
+    buffer: Vec<u8>,
+    /// `buffer[offsets[ord]..offsets[ord + 1]]` is the term of `ord`.
+    ///
+    /// Invariant: `offsets[0] == 0`, and `offsets.len() == num_terms() + 1`.
+    offsets: Vec<usize>,
+    /// One item per term, hashed with `table_hash(item.hash)`.
+    term_ords: HashTable<TermOrdItem>,
+}
+
+/// An item of the hash table: a term ord and the 32-bit hash of its term.
+#[derive(Clone, Copy)]
+struct TermOrdItem {
+    ord: u32,
+    hash: u32,
+}
+
+/// Returns the 32-bit hash of `term` stored in the hash table.
+#[inline]
+fn term_hash(term: &[u8]) -> u32 {
+    let hash: u64 = FxBuildHasher.hash_one(term);
+    // Both halves contribute to the folded hash.
+    (hash ^ (hash >> 32)) as u32
+}
+
+/// Returns the hash used by the hash table for an item with the 32-bit hash `hash`.
+///
+/// hashbrown takes the bucket index from the low bits of the hash, and its 7-bit tag from the top
+/// bits. The multiplication spreads the 32 bits of `hash` to the top bits: without it, all of the
+/// tags would be 0.
+///
+/// Contract: this must only depend on the stored hash, so that the table can compute it again when
+/// it grows without reading the term.
+#[inline]
+fn table_hash(hash: u32) -> u64 {
+    (hash as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
+}
+
+impl Default for DynamicTermDictionary {
+    fn default() -> Self {
+        DynamicTermDictionary {
+            buffer: Vec::new(),
+            offsets: vec![0],
+            term_ords: HashTable::new(),
+        }
+    }
+}
+
+impl fmt::Debug for DynamicTermDictionary {
+    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+        formatter
+            .debug_struct("DynamicTermDictionary")
+            .field("num_terms", &self.num_terms())
+            .finish()
+    }
+}
+
+/// Returns the term of `ord`.
+///
+/// Precondition: `ord < offsets.len() - 1`.
+#[inline]
+fn term_of<'a>(buffer: &'a [u8], offsets: &[usize], ord: usize) -> &'a [u8] {
+    &buffer[offsets[ord]..offsets[ord + 1]]
+}
+
+impl DynamicTermDictionary {
+    /// Returns the term ord of `term`. A new term gets the next term ord.
+    pub(crate) fn intern(&mut self, term: &[u8]) -> u64 {
+        let DynamicTermDictionary {
+            buffer,
+            offsets,
+            term_ords,
+        } = self;
+        let hash: u32 = term_hash(term);
+        let entry = term_ords.entry(
+            table_hash(hash),
+            |item| item.hash == hash && term_of(buffer, offsets, item.ord as usize) == term,
+            |item| table_hash(item.hash),
+        );
+        match entry {
+            Entry::Occupied(occupied_entry) => occupied_entry.get().ord as u64,
+            Entry::Vacant(vacant_entry) => {
+                let ord: u32 = u32::try_from(offsets.len() - 1)
+                    .expect("a dynamic term dictionary holds at most u32::MAX terms");
+                buffer.extend_from_slice(term);
+                offsets.push(buffer.len());
+                vacant_entry.insert(TermOrdItem { ord, hash });
+                ord as u64
+            }
+        }
+    }
+
+    /// Returns the term of `ord`, or `None` if `ord` is out of range.
+    pub(crate) fn term(&self, ord: u64) -> Option<&[u8]> {
+        if ord >= self.num_terms() {
+            return None;
+        }
+        Some(term_of(&self.buffer, &self.offsets, ord as usize))
+    }
+
+    pub(crate) fn num_terms(&self) -> u64 {
+        (self.offsets.len() - 1) as u64
+    }
+}
+
+impl ValueSourceDictionary for DynamicTermDictionary {
+    fn ords_sorted_with_terms(&self) -> bool {
+        false
+    }
+
+    fn num_terms(&self) -> u64 {
+        DynamicTermDictionary::num_terms(self)
+    }
+
+    fn sorted_ords_to_term_cb(
+        &self,
+        sorted_ords: &[u64],
+        callback: &mut dyn FnMut(&[u8]),
+    ) -> io::Result<bool> {
+        // Random access: the ords do not need to be sorted.
+        for &ord in sorted_ords {
+            let Some(term) = self.term(ord) else {
+                return Ok(false);
+            };
+            callback(term);
+        }
+        Ok(true)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_intern_assigns_ords_in_first_seen_order() {
+        let mut dictionary = DynamicTermDictionary::default();
+        assert_eq!(dictionary.num_terms(), 0);
+        assert_eq!(dictionary.intern(b"zebra"), 0);
+        assert_eq!(dictionary.intern(b"apple"), 1);
+        assert_eq!(dictionary.intern(b"zebra"), 0);
+        assert_eq!(dictionary.intern(b""), 2);
+        assert_eq!(dictionary.intern(b"apple"), 1);
+        assert_eq!(dictionary.num_terms(), 3);
+        assert_eq!(dictionary.term(0), Some(&b"zebra"[..]));
+        assert_eq!(dictionary.term(1), Some(&b"apple"[..]));
+        assert_eq!(dictionary.term(2), Some(&b""[..]));
+        assert_eq!(dictionary.term(3), None);
+        assert!(!dictionary.ords_sorted_with_terms());
+    }
+
+    #[test]
+    fn test_offsets_layout() {
+        let mut dictionary = DynamicTermDictionary::default();
+        dictionary.intern(b"ab");
+        dictionary.intern(b"");
+        dictionary.intern(b"cde");
+        assert_eq!(&dictionary.buffer, b"abcde");
+        assert_eq!(&dictionary.offsets, &[0, 2, 2, 5]);
+    }
+
+    #[test]
+    fn test_sorted_ords_to_term_cb() {
+        let mut dictionary = DynamicTermDictionary::default();
+        for term in ["b", "a", "c"] {
+            dictionary.intern(term.as_bytes());
+        }
+        let mut terms: Vec<String> = Vec::new();
+        let all_found = dictionary
+            .sorted_ords_to_term_cb(&[0, 2], &mut |term| {
+                terms.push(String::from_utf8(term.to_vec()).unwrap())
+            })
+            .unwrap();
+        assert!(all_found);
+        assert_eq!(terms, vec!["b".to_string(), "c".to_string()]);
+        let all_found = dictionary
+            .sorted_ords_to_term_cb(&[3], &mut |_| {})
+            .unwrap();
+        assert!(!all_found);
+    }
+
+    #[test]
+    fn test_many_terms() {
+        // Several table resizes.
+        let mut dictionary = DynamicTermDictionary::default();
+        for i in 0..10_000u64 {
+            assert_eq!(dictionary.intern(i.to_string().as_bytes()), i);
+        }
+        for i in 0..10_000u64 {
+            assert_eq!(dictionary.intern(i.to_string().as_bytes()), i);
+            assert_eq!(dictionary.term(i), Some(i.to_string().as_bytes()));
+        }
+        assert_eq!(dictionary.num_terms(), 10_000);
+    }
+
+    #[test]
+    fn test_colliding_hashes_get_distinct_ords() {
+        // Brute-force two distinct terms with the same 32-bit hash. With 300k terms, about 10
+        // colliding pairs are expected.
+        let num_candidates = 300_000;
+        let mut term_by_hash: std::collections::HashMap<u32, String> =
+            std::collections::HashMap::with_capacity(num_candidates);
+        let mut colliding_pair: Option<(String, String)> = None;
+        for i in 0..num_candidates {
+            let term = format!("t{i}");
+            if let Some(other_term) = term_by_hash.insert(term_hash(term.as_bytes()), term.clone())
+            {
+                colliding_pair = Some((other_term, term));
+                break;
+            }
+        }
+        let (first_term, second_term) = colliding_pair.expect("no 32-bit hash collision found");
+        assert_ne!(first_term, second_term);
+
+        let mut dictionary = DynamicTermDictionary::default();
+        assert_eq!(dictionary.intern(first_term.as_bytes()), 0);
+        assert_eq!(dictionary.intern(second_term.as_bytes()), 1);
+        assert_eq!(dictionary.intern(first_term.as_bytes()), 0);
+        assert_eq!(dictionary.intern(second_term.as_bytes()), 1);
+        assert_eq!(dictionary.term(0), Some(first_term.as_bytes()));
+        assert_eq!(dictionary.term(1), Some(second_term.as_bytes()));
+    }
+}

src/aggregation/value_source/jitexpr_value_source/mod.rs

@@ -0,0 +1,733 @@
+//! A [`ValueSource`] computing its values with a jitexpr expression.
+
+mod dynamic_term_dictionary;
+mod segment_source;
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use columnar::{Column, ColumnType, DynamicColumnHandle};
+use jitexpr::ast::{infer_types, infer_types_with_target, InferredTypeSet, TypeError, UntypedExpr};
+use jitexpr::compile::CompiledFn;
+use jitexpr::types::VarType;
+
+use self::segment_source::{ConstantValueSource, InputColumn, JitExprValueSource};
+use super::{ValueSource, ValueSourceProvider};
+use crate::query::doc_predicate_query::{find_input_column_handle, var_type_for_column_type};
+use crate::{SegmentReader, TantivyError};
+
+/// A [`ValueSourceProvider`] computing values with a jitexpr expression over the fast fields of a
+/// segment.
+///
+/// Requires the `jitexpr` feature. Variable names are resolved as fast-field names.
+///
+/// # Types
+///
+/// The expression is untyped. Its types are resolved for each segment, from the column types of
+/// the segment and the column types accepted by the aggregation. The result type can therefore
+/// differ from one segment to another: e.g. `(ADD a 1i64)` is `i64` on a segment where `a` is an
+/// `i64` column, and `f64` on a segment where `a` is an `f64` column.
+///
+/// - A variable is bound to the first column of its name with a type it accepts. A variable without
+///   such a column is null.
+/// - If the expression can never produce a type accepted by the aggregation, `for_segment` returns
+///   an error. jitexpr has no date type: date aggregations always fail.
+/// - If the expression does not produce an accepted type on a given segment, the source has no
+///   value on that segment.
+///
+/// # Values
+///
+/// Each input takes the first value of the doc. A null result means that the doc has no value. See
+/// `JitExprValueSource`.
+///
+/// Values use the monotonic `u64` mapping of fast fields. Text results are term ords of a
+/// dictionary built during the evaluation. These ords are not sorted with the terms, so the terms
+/// aggregation rejects the options relying on that (e.g. ordering by `_key`).
+///
+/// # Compilation
+///
+/// The expression is compiled for each segment, through the expression compilation cache of the
+/// index: segments with the same column types share the same compiled function.
+#[derive(Clone, Debug)]
+pub struct JitExprValueSourceProvider {
+    expression: UntypedExpr,
+}
+
+impl JitExprValueSourceProvider {
+    /// Creates a provider. Returns an error if the expression is not well typed.
+    pub fn new(expression: UntypedExpr) -> Result<JitExprValueSourceProvider, TypeError> {
+        infer_types(&expression)?;
+        Ok(JitExprValueSourceProvider { expression })
+    }
+
+    /// Returns the expression evaluated by this provider.
+    pub fn expression(&self) -> &UntypedExpr {
+        &self.expression
+    }
+}
+
+impl ValueSourceProvider for JitExprValueSourceProvider {
+    fn for_segment(
+        &self,
+        reader: &SegmentReader,
+        allowed_column_types: Option<&[ColumnType]>,
+    ) -> crate::Result<Box<dyn ValueSource>> {
+        // The checks below only depend on the expression and on the aggregation, so they fail on
+        // all of the segments alike.
+        let target_types: InferredTypeSet = target_type_set(allowed_column_types);
+        if target_types == InferredTypeSet::NONE {
+            return Err(TantivyError::InvalidArgument(format!(
+                "the calculated field `{}` cannot produce any of the column types accepted by the \
+                 aggregation: {allowed_column_types:?}",
+                self.expression
+            )));
+        }
+        let inferred_types: HashMap<&str, InferredTypeSet> =
+            infer_types_with_target(&self.expression, target_types).map_err(|type_error| {
+                TantivyError::InvalidArgument(format!(
+                    "the calculated field `{}` cannot produce any of the types accepted by the \
+                     aggregation {target_types}: {type_error}",
+                    self.expression
+                ))
+            })?;
+
+        let mut variable_types: HashMap<&str, VarType> =
+            HashMap::with_capacity(inferred_types.len());
+        let mut column_handles: HashMap<&str, DynamicColumnHandle> =
+            HashMap::with_capacity(inferred_types.len());
+        for (variable_name, accepted_types) in inferred_types {
+            let Some(handle) = find_input_column_handle(reader, variable_name, accepted_types)
+            else {
+                // The compiler treats unbound variables as null.
+                continue;
+            };
+            let Some(var_type) = var_type_for_column_type(handle.column_type()) else {
+                continue;
+            };
+            variable_types.insert(variable_name, var_type);
+            column_handles.insert(variable_name, handle);
+        }
+
+        let compiled_fn: Arc<CompiledFn> = reader
+            .index()
+            .expr_compilation_cache()
+            .compile(&self.expression, &variable_types)
+            .map_err(|compile_error| {
+                TantivyError::InvalidArgument(format!(
+                    "the compilation of the calculated field `{}` failed: {compile_error}",
+                    self.expression
+                ))
+            })?;
+
+        let Some(column_type) = column_type_for_result_type(compiled_fn.result_type()) else {
+            // The expression is null for all of the docs of the segment.
+            return Ok(empty_source(reader, ColumnType::U64, allowed_column_types));
+        };
+        if let Some(allowed_column_types) = allowed_column_types {
+            if !allowed_column_types.contains(&column_type) {
+                return Ok(empty_source(
+                    reader,
+                    column_type,
+                    Some(allowed_column_types),
+                ));
+            }
+        }
+
+        if compiled_fn.inputs().is_empty() {
+            // The expression does not depend on the doc: we evaluate it once.
+            let Some(constant_source) =
+                ConstantValueSource::evaluate(compiled_fn.context(), column_type)
+            else {
+                return Ok(empty_source(reader, column_type, allowed_column_types));
+            };
+            return Ok(Box::new(constant_source));
+        }
+
+        // The compiler defines the order of the inputs. Do not rely on the inference or on the
+        // HashMap iteration order.
+        let mut inputs: Vec<InputColumn> = Vec::with_capacity(compiled_fn.inputs().len());
+        for compiled_input in compiled_fn.inputs() {
+            let Some(handle) = column_handles.remove(compiled_input.variable_name.as_ref()) else {
+                return Err(TantivyError::InternalError(format!(
+                    "the compiled input `{}` is not bound to a column",
+                    compiled_input.variable_name
+                )));
+            };
+            inputs.push(InputColumn::open(&handle, compiled_input.r#type)?);
+        }
+        Ok(Box::new(JitExprValueSource::new(
+            compiled_fn.context(),
+            column_type,
+            inputs,
+        )))
+    }
+}
+
+/// Returns the result types matching the column types accepted by the aggregation.
+fn target_type_set(allowed_column_types: Option<&[ColumnType]>) -> InferredTypeSet {
+    let Some(allowed_column_types) = allowed_column_types else {
+        return InferredTypeSet::ALL;
+    };
+    let mut target_types = InferredTypeSet::NONE;
+    for column_type in allowed_column_types {
+        match column_type {
+            ColumnType::Bool => target_types.boolean = true,
+            ColumnType::I64 => target_types.i64 = true,
+            ColumnType::U64 => target_types.u64 = true,
+            ColumnType::F64 => target_types.f64 = true,
+            ColumnType::Str => target_types.string = true,
+            // jitexpr has no such types.
+            ColumnType::DateTime | ColumnType::IpAddr | ColumnType::Bytes => {}
+        }
+    }
+    target_types
+}
+
+/// Returns `None` for an expression that is always null.
+fn column_type_for_result_type(result_type: VarType) -> Option<ColumnType> {
+    match result_type {
+        VarType::Bool => Some(ColumnType::Bool),
+        VarType::I64 => Some(ColumnType::I64),
+        VarType::U64 => Some(ColumnType::U64),
+        VarType::F64 => Some(ColumnType::F64),
+        VarType::Str => Some(ColumnType::Str),
+        VarType::None => None,
+    }
+}
+
+/// Returns a source without any value.
+///
+/// The source gets `column_type` if the aggregation accepts it, and an accepted type otherwise:
+/// with a type the aggregation does not accept, the aggregation would use the fast field with the
+/// same name instead.
+fn empty_source(
+    reader: &SegmentReader,
+    column_type: ColumnType,
+    allowed_column_types: Option<&[ColumnType]>,
+) -> Box<dyn ValueSource> {
+    let column_type = match allowed_column_types {
+        Some(allowed_column_types) if !allowed_column_types.contains(&column_type) => {
+            allowed_column_types.first().copied().unwrap_or(column_type)
+        }
+        _ => column_type,
+    };
+    Box::new((
+        Column::<u64>::build_empty_column(reader.max_doc()),
+        column_type,
+    ))
+}
+
+#[cfg(test)]
+mod tests {
+    use serde_json::{json, Value};
+
+    use super::*;
+    use crate::aggregation::agg_req::Aggregations;
+    use crate::aggregation::{AggContextParams, AggregationCollector, ValueSourceRegistry};
+    use crate::query::AllQuery;
+    use crate::schema::{Schema, FAST, STRING};
+    use crate::{Index, TantivyDocument};
+
+    fn registry(calculated_fields: &[(&str, &str)]) -> ValueSourceRegistry {
+        let mut registry = ValueSourceRegistry::default();
+        for (name, expression) in calculated_fields {
+            let expression = jitexpr::ast::deserialize(expression).unwrap();
+            let provider = JitExprValueSourceProvider::new(expression).unwrap();
+            registry.register(name, Arc::new(provider));
+        }
+        registry
+    }
+
+    fn try_run_agg(
+        index: &Index,
+        calculated_fields: &[(&str, &str)],
+        aggs: Value,
+    ) -> crate::Result<Value> {
+        let context =
+            AggContextParams::default().with_value_sources(Arc::new(registry(calculated_fields)));
+        let aggs: Aggregations = serde_json::from_value(aggs).unwrap();
+        let collector = AggregationCollector::from_aggs(aggs, context);
+        let searcher = index.reader().unwrap().searcher();
+        let result = searcher.search(&AllQuery, &collector)?;
+        Ok(serde_json::to_value(result).unwrap())
+    }
+
+    fn run_agg(index: &Index, calculated_fields: &[(&str, &str)], aggs: Value) -> Value {
+        try_run_agg(index, calculated_fields, aggs).unwrap()
+    }
+
+    /// Returns the `(key, doc_count)` of the buckets, sorted by key.
+    fn sorted_bucket_counts(result: &Value, agg_name: &str) -> Vec<(String, u64)> {
+        let mut bucket_counts: Vec<(String, u64)> = result[agg_name]["buckets"]
+            .as_array()
+            .unwrap()
+            .iter()
+            .map(|bucket| {
+                (
+                    bucket["key"].to_string(),
+                    bucket["doc_count"].as_u64().unwrap(),
+                )
+            })
+            .collect();
+        bucket_counts.sort();
+        bucket_counts
+    }
+
+    /// Two segments. The last doc of each segment has no value.
+    fn create_index() -> Index {
+        let mut schema_builder = Schema::builder();
+        let number = schema_builder.add_u64_field("number", FAST);
+        let price = schema_builder.add_f64_field("price", FAST);
+        let delta = schema_builder.add_i64_field("delta", FAST);
+        let flag = schema_builder.add_bool_field("flag", FAST);
+        let label = schema_builder.add_text_field("label", STRING | FAST);
+        let index = Index::create_in_ram(schema_builder.build());
+        let mut writer = index.writer_for_tests().unwrap();
+        writer
+            .add_document(
+                doc!(number => 1u64, price => 1.5f64, delta => -3i64, flag => true, label => "a"),
+            )
+            .unwrap();
+        writer
+            .add_document(
+                doc!(number => 2u64, price => -2.5f64, delta => 4i64, flag => false, label => "b"),
+            )
+            .unwrap();
+        writer.add_document(doc!()).unwrap();
+        writer.commit().unwrap();
+        writer
+            .add_document(
+                doc!(number => 3u64, price => 10.0f64, delta => -1i64, flag => true, label => "a"),
+            )
+            .unwrap();
+        writer.add_document(doc!()).unwrap();
+        writer.commit().unwrap();
+        index
+    }
+
+    #[test]
+    fn test_metrics_over_calculated_field() {
+        let index = create_index();
+        let result = run_agg(
+            &index,
+            &[("doubled", "(MULTIPLY number 2u64)")],
+            json!({
+                "s": { "stats": { "field": "doubled" } },
+                "with_missing": { "stats": { "field": "doubled", "missing": 100.0 } }
+            }),
+        );
+        // The docs without `number` evaluate to null, and have no value.
+        assert_eq!(result["s"]["count"], 3);
+        assert_eq!(result["s"]["sum"], 12.0);
+        assert_eq!(result["s"]["min"], 2.0);
+        assert_eq!(result["s"]["max"], 6.0);
+        assert_eq!(result["with_missing"]["count"], 5);
+        assert_eq!(result["with_missing"]["sum"], 212.0);
+    }
+
+    #[test]
+    fn test_histogram_and_range_over_calculated_field() {
+        let index = create_index();
+        let result = run_agg(
+            &index,
+            &[("shifted", "(ADD price 1.0f64)")],
+            json!({
+                "h": { "histogram": { "field": "shifted", "interval": 5.0 } },
+                "r": { "range": { "field": "shifted", "ranges": [{ "to": 0.0 }, { "from": 0.0 }] } }
+            }),
+        );
+        let histogram_counts: Vec<(f64, u64)> = result["h"]["buckets"]
+            .as_array()
+            .unwrap()
+            .iter()
+            .map(|bucket| {
+                (
+                    bucket["key"].as_f64().unwrap(),
+                    bucket["doc_count"].as_u64().unwrap(),
+                )
+            })
+            .collect();
+        assert_eq!(
+            histogram_counts,
+            vec![(-5.0, 1), (0.0, 1), (5.0, 0), (10.0, 1)]
+        );
+        let range_buckets = result["r"]["buckets"].as_array().unwrap();
+        assert_eq!(range_buckets[0]["doc_count"], 1);
+        assert_eq!(range_buckets[1]["doc_count"], 2);
+    }
+
+    /// An aggregation over the identity calculated field must match the aggregation over the
+    /// field itself. This checks the monotonic mapping of each type.
+    #[test]
+    fn test_identity_matches_fast_field_for_each_type() {
+        let index = create_index();
+        for (field, agg_kinds) in [
+            ("number", &["terms", "stats"][..]),
+            ("price", &["terms", "stats"][..]),
+            ("delta", &["terms", "stats"][..]),
+            ("flag", &["terms"][..]),
+            ("label", &["terms", "cardinality"][..]),
+        ] {
+            for agg_kind in agg_kinds {
+                let agg_on =
+                    |field_name: &str| json!({ "agg": { *agg_kind: { "field": field_name } } });
+                let expected = run_agg(&index, &[], agg_on(field));
+                let actual = run_agg(&index, &[("calculated", field)], agg_on("calculated"));
+                assert_eq!(actual, expected, "{agg_kind} over {field}");
+            }
+        }
+    }
+
+    #[test]
+    fn test_registered_name_shadows_fast_field() {
+        let index = create_index();
+        // The variable `number` is the fast field, the aggregated `number` is the calculated
+        // field.
+        let result = run_agg(
+            &index,
+            &[("number", "(MULTIPLY number 10u64)")],
+            json!({ "s": { "sum": { "field": "number" } } }),
+        );
+        assert_eq!(result["s"]["value"], 60.0);
+    }
+
+    #[test]
+    fn test_constant_calculated_field() {
+        let index = create_index();
+        let result = run_agg(
+            &index,
+            &[("three", "(ADD 1u64 2u64)"), ("constant_text", r#""x""#)],
+            json!({
+                "s": { "stats": { "field": "three" } },
+                "t": { "terms": { "field": "constant_text" } }
+            }),
+        );
+        assert_eq!(result["s"]["count"], 5);
+        assert_eq!(result["s"]["sum"], 15.0);
+        assert_eq!(
+            sorted_bucket_counts(&result, "t"),
+            vec![(r#""x""#.to_string(), 5)]
+        );
+    }
+
+    #[test]
+    fn test_unknown_variable_has_no_value() {
+        let index = create_index();
+        let result = run_agg(
+            &index,
+            &[("unknown", "(ADD absent_field 1u64)")],
+            json!({ "s": { "stats": { "field": "unknown" } } }),
+        );
+        assert_eq!(result["s"]["count"], 0);
+    }
+
+    #[test]
+    fn test_incompatible_calculated_field_is_an_error() {
+        let index = create_index();
+        let err = try_run_agg(
+            &index,
+            &[("upper", "(UPPER label)")],
+            json!({ "s": { "stats": { "field": "upper" } } }),
+        )
+        .unwrap_err();
+        assert!(err.to_string().contains("cannot produce"), "{err}");
+        // jitexpr has no date type.
+        let err = try_run_agg(
+            &index,
+            &[("doubled", "(MULTIPLY number 2u64)")],
+            json!({ "h": { "date_histogram": { "field": "doubled", "fixed_interval": "1d" } } }),
+        )
+        .unwrap_err();
+        assert!(err.to_string().contains("cannot produce"), "{err}");
+    }
+
+    #[test]
+    fn test_multivalued_inputs_use_first_value() {
+        let mut schema_builder = Schema::builder();
+        let number = schema_builder.add_u64_field("number", FAST);
+        let other = schema_builder.add_u64_field("other", FAST);
+        let index = Index::create_in_ram(schema_builder.build());
+        let mut writer = index.writer_for_tests().unwrap();
+        writer
+            .add_document(doc!(number => 1u64, number => 2u64, other => 10u64, other => 20u64))
+            .unwrap();
+        writer
+            .add_document(doc!(number => 3u64, other => 30u64))
+            .unwrap();
+        writer.add_document(doc!(other => 40u64)).unwrap();
+        writer.commit().unwrap();
+
+        // Each input takes the first value of the doc. The doc without `number` is null.
+        let result = run_agg(
+            &index,
+            &[("sum", "(ADD number other)")],
+            json!({
+                "s": { "stats": { "field": "sum" } },
+                "t": { "terms": { "field": "sum" } }
+            }),
+        );
+        assert_eq!(result["s"]["count"], 2);
+        assert_eq!(result["s"]["sum"], 11.0 + 33.0);
+        assert_eq!(
+            sorted_bucket_counts(&result, "t"),
+            vec![("11".to_string(), 1), ("33".to_string(), 1)]
+        );
+
+        // A doc without a value for an input is evaluated with null.
+        let result = run_agg(
+            &index,
+            &[("is_null", "(IS_NULL number)")],
+            json!({ "t": { "terms": { "field": "is_null" } } }),
+        );
+        // Bool keys are serialized as 0 / 1, like for a bool fast field.
+        assert_eq!(
+            sorted_bucket_counts(&result, "t"),
+            vec![("0".to_string(), 2), ("1".to_string(), 1)]
+        );
+    }
+
+    /// Two segments, in which the labels are first seen in a different order.
+    fn create_label_index() -> Index {
+        let mut schema_builder = Schema::builder();
+        let label = schema_builder.add_text_field("label", STRING | FAST);
+        let score = schema_builder.add_u64_field("score", FAST);
+        let index = Index::create_in_ram(schema_builder.build());
+        let mut writer = index.writer_for_tests().unwrap();
+        for (label_value, score_value) in [("b", 1u64), ("a", 2), ("b", 3), ("c", 4)] {
+            writer
+                .add_document(doc!(label => label_value, score => score_value))
+                .unwrap();
+        }
+        writer.add_document(doc!(score => 100u64)).unwrap();
+        writer.commit().unwrap();
+        for (label_value, score_value) in [("c", 5u64), ("a", 6), ("b", 7)] {
+            writer
+                .add_document(doc!(label => label_value, score => score_value))
+                .unwrap();
+        }
+        // A multivalued doc: only its first value, `a`, is used.
+        writer
+            .add_document(doc!(label => "a", label => "d", score => 8u64))
+            .unwrap();
+        writer.commit().unwrap();
+        index
+    }
+
+    #[test]
+    fn test_terms_over_calculated_text_field() {
+        let index = create_label_index();
+        let result = run_agg(
+            &index,
+            &[("upper", "(UPPER label)")],
+            json!({
+                "t": {
+                    "terms": { "field": "upper" },
+                    "aggs": { "score_sum": { "sum": { "field": "score" } } }
+                }
+            }),
+        );
+        let keys_counts_sums: Vec<(&str, u64, f64)> = result["t"]["buckets"]
+            .as_array()
+            .unwrap()
+            .iter()
+            .map(|bucket| {
+                (
+                    bucket["key"].as_str().unwrap(),
+                    bucket["doc_count"].as_u64().unwrap(),
+                    bucket["score_sum"]["value"].as_f64().unwrap(),
+                )
+            })
+            .collect();
+        assert_eq!(
+            keys_counts_sums,
+            vec![("A", 3, 16.0), ("B", 3, 11.0), ("C", 2, 9.0)]
+        );
+    }
+
+    #[test]
+    fn test_terms_ordered_by_sub_aggregation_over_calculated_text_field() {
+        let index = create_label_index();
+        let result = run_agg(
+            &index,
+            &[("upper", "(UPPER label)")],
+            json!({
+                "t": {
+                    "terms": { "field": "upper", "order": { "score_min": "asc" } },
+                    "aggs": { "score_min": { "min": { "field": "score" } } }
+                }
+            }),
+        );
+        let keys: Vec<&str> = result["t"]["buckets"]
+            .as_array()
+            .unwrap()
+            .iter()
+            .map(|bucket| bucket["key"].as_str().unwrap())
+            .collect();
+        assert_eq!(keys, vec!["B", "A", "C"]);
+    }
+
+    #[test]
+    fn test_cardinality_over_calculated_text_field() {
+        let index = create_label_index();
+        let result = run_agg(
+            &index,
+            &[("first_letter", "(LEFT (UPPER label) 1u64)")],
+            json!({ "c": { "cardinality": { "field": "first_letter" } } }),
+        );
+        assert_eq!(result["c"]["value"], 3.0);
+    }
+
+    #[test]
+    fn test_unsupported_terms_options_on_calculated_text_field() {
+        let index = create_label_index();
+        let calculated_fields = [("upper", "(UPPER label)")];
+        for (terms_req, expected_message) in [
+            (
+                json!({ "field": "upper", "order": { "_key": "asc" } }),
+                "ordered by `_key`",
+            ),
+            (
+                json!({ "field": "upper", "min_doc_count": 0 }),
+                "`min_doc_count: 0`",
+            ),
+            (
+                json!({ "field": "upper", "include": "A.*" }),
+                "`include` / `exclude`",
+            ),
+            (
+                json!({ "field": "upper", "exclude": ["A"] }),
+                "`include` / `exclude`",
+            ),
+            (json!({ "field": "upper", "missing": "NONE" }), "`missing`"),
+        ] {
+            let err = try_run_agg(
+                &index,
+                &calculated_fields,
+                json!({ "t": { "terms": terms_req } }),
+            )
+            .unwrap_err();
+            assert!(
+                err.to_string().contains(expected_message),
+                "{terms_req}: {err}"
+            );
+        }
+    }
+
+    #[test]
+    fn test_result_type_is_resolved_per_segment() {
+        let mut schema_builder = Schema::builder();
+        let json_field = schema_builder.add_json_field("j", FAST);
+        let index = Index::create_in_ram(schema_builder.build());
+        let mut writer = index.writer_for_tests().unwrap();
+        writer
+            .add_document(doc!(json_field => json!({ "x": 2 })))
+            .unwrap();
+        writer.commit().unwrap();
+        writer
+            .add_document(doc!(json_field => json!({ "x": 1.5 })))
+            .unwrap();
+        writer.commit().unwrap();
+
+        let expression = jitexpr::ast::deserialize("(ADD j.x 1i64)").unwrap();
+        let provider = JitExprValueSourceProvider::new(expression).unwrap();
+        let searcher = index.reader().unwrap().searcher();
+        let mut column_types: Vec<ColumnType> = Vec::new();
+        for segment_reader in searcher.segment_readers() {
+            let source = provider.for_segment(segment_reader, None).unwrap();
+            column_types.push(source.column_type());
+        }
+        column_types.sort_by_key(|column_type| column_type.to_code());
+        assert_eq!(column_types, vec![ColumnType::I64, ColumnType::F64]);
+
+        let result = run_agg(
+            &index,
+            &[("x_plus_one", "(ADD j.x 1i64)")],
+            json!({ "s": { "stats": { "field": "x_plus_one" } } }),
+        );
+        assert_eq!(result["s"]["count"], 2);
+        assert_eq!(result["s"]["sum"], 5.5);
+    }
+
+    #[test]
+    fn test_compilation_uses_the_index_cache() {
+        let index = create_index();
+        let cache = index.expr_compilation_cache();
+        assert!(cache.is_empty());
+        run_agg(
+            &index,
+            &[("doubled", "(MULTIPLY number 2u64)")],
+            json!({ "s": { "stats": { "field": "doubled" } } }),
+        );
+        // Both segments have the same column types, and share the same compiled function.
+        assert_eq!(cache.len(), 1);
+        run_agg(
+            &index,
+            &[("doubled", "(MULTIPLY number 2u64)")],
+            json!({ "s": { "sum": { "field": "doubled" } } }),
+        );
+        assert_eq!(cache.len(), 1);
+    }
+
+    /// Calculated fields over many docs (several blocks), with sparse inputs and deletes, must
+    /// match the same values materialized in fast fields.
+    #[test]
+    fn test_matches_materialized_fields() {
+        let mut schema_builder = Schema::builder();
+        let id = schema_builder.add_u64_field("id", FAST | crate::schema::INDEXED);
+        let number = schema_builder.add_u64_field("number", FAST);
+        let label = schema_builder.add_text_field("label", STRING | FAST);
+        let tripled = schema_builder.add_u64_field("tripled", FAST);
+        let upper = schema_builder.add_text_field("upper", STRING | FAST);
+        let index = Index::create_in_ram(schema_builder.build());
+        let mut writer = index.writer_for_tests().unwrap();
+        for segment_ord in 0..2u64 {
+            for i in 0..300u64 {
+                let mut doc = TantivyDocument::default();
+                doc.add_u64(id, segment_ord * 1000 + i);
+                if i % 3 != 0 {
+                    doc.add_u64(number, i % 17);
+                    doc.add_u64(tripled, (i % 17) * 3);
+                }
+                if i % 2 == 0 {
+                    let label_value = ["a", "bb", "ccc", "dd"][((i + segment_ord) % 4) as usize];
+                    doc.add_text(label, label_value);
+                    doc.add_text(upper, label_value.to_uppercase());
+                }
+                writer.add_document(doc).unwrap();
+            }
+            writer.commit().unwrap();
+        }
+        writer.delete_term(crate::Term::from_field_u64(id, 4));
+        writer.delete_term(crate::Term::from_field_u64(id, 1002));
+        writer.commit().unwrap();
+
+        let calculated_fields = [
+            ("tripled_calc", "(MULTIPLY number 3u64)"),
+            ("upper_calc", "(UPPER label)"),
+        ];
+        let aggs_on = |tripled_field: &str, upper_field: &str| {
+            json!({
+                "tripled_stats": { "stats": { "field": tripled_field } },
+                "tripled_terms": { "terms": { "field": tripled_field, "size": 100 } },
+                "upper_terms": { "terms": { "field": upper_field, "size": 100 } },
+                "upper_cardinality": { "cardinality": { "field": upper_field } }
+            })
+        };
+        let expected = run_agg(&index, &[], aggs_on("tripled", "upper"));
+        let actual = run_agg(
+            &index,
+            &calculated_fields,
+            aggs_on("tripled_calc", "upper_calc"),
+        );
+        assert_eq!(actual["tripled_stats"], expected["tripled_stats"]);
+        assert_eq!(actual["upper_cardinality"], expected["upper_cardinality"]);
+        for agg_name in ["tripled_terms", "upper_terms"] {
+            assert_eq!(
+                sorted_bucket_counts(&actual, agg_name),
+                sorted_bucket_counts(&expected, agg_name),
+                "{agg_name}"
+            );
+        }
+        assert_eq!(expected["tripled_stats"]["count"], 398);
+    }
+}

src/aggregation/value_source/jitexpr_value_source/segment_source.rs

@@ -0,0 +1,462 @@
+//! The value sources of a jitexpr expression on a segment.
+
+use columnar::{
+    Cardinality, Column, ColumnType, DynamicColumn, DynamicColumnHandle,
+    MonotonicallyMappableToU64, RowId, StrColumn,
+};
+use jitexpr::compile::CompiledFnCtx;
+use jitexpr::types::{VarType, VariableValue};
+
+use super::dynamic_term_dictionary::DynamicTermDictionary;
+use crate::aggregation::value_source::{ValueSource, ValueSourceDictionary};
+use crate::query::doc_predicate_query::var_type_for_column_type;
+use crate::{DocId, TantivyError, COLLECT_BLOCK_BUFFER_LEN};
+
+/// The value source of a jitexpr expression on a segment.
+///
+/// Each input takes the first value of the doc. A doc without a value for an input gets null for
+/// that input. The expression is evaluated once per doc, and a null result means that the doc has
+/// no value. A doc therefore has at most one value.
+///
+/// Why the first value only: using all of the values of multivalued inputs would require
+/// evaluating the cartesian product of the values of the inputs, which explodes with several
+/// multivalued inputs.
+///
+/// Values use the monotonic `u64` mapping of fast fields. `Str` results are interned in a
+/// [`DynamicTermDictionary`] owned by the source: their values are term ords in first-seen order.
+pub(super) struct JitExprValueSource {
+    compiled_fn: CompiledFnCtx,
+    column_type: ColumnType,
+    /// Invariant: `inputs[i]` is bound to `compiled_fn.inputs()[i]`, and has the same type.
+    inputs: Vec<InputColumn>,
+    /// Interns the `Str` results.
+    term_dictionary: DynamicTermDictionary,
+}
+
+impl JitExprValueSource {
+    /// Precondition: `column_type` is the column type of the result type of `compiled_fn`.
+    ///
+    /// Panics if `inputs` do not match the inputs of `compiled_fn`.
+    pub(super) fn new(
+        compiled_fn: CompiledFnCtx,
+        column_type: ColumnType,
+        inputs: Vec<InputColumn>,
+    ) -> JitExprValueSource {
+        // Evaluating the function with inputs of the wrong type is undefined behavior.
+        assert_eq!(compiled_fn.inputs().len(), inputs.len());
+        for (compiled_input, input) in compiled_fn.inputs().iter().zip(&inputs) {
+            assert_eq!(compiled_input.r#type, input.var_type);
+        }
+        JitExprValueSource {
+            compiled_fn,
+            column_type,
+            inputs,
+            term_dictionary: DynamicTermDictionary::default(),
+        }
+    }
+}
+
+impl std::fmt::Debug for JitExprValueSource {
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        formatter
+            .debug_struct("JitExprValueSource")
+            .field("column_type", &self.column_type)
+            .field("num_inputs", &self.inputs.len())
+            .finish()
+    }
+}
+
+impl ValueSource for JitExprValueSource {
+    fn column_type(&self) -> ColumnType {
+        self.column_type
+    }
+
+    fn load_block(
+        &mut self,
+        docs: &[DocId],
+        values: &mut Vec<u64>,
+        docids: &mut Vec<DocId>,
+        _row_ids: &mut Vec<RowId>,
+    ) -> Cardinality {
+        values.clear();
+        docids.clear();
+        for input in &mut self.inputs {
+            input.load_block(docs);
+        }
+        let JitExprValueSource {
+            compiled_fn,
+            inputs,
+            term_dictionary,
+            ..
+        } = self;
+        let result_type: VarType = compiled_fn.result_type();
+        let mut args: Vec<VariableValue> = Vec::with_capacity(inputs.len());
+        for (doc_pos, &doc) in docs.iter().enumerate() {
+            args.clear();
+            for input in inputs.iter() {
+                args.push(input.value(doc_pos));
+            }
+            // SAFETY: `args` follow the compiled inputs, with their types (see `new`). The strings
+            // are borrowed from `inputs`, which are not modified during the call.
+            let result: VariableValue = unsafe { compiled_fn.call(&args) };
+            // SAFETY: `result` has the result type of the compiled function.
+            let Some(value) = (unsafe { encode_result(result, result_type, term_dictionary) })
+            else {
+                continue;
+            };
+            values.push(value);
+            docids.push(doc);
+        }
+        if values.len() == docs.len() {
+            // Each doc has exactly one value, in the order of `docs`.
+            return Cardinality::Full;
+        }
+        Cardinality::Optional
+    }
+
+    fn term_dictionary(&self) -> Option<&dyn ValueSourceDictionary> {
+        if self.column_type != ColumnType::Str {
+            return None;
+        }
+        Some(&self.term_dictionary)
+    }
+}
+
+/// A source with the same value for all of the docs.
+#[derive(Debug)]
+pub(super) struct ConstantValueSource {
+    column_type: ColumnType,
+    /// In the monotonic `u64` mapping. For `Str`, the term ord of the only term of
+    /// `term_dictionary`.
+    value: u64,
+    term_dictionary: DynamicTermDictionary,
+}
+
+impl ConstantValueSource {
+    /// Evaluates a compiled function without any input. Returns `None` if the result is null.
+    ///
+    /// Precondition: `column_type` is the column type of the result type of `compiled_fn`.
+    pub(super) fn evaluate(
+        mut compiled_fn: CompiledFnCtx,
+        column_type: ColumnType,
+    ) -> Option<ConstantValueSource> {
+        assert!(compiled_fn.inputs().is_empty());
+        let result_type: VarType = compiled_fn.result_type();
+        let mut term_dictionary = DynamicTermDictionary::default();
+        // SAFETY: the function has no input.
+        let result: VariableValue = unsafe { compiled_fn.call(&[]) };
+        // SAFETY: `result` has the result type of the compiled function.
+        let value = unsafe { encode_result(result, result_type, &mut term_dictionary) }?;
+        Some(ConstantValueSource {
+            column_type,
+            value,
+            term_dictionary,
+        })
+    }
+}
+
+impl ValueSource for ConstantValueSource {
+    fn column_type(&self) -> ColumnType {
+        self.column_type
+    }
+
+    fn load_block(
+        &mut self,
+        docs: &[DocId],
+        values: &mut Vec<u64>,
+        _docids: &mut Vec<DocId>,
+        _row_ids: &mut Vec<RowId>,
+    ) -> Cardinality {
+        values.clear();
+        values.resize(docs.len(), self.value);
+        Cardinality::Full
+    }
+
+    fn term_dictionary(&self) -> Option<&dyn ValueSourceDictionary> {
+        if self.column_type != ColumnType::Str {
+            return None;
+        }
+        Some(&self.term_dictionary)
+    }
+}
+
+/// Encodes a result in the monotonic `u64` mapping. Returns `None` for a null result.
+///
+/// `Str` results are interned in `term_dictionary`, and encoded as their term ord.
+///
+/// # Safety
+///
+/// `result` must have the type `result_type`.
+unsafe fn encode_result(
+    result: VariableValue,
+    result_type: VarType,
+    term_dictionary: &mut DynamicTermDictionary,
+) -> Option<u64> {
+    // SAFETY: guaranteed by the caller.
+    unsafe {
+        match result_type {
+            VarType::Bool => result.as_bool().map(bool::to_u64),
+            VarType::I64 => result.as_i64().map(i64::to_u64),
+            VarType::U64 => result.as_u64(),
+            VarType::F64 => result.as_f64().map(f64::to_u64),
+            VarType::Str => result
+                .as_str()
+                .map(|term: &str| term_dictionary.intern(term.as_bytes())),
+            VarType::None => None,
+        }
+    }
+}
+
+/// A fast-field column bound to an input of the compiled expression.
+pub(super) struct InputColumn {
+    var_type: VarType,
+    /// Values in their monotonic `u64` mapping. Term ords for a `Str` input.
+    column: Column<u64>,
+    /// The dictionary of the term ords. `Some` if and only if `var_type` is `Str`.
+    str_column: Option<StrColumn>,
+    /// `first_values[i]` is the first value of `docs[i]`, for the docs of the last block.
+    first_values: Vec<Option<u64>>,
+    /// The terms of `first_values`, for a `Str` input.
+    block_terms: BlockTerms,
+}
+
+impl InputColumn {
+    /// Opens the column of `handle`, for an input of type `var_type`.
+    pub(super) fn open(
+        handle: &DynamicColumnHandle,
+        var_type: VarType,
+    ) -> crate::Result<InputColumn> {
+        let column_type = handle.column_type();
+        if var_type_for_column_type(column_type) != Some(var_type) {
+            return Err(TantivyError::InternalError(format!(
+                "an input of type {var_type:?} cannot be bound to a column of type {column_type}"
+            )));
+        }
+        let (column, str_column) = if var_type == VarType::Str {
+            let DynamicColumn::Str(str_column) = handle.open()? else {
+                return Err(TantivyError::InternalError(
+                    "a text column could not be opened as a text column".to_string(),
+                ));
+            };
+            (str_column.ords().clone(), Some(str_column))
+        } else {
+            let Some(column) = handle.open_u64_lenient()? else {
+                return Err(TantivyError::InternalError(format!(
+                    "a column of type {column_type} could not be opened as u64"
+                )));
+            };
+            (column, None)
+        };
+        Ok(InputColumn {
+            var_type,
+            column,
+            str_column,
+            first_values: Vec::with_capacity(COLLECT_BLOCK_BUFFER_LEN),
+            block_terms: BlockTerms::default(),
+        })
+    }
+
+    /// Loads the first value of each doc of `docs`.
+    ///
+    /// Panics if the dictionary of a text column is corrupted.
+    fn load_block(&mut self, docs: &[DocId]) {
+        // `first_vals` does not write the entries of the docs without a value: we reset them.
+        self.first_values.clear();
+        self.first_values.resize(docs.len(), None);
+        self.column.first_vals(docs, &mut self.first_values);
+        if let Some(str_column) = &self.str_column {
+            self.block_terms.load(str_column, &self.first_values);
+        }
+    }
+
+    /// Returns the value of `docs[doc_pos]`, for the `docs` of the last `load_block` call.
+    #[inline]
+    fn value(&self, doc_pos: usize) -> VariableValue<'_> {
+        let Some(value) = self.first_values[doc_pos] else {
+            return VariableValue::none();
+        };
+        match self.var_type {
+            VarType::Bool => VariableValue::from(bool::from_u64(value)),
+            VarType::I64 => VariableValue::from(i64::from_u64(value)),
+            VarType::U64 => VariableValue::from(value),
+            VarType::F64 => VariableValue::from(f64::from_u64(value)),
+            VarType::Str => VariableValue::from(self.block_terms.term(value)),
+            VarType::None => VariableValue::none(),
+        }
+    }
+}
+
+/// The terms of the term ords of a block, fetched with a single dictionary lookup.
+#[derive(Default)]
+struct BlockTerms {
+    /// The distinct term ords of the block, sorted.
+    sorted_ords: Vec<u64>,
+    /// The concatenated terms, in the order of `sorted_ords`.
+    buffer: String,
+    /// `buffer[offsets[i]..offsets[i + 1]]` is the term of `sorted_ords[i]`.
+    offsets: Vec<usize>,
+}
+
+impl BlockTerms {
+    /// Panics if the dictionary is corrupted: `load_block` cannot return an error.
+    fn load(&mut self, str_column: &StrColumn, term_ords: &[Option<u64>]) {
+        let BlockTerms {
+            sorted_ords,
+            buffer,
+            offsets,
+        } = self;
+        sorted_ords.clear();
+        sorted_ords.extend(term_ords.iter().flatten());
+        sorted_ords.sort_unstable();
+        sorted_ords.dedup();
+        buffer.clear();
+        offsets.clear();
+        offsets.push(0);
+        let all_found = str_column
+            .dictionary()
+            .sorted_ords_to_term_cb(sorted_ords, |term| {
+                let term_str =
+                    std::str::from_utf8(term).expect("fast-field string dictionary is corrupted");
+                buffer.push_str(term_str);
+                offsets.push(buffer.len());
+            })
+            .expect("fast-field string dictionary is corrupted");
+        assert!(all_found, "fast-field string dictionary is corrupted");
+    }
+
+    /// Precondition: `term_ord` is one of the term ords of the last `load` call.
+    #[inline]
+    fn term(&self, term_ord: u64) -> &str {
+        let idx = self
+            .sorted_ords
+            .binary_search(&term_ord)
+            .expect("the term ord belongs to the block");
+        &self.buffer[self.offsets[idx]..self.offsets[idx + 1]]
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::aggregation::{JitExprValueSourceProvider, ValueSourceProvider};
+    use crate::schema::{Schema, FAST, STRING};
+    use crate::Index;
+
+    /// One segment: doc `i` has `number = i` and `label = labels[i]`, except for the docs of
+    /// `null_docs`, which have neither.
+    fn create_index(labels: &[&str], null_docs: &[usize]) -> Index {
+        let mut schema_builder = Schema::builder();
+        let number = schema_builder.add_u64_field("number", FAST);
+        let label = schema_builder.add_text_field("label", STRING | FAST);
+        let index = Index::create_in_ram(schema_builder.build());
+        let mut writer = index.writer_for_tests().unwrap();
+        for (doc_pos, label_value) in labels.iter().enumerate() {
+            if null_docs.contains(&doc_pos) {
+                writer.add_document(doc!()).unwrap();
+            } else {
+                writer
+                    .add_document(doc!(number => doc_pos as u64, label => *label_value))
+                    .unwrap();
+            }
+        }
+        writer.commit().unwrap();
+        index
+    }
+
+    fn open_source(index: &Index, expression: &str) -> Box<dyn ValueSource> {
+        let expression = jitexpr::ast::deserialize(expression).unwrap();
+        let provider = JitExprValueSourceProvider::new(expression).unwrap();
+        let searcher = index.reader().unwrap().searcher();
+        provider
+            .for_segment(searcher.segment_reader(0), None)
+            .unwrap()
+    }
+
+    /// Returns the values and docids of a block. `docids` is only filled for a non-full block.
+    fn load_block(
+        source: &mut dyn ValueSource,
+        docs: &[DocId],
+    ) -> (Cardinality, Vec<u64>, Vec<DocId>) {
+        let mut values = Vec::new();
+        let mut docids = Vec::new();
+        let mut row_ids = Vec::new();
+        let cardinality = source.load_block(docs, &mut values, &mut docids, &mut row_ids);
+        if cardinality.is_full() {
+            docids.clear();
+        }
+        (cardinality, values, docids)
+    }
+
+    fn terms(source: &dyn ValueSource, term_ords: &[u64]) -> Vec<String> {
+        let mut terms = Vec::new();
+        let all_found = source
+            .term_dictionary()
+            .unwrap()
+            .sorted_ords_to_term_cb(term_ords, &mut |term| {
+                terms.push(String::from_utf8(term.to_vec()).unwrap())
+            })
+            .unwrap();
+        assert!(all_found);
+        terms
+    }
+
+    #[test]
+    fn test_block_cardinality() {
+        let index = create_index(&["a", "b", "c", "d"], &[2]);
+        let mut source = open_source(&index, "(ADD number 10u64)");
+        assert_eq!(source.column_type(), ColumnType::U64);
+        assert!(source.term_dictionary().is_none());
+
+        // All of the docs have a value.
+        let (cardinality, values, _) = load_block(&mut *source, &[0, 1, 3]);
+        assert_eq!(cardinality, Cardinality::Full);
+        assert_eq!(values, vec![10, 11, 13]);
+
+        // Doc 2 has no value. The buffers of the previous block are cleared.
+        let (cardinality, values, docids) = load_block(&mut *source, &[1, 2, 3]);
+        assert_eq!(cardinality, Cardinality::Optional);
+        assert_eq!(values, vec![11, 13]);
+        assert_eq!(docids, vec![1, 3]);
+
+        let (cardinality, values, docids) = load_block(&mut *source, &[2]);
+        assert_eq!(cardinality, Cardinality::Optional);
+        assert!(values.is_empty());
+        assert!(docids.is_empty());
+    }
+
+    #[test]
+    fn test_text_term_ords_are_stable_across_blocks() {
+        let index = create_index(&["b", "a", "b", "c", "a"], &[]);
+        let mut source = open_source(&index, "(UPPER label)");
+        assert_eq!(source.column_type(), ColumnType::Str);
+        assert!(!source.term_dictionary().unwrap().ords_sorted_with_terms());
+
+        let (cardinality, first_block_ords, _) = load_block(&mut *source, &[0, 1, 2]);
+        assert_eq!(cardinality, Cardinality::Full);
+        // Ords are assigned in first-seen order.
+        assert_eq!(first_block_ords, vec![0, 1, 0]);
+
+        let (_, second_block_ords, _) = load_block(&mut *source, &[3, 4]);
+        assert_eq!(second_block_ords, vec![2, 1]);
+        assert_eq!(source.term_dictionary().unwrap().num_terms(), 3);
+        assert_eq!(terms(&*source, &[0, 1, 2]), vec!["B", "A", "C"]);
+    }
+
+    #[test]
+    fn test_constant_source() {
+        let index = create_index(&["a", "b"], &[]);
+        let mut source = open_source(&index, r#"(UPPER "x")"#);
+        assert_eq!(source.column_type(), ColumnType::Str);
+        let (cardinality, values, _) = load_block(&mut *source, &[0, 1]);
+        assert_eq!(cardinality, Cardinality::Full);
+        assert_eq!(values, vec![0, 0]);
+        assert_eq!(terms(&*source, &[0]), vec!["X"]);
+
+        // A null constant is a source without any value.
+        let mut source = open_source(&index, "(ADD absent 1u64)");
+        let (cardinality, values, docids) = load_block(&mut *source, &[0, 1]);
+        assert!(!cardinality.is_full());
+        assert!(values.is_empty());
+        assert!(docids.is_empty());
+    }
+}

src/aggregation/value_source/mod.rs

@@ -1,4 +1,6 @@
 mod block_accessor;
+#[cfg(feature = "jitexpr")]
+mod jitexpr_value_source;
 mod value_source_registry;
 
 #[cfg(test)]
@@ -9,6 +11,8 @@ use std::io;
 
 pub(crate) use block_accessor::ColumnBlockAccessor;
 use columnar::{Cardinality, Column, ColumnType, ColumnValues, Dictionary, RowId, StrColumn};
+#[cfg(feature = "jitexpr")]
+pub use jitexpr_value_source::JitExprValueSourceProvider;
 pub use value_source_registry::{ValueSourceProvider, ValueSourceRegistry};
 
 use crate::DocId;

src/query/doc_predicate_query/jitexpr_predicate.rs

@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 use std::io;
 
-use columnar::{ColumnIndex, ColumnType, DynamicColumn, StrColumn};
+use columnar::{ColumnIndex, ColumnType, DynamicColumn, DynamicColumnHandle, StrColumn};
 use common::VersatileBuffer;
 use jitexpr::ast::{
     infer_types_with_target, required_presence_for_true, InferredTypeSet, TypeError, UntypedExpr,
@@ -245,10 +245,24 @@ fn open_input_column(
     name: &str,
     accepted_types: InferredTypeSet,
 ) -> io::Result<Option<DynamicColumn>> {
+    let Some(handle) = find_input_column_handle(reader, name, accepted_types) else {
+        return Ok(None);
+    };
+    Ok(Some(handle.open()?))
+}
+
+/// Returns the handle of the first column of `name` with a type in `accepted_types`.
+///
+/// Returns `None` if `name` is not a fast field, or has no column of an accepted type.
+pub(crate) fn find_input_column_handle(
+    reader: &SegmentReader,
+    name: &str,
+    accepted_types: InferredTypeSet,
+) -> Option<DynamicColumnHandle> {
     let Ok(column_handles) = reader.fast_fields().dynamic_column_handles(name) else {
         // If the call to dynamic_column_handles fails (for instance because the column is not a
         // fast field) we choose to act as if the column was absent.
-        return Ok(None);
+        return None;
     };
     for handle in column_handles {
         // We return the first column that could be accepted
@@ -256,13 +270,13 @@ fn open_input_column(
             continue;
         };
         if accepted_types.contains(var_type) {
-            return Ok(Some(handle.open()?));
+            return Some(handle);
         }
     }
-    Ok(None)
+    None
 }
 
-fn var_type_for_column_type(column_type: ColumnType) -> Option<VarType> {
+pub(crate) fn var_type_for_column_type(column_type: ColumnType) -> Option<VarType> {
     match column_type {
         ColumnType::Bool => Some(VarType::Bool),
         ColumnType::I64 => Some(VarType::I64),

src/query/doc_predicate_query/mod.rs

@@ -7,6 +7,8 @@ mod jitexpr_predicate;
 
 pub use function_predicate::FunctionPredicate;
 #[cfg(feature = "jitexpr")]
+pub(crate) use jitexpr_predicate::{find_input_column_handle, var_type_for_column_type};
+#[cfg(feature = "jitexpr")]
 pub use jitexpr_predicate::{JitExprEvalState, JitExprPredicate};
 
 use crate::docset::{SeekDangerResult, TERMINATED};