From 3f1be50ff0aed6c79d629e300fb2a8268903dead Mon Sep 17 00:00:00 2001 From: Paul Masurel Date: Thu, 8 Oct 2026 10:55:36 +0200 Subject: [PATCH] First stab --- Cargo.toml | 4 +- src/aggregation/agg_data.rs | 36 +- src/aggregation/mod.rs | 2 + .../dynamic_term_dictionary.rs | 249 ++++++ .../value_source/jitexpr_value_source/mod.rs | 733 ++++++++++++++++++ .../jitexpr_value_source/segment_source.rs | 462 +++++++++++ src/aggregation/value_source/mod.rs | 4 + .../doc_predicate_query/jitexpr_predicate.rs | 24 +- src/query/doc_predicate_query/mod.rs | 2 + 9 files changed, 1507 insertions(+), 9 deletions(-) create mode 100644 src/aggregation/value_source/jitexpr_value_source/dynamic_term_dictionary.rs create mode 100644 src/aggregation/value_source/jitexpr_value_source/mod.rs create mode 100644 src/aggregation/value_source/jitexpr_value_source/segment_source.rs diff --git a/Cargo.toml b/Cargo.toml index 21521b9ac..e23d8d7d1 100644 --- a/Cargo.toml +++ b/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. diff --git a/src/aggregation/agg_data.rs b/src/aggregation/agg_data.rs index 7fc50f92b..12c56382d 100644 --- a/src/aggregation/agg_data.rs +++ b/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. /// diff --git a/src/aggregation/mod.rs b/src/aggregation/mod.rs index ba9e14e53..f076caa15 100644 --- a/src/aggregation/mod.rs +++ b/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, }; diff --git a/src/aggregation/value_source/jitexpr_value_source/dynamic_term_dictionary.rs b/src/aggregation/value_source/jitexpr_value_source/dynamic_term_dictionary.rs new file mode 100644 index 000000000..865828a72 --- /dev/null +++ b/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, + /// `buffer[offsets[ord]..offsets[ord + 1]]` is the term of `ord`. + /// + /// Invariant: `offsets[0] == 0`, and `offsets.len() == num_terms() + 1`. + offsets: Vec, + /// One item per term, hashed with `table_hash(item.hash)`. + term_ords: HashTable, +} + +/// 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 { + // 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 = 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 = + 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())); + } +} diff --git a/src/aggregation/value_source/jitexpr_value_source/mod.rs b/src/aggregation/value_source/jitexpr_value_source/mod.rs new file mode 100644 index 000000000..ae794606d --- /dev/null +++ b/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 { + 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> { + // 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 = 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 = 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 { + 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 { + 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::::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 { + 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 = 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); + } +} diff --git a/src/aggregation/value_source/jitexpr_value_source/segment_source.rs b/src/aggregation/value_source/jitexpr_value_source/segment_source.rs new file mode 100644 index 000000000..23dbeeef6 --- /dev/null +++ b/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, + /// 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, + ) -> 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, + docids: &mut Vec, + _row_ids: &mut Vec, + ) -> 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 = 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 { + 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, + _docids: &mut Vec, + _row_ids: &mut Vec, + ) -> 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 { + // 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, + /// The dictionary of the term ords. `Some` if and only if `var_type` is `Str`. + str_column: Option, + /// `first_values[i]` is the first value of `docs[i]`, for the docs of the last block. + first_values: Vec>, + /// 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 { + 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, + /// 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, +} + +impl BlockTerms { + /// Panics if the dictionary is corrupted: `load_block` cannot return an error. + fn load(&mut self, str_column: &StrColumn, term_ords: &[Option]) { + 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 { + 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, Vec) { + 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 { + 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()); + } +} diff --git a/src/aggregation/value_source/mod.rs b/src/aggregation/value_source/mod.rs index 930ca35c5..b64193272 100644 --- a/src/aggregation/value_source/mod.rs +++ b/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; diff --git a/src/query/doc_predicate_query/jitexpr_predicate.rs b/src/query/doc_predicate_query/jitexpr_predicate.rs index 20e62cb19..f8b1e29a3 100644 --- a/src/query/doc_predicate_query/jitexpr_predicate.rs +++ b/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> { + 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 { 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 { +pub(crate) fn var_type_for_column_type(column_type: ColumnType) -> Option { match column_type { ColumnType::Bool => Some(VarType::Bool), ColumnType::I64 => Some(VarType::I64), diff --git a/src/query/doc_predicate_query/mod.rs b/src/query/doc_predicate_query/mod.rs index e8c74e106..ddedb1d2e 100644 --- a/src/query/doc_predicate_query/mod.rs +++ b/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};