use std::sync::Arc; use columnar::{Column, ColumnType, StrColumn}; use common::BitSet; use rustc_hash::FxHashSet; use serde::Serialize; use tantivy_fst::Regex; use crate::aggregation::accessor_helpers::{ get_all_ff_reader_or_empty, get_dynamic_columns, get_ff_reader, get_missing_val_as_u64_lenient, get_numeric_or_date_column_types, }; use crate::aggregation::agg_req::{Aggregation, AggregationVariants, Aggregations}; use crate::aggregation::bucket::{ build_segment_filter_collector, build_segment_histogram_collector, build_segment_multi_terms_collector, build_segment_range_collector, CompositeAggReqData, CompositeAggregation, CompositeSourceAccessors, FilterAggReqData, HistogramAggReqData, HistogramBounds, IncludeExcludeParam, MissingTermAggReqData, MultiTermsAggReqData, MultiTermsAggregation, MultiTermsFieldAccessor, MultiTermsMissingAccessor, RangeAggReqData, TermMissingAgg, TermsAggReqData, TermsAggregation, TermsAggregationInternal, }; use crate::aggregation::metric::{ build_segment_stats_collector, AverageAggregation, CardinalityAggReqData, CardinalityAggregationReq, CountAggregation, ExtendedStatsAggregation, MaxAggregation, MetricAggReqData, MinAggregation, SegmentCardinalityCollector, SegmentExtendedStatsCollector, SegmentPercentilesCollector, StatsAggregation, StatsType, SumAggregation, TermOrdSet, TopHitsAggReqData, TopHitsSegmentCollector, BITSET_MAX_TERM_ORD, }; use crate::aggregation::segment_agg_result::{ GenericSegmentAggregationResultsCollector, SegmentAggregationCollector, }; use crate::aggregation::{f64_to_fastfield_u64, AggContextParams, ColumnBlockAccessor, Key}; use crate::{SegmentOrdinal, SegmentReader}; #[derive(Default)] /// Datastructure holding all request data for executing aggregations on a segment. /// It is passed to the collectors during collection. pub struct AggregationsSegmentCtx { /// Request data for each aggregation type. pub per_request: PerRequestAggSegCtx, pub context: AggContextParams, pub(crate) column_block_accessor: ColumnBlockAccessor, } impl AggregationsSegmentCtx { pub(crate) fn push_term_req_data(&mut self, data: TermsAggReqData) -> usize { self.per_request.term_req_data.push(data); self.per_request.term_req_data.len() - 1 } pub(crate) fn push_cardinality_req_data(&mut self, data: CardinalityAggReqData) -> usize { self.per_request.cardinality_req_data.push(data); self.per_request.cardinality_req_data.len() - 1 } pub(crate) fn push_metric_req_data(&mut self, data: MetricAggReqData) -> usize { self.per_request.stats_metric_req_data.push(data); self.per_request.stats_metric_req_data.len() - 1 } pub(crate) fn push_top_hits_req_data(&mut self, data: TopHitsAggReqData) -> usize { self.per_request.top_hits_req_data.push(data); self.per_request.top_hits_req_data.len() - 1 } pub(crate) fn push_missing_term_req_data(&mut self, data: MissingTermAggReqData) -> usize { self.per_request.missing_term_req_data.push(data); self.per_request.missing_term_req_data.len() - 1 } pub(crate) fn push_histogram_req_data(&mut self, data: HistogramAggReqData) -> usize { self.per_request.histogram_req_data.push(data); self.per_request.histogram_req_data.len() - 1 } pub(crate) fn push_range_req_data(&mut self, data: RangeAggReqData) -> usize { self.per_request.range_req_data.push(data); self.per_request.range_req_data.len() - 1 } pub(crate) fn push_filter_req_data(&mut self, data: FilterAggReqData) -> usize { self.per_request.filter_req_data.push(data); self.per_request.filter_req_data.len() - 1 } pub(crate) fn push_composite_req_data(&mut self, data: CompositeAggReqData) -> usize { self.per_request.composite_req_data.push(data); self.per_request.composite_req_data.len() - 1 } pub(crate) fn push_multi_terms_req_data(&mut self, data: MultiTermsAggReqData) -> usize { self.per_request.multi_terms_req_data.push(data); self.per_request.multi_terms_req_data.len() - 1 } #[inline] pub(crate) fn get_term_req_data(&self, idx: usize) -> &TermsAggReqData { &self.per_request.term_req_data[idx] } #[inline] pub(crate) fn get_cardinality_req_data(&self, idx: usize) -> &CardinalityAggReqData { &self.per_request.cardinality_req_data[idx] } #[inline] pub(crate) fn get_metric_req_data(&self, idx: usize) -> &MetricAggReqData { &self.per_request.stats_metric_req_data[idx] } #[inline] pub(crate) fn get_top_hits_req_data(&self, idx: usize) -> &TopHitsAggReqData { &self.per_request.top_hits_req_data[idx] } #[inline] pub(crate) fn get_missing_term_req_data(&self, idx: usize) -> &MissingTermAggReqData { &self.per_request.missing_term_req_data[idx] } } /// Each type of aggregation has its own request data struct. This struct holds /// all request data to execute the aggregation request on a single segment. /// /// The request tree is represented by `agg_tree`. Tree nodes contain the index /// of their context in corresponding request data vector (e.g. `term_req_data` /// for a node with [AggKind::Terms]). #[derive(Default)] pub struct PerRequestAggSegCtx { /// TermsAggReqData contains the request data for a terms aggregation. pub term_req_data: Vec, /// HistogramAggReqData contains the request data for a histogram aggregation. pub histogram_req_data: Vec, /// RangeAggReqData contains the request data for a range aggregation. pub range_req_data: Vec, /// FilterAggReqData contains the request data for a filter aggregation. pub filter_req_data: Vec, /// Shared by avg, min, max, sum, stats, extended_stats, count pub stats_metric_req_data: Vec, /// CardinalityAggReqData contains the request data for a cardinality aggregation. pub cardinality_req_data: Vec, /// TopHitsAggReqData contains the request data for a top_hits aggregation. pub top_hits_req_data: Vec, /// MissingTermAggReqData contains the request data for a missing term aggregation. pub missing_term_req_data: Vec, /// CompositeAggReqData contains the request data for a composite aggregation. pub composite_req_data: Vec, /// MultiTermsAggReqData contains the request data for a multi_terms aggregation. pub multi_terms_req_data: Vec, /// Request tree used to build collectors. pub agg_tree: Vec, } impl PerRequestAggSegCtx { /// Estimate the memory consumption of this struct in bytes. fn get_memory_consumption(&self) -> usize { self.term_req_data .iter() .map(|t| t.get_memory_consumption()) .sum::() + self .histogram_req_data .iter() .map(|t| t.get_memory_consumption()) .sum::() + self .range_req_data .iter() .map(|t| t.get_memory_consumption()) .sum::() + self .filter_req_data .iter() .map(|t| t.get_memory_consumption()) .sum::() + self .stats_metric_req_data .iter() .map(|t| t.get_memory_consumption()) .sum::() + self .cardinality_req_data .iter() .map(|t| t.get_memory_consumption()) .sum::() + self .top_hits_req_data .iter() .map(|t| t.get_memory_consumption()) .sum::() + self .missing_term_req_data .iter() .map(|t| t.get_memory_consumption()) .sum::() + self .composite_req_data .iter() .map(|t| t.get_memory_consumption()) .sum::() + self .multi_terms_req_data .iter() .map(|t| t.get_memory_consumption()) .sum::() + self.agg_tree.len() * std::mem::size_of::() } pub fn get_name(&self, node: &AggRefNode) -> &str { let idx = node.idx_in_req_data; let kind = node.kind; match kind { AggKind::Terms => self.term_req_data[idx].name.as_str(), AggKind::Cardinality => &self.cardinality_req_data[idx].name, AggKind::StatsKind(_) => &self.stats_metric_req_data[idx].name, AggKind::TopHits => &self.top_hits_req_data[idx].name, AggKind::MissingTerm => &self.missing_term_req_data[idx].name, AggKind::Histogram => self.histogram_req_data[idx].name.as_str(), AggKind::DateHistogram => self.histogram_req_data[idx].name.as_str(), AggKind::Range => self.range_req_data[idx].name.as_str(), AggKind::Filter => self.filter_req_data[idx].name.as_str(), AggKind::Composite => self.composite_req_data[idx].name.as_str(), AggKind::MultiTerms => self.multi_terms_req_data[idx].name.as_str(), } } /// Convert the aggregation tree into a serializable struct representation. /// Each node contains: { name, kind, children }. #[allow(dead_code)] pub fn get_view_tree(&self) -> Vec { fn node_to_view(node: &AggRefNode, pr: &PerRequestAggSegCtx) -> AggTreeViewNode { let mut children: Vec = node.children.iter().map(|c| node_to_view(c, pr)).collect(); children.sort_by_key(|v| serde_json::to_string(v).unwrap()); AggTreeViewNode { name: pr.get_name(node).to_string(), kind: node.kind.as_str().to_string(), children, } } let mut roots: Vec = self .agg_tree .iter() .map(|n| node_to_view(n, self)) .collect(); roots.sort_by_key(|v| serde_json::to_string(v).unwrap()); roots } } pub(crate) fn build_segment_agg_collectors_root( req: &mut AggregationsSegmentCtx, ) -> crate::Result> { build_segment_agg_collectors_generic(req, &req.per_request.agg_tree.clone()) } pub(crate) fn build_segment_agg_collectors( req: &mut AggregationsSegmentCtx, nodes: &[AggRefNode], ) -> crate::Result> { build_segment_agg_collectors_generic(req, nodes) } fn build_segment_agg_collectors_generic( req: &mut AggregationsSegmentCtx, nodes: &[AggRefNode], ) -> crate::Result> { let mut collectors = Vec::new(); for node in nodes.iter() { collectors.push(build_segment_agg_collector(req, node)?); } req.context .limits .add_memory_consumed(req.per_request.get_memory_consumption() as u64)?; // Single collector special case if collectors.len() == 1 { return Ok(collectors.pop().unwrap()); } let agg = GenericSegmentAggregationResultsCollector { aggs: collectors }; Ok(Box::new(agg)) } pub(crate) fn build_segment_agg_collector( req: &mut AggregationsSegmentCtx, node: &AggRefNode, ) -> crate::Result> { match node.kind { AggKind::Terms => crate::aggregation::bucket::build_segment_term_collector(req, node), AggKind::MissingTerm => { let req_data = &mut req.per_request.missing_term_req_data[node.idx_in_req_data]; if req_data.accessors.is_empty() { return Err(crate::TantivyError::InternalError( "MissingTerm aggregation requires at least one field accessor.".to_string(), )); } Ok(Box::new(TermMissingAgg::new(req, node)?)) } AggKind::Cardinality => { let req_data = req.get_cardinality_req_data(node.idx_in_req_data); // For str columns, choose the per-bucket entries representation // based on the segment's column.max_value(): // * small (< BITSET_MAX_TERM_ORD): `BitSet`, pre-allocated, no promotion machinery. // * large: `TermOrdSet` (sparse FxHashSet that promotes to a paged bitset). // For non-str columns the `entries` field is unused (values go // straight into the HLL sketch); we still pick `TermOrdSet` // because its empty Sparse(FxHashSet) costs nothing. let is_str = req_data.column_type == ColumnType::Str; let max_term_ord_inclusive = if is_str { req_data.accessor.max_value() } else { 0 }; let collector: Box = if is_str && max_term_ord_inclusive < BITSET_MAX_TERM_ORD { Box::new(SegmentCardinalityCollector::::from_req( req_data.column_type, node.idx_in_req_data, req_data.accessor.clone(), req_data.missing_value_for_accessor, max_term_ord_inclusive, )) } else { Box::new(SegmentCardinalityCollector::::from_req( req_data.column_type, node.idx_in_req_data, req_data.accessor.clone(), req_data.missing_value_for_accessor, max_term_ord_inclusive, )) }; Ok(collector) } AggKind::StatsKind(stats_type) => { let req_data = &mut req.per_request.stats_metric_req_data[node.idx_in_req_data]; match stats_type { StatsType::Sum | StatsType::Average | StatsType::Count | StatsType::Max | StatsType::Min | StatsType::Stats => build_segment_stats_collector(req_data), StatsType::ExtendedStats(sigma) => Ok(Box::new( SegmentExtendedStatsCollector::from_req(req_data, sigma), )), StatsType::Percentiles => { let req_data = req.get_metric_req_data(node.idx_in_req_data); Ok(Box::new( SegmentPercentilesCollector::from_req_and_validate( req_data.field_type, req_data.missing_u64, req_data.accessor.clone(), node.idx_in_req_data, ), )) } } } AggKind::TopHits => { let req_data = &mut req.per_request.top_hits_req_data[node.idx_in_req_data]; Ok(Box::new(TopHitsSegmentCollector::from_req( &req_data.req, node.idx_in_req_data, req_data.segment_ordinal, ))) } AggKind::Histogram => build_segment_histogram_collector(req, node), AggKind::DateHistogram => build_segment_histogram_collector(req, node), AggKind::Range => Ok(build_segment_range_collector(req, node)?), AggKind::Filter => build_segment_filter_collector(req, node), AggKind::Composite => Ok(Box::new( crate::aggregation::bucket::SegmentCompositeCollector::from_req_and_validate( req, node, )?, )), AggKind::MultiTerms => build_segment_multi_terms_collector(req, node), } } /// See [PerRequestAggSegCtx] #[derive(Debug, Clone)] pub struct AggRefNode { pub kind: AggKind, pub idx_in_req_data: usize, pub children: Vec, } impl AggRefNode { pub fn get_sub_agg(&self, name: &str, pr: &PerRequestAggSegCtx) -> Option<&AggRefNode> { self.children .iter() .find(|&child| pr.get_name(child) == name) } } #[derive(Copy, Clone, Debug)] pub enum AggKind { Terms, Cardinality, /// One of: Statistics, Average, Min, Max, Sum, Count, Stats, ExtendedStats StatsKind(StatsType), TopHits, MissingTerm, Histogram, DateHistogram, Range, Filter, Composite, MultiTerms, } impl AggKind { #[cfg_attr(not(test), allow(dead_code))] fn as_str(&self) -> &'static str { match self { AggKind::Terms => "Terms", AggKind::Cardinality => "Cardinality", AggKind::StatsKind(_) => "Metric", AggKind::TopHits => "TopHits", AggKind::MissingTerm => "MissingTerm", AggKind::Histogram => "Histogram", AggKind::DateHistogram => "DateHistogram", AggKind::Range => "Range", AggKind::Filter => "Filter", AggKind::Composite => "Composite", AggKind::MultiTerms => "MultiTerms", } } } /// Build AggregationsData by walking the request tree. pub(crate) fn build_aggregations_data_from_req( aggs: &Aggregations, reader: &SegmentReader, segment_ordinal: SegmentOrdinal, context: AggContextParams, ) -> crate::Result { let mut data = AggregationsSegmentCtx { per_request: Default::default(), context, column_block_accessor: ColumnBlockAccessor::default(), }; for (name, agg) in aggs.iter() { let nodes = build_nodes(name, agg, reader, segment_ordinal, &mut data, true)?; data.per_request.agg_tree.extend(nodes); } Ok(data) } fn build_nodes( agg_name: &str, req: &Aggregation, reader: &SegmentReader, segment_ordinal: SegmentOrdinal, data: &mut AggregationsSegmentCtx, is_top_level: bool, ) -> crate::Result> { use AggregationVariants::*; match &req.agg { Range(range_req) => { let (accessor, field_type) = get_ff_reader( reader, &range_req.field, Some(get_numeric_or_date_column_types()), )?; let idx_in_req_data = data.push_range_req_data(RangeAggReqData { accessor, field_type, name: agg_name.to_string(), req: range_req.clone(), is_top_level, }); let children = build_children(&req.sub_aggregation, reader, segment_ordinal, data)?; Ok(vec![AggRefNode { kind: AggKind::Range, idx_in_req_data, children, }]) } Histogram(histo_req) => { let (accessor, field_type) = get_ff_reader( reader, &histo_req.field, Some(get_numeric_or_date_column_types()), )?; let idx_in_req_data = data.push_histogram_req_data(HistogramAggReqData { accessor, field_type, name: agg_name.to_string(), req: histo_req.clone(), is_date_histogram: false, bounds: HistogramBounds { min: f64::MIN, max: f64::MAX, }, offset: 0.0, }); let children = build_children(&req.sub_aggregation, reader, segment_ordinal, data)?; Ok(vec![AggRefNode { kind: AggKind::Histogram, idx_in_req_data, children, }]) } DateHistogram(date_req) => { let (accessor, field_type) = get_ff_reader(reader, &date_req.field, Some(&[ColumnType::DateTime]))?; // Convert to histogram request, normalize to ns precision let mut histo_req = date_req.to_histogram_req()?; histo_req.normalize_date_time(); let idx_in_req_data = data.push_histogram_req_data(HistogramAggReqData { accessor, field_type, name: agg_name.to_string(), req: histo_req, is_date_histogram: true, bounds: HistogramBounds { min: f64::MIN, max: f64::MAX, }, offset: 0.0, }); let children = build_children(&req.sub_aggregation, reader, segment_ordinal, data)?; Ok(vec![AggRefNode { kind: AggKind::DateHistogram, idx_in_req_data, children, }]) } Terms(terms_req) => build_terms_or_cardinality_nodes( agg_name, &terms_req.field, &terms_req.missing, reader, segment_ordinal, data, &req.sub_aggregation, TermsOrCardinalityRequest::Terms(terms_req.clone()), is_top_level, ), Cardinality(card_req) => build_terms_or_cardinality_nodes( agg_name, &card_req.field, &card_req.missing, reader, segment_ordinal, data, &req.sub_aggregation, TermsOrCardinalityRequest::Cardinality(card_req.clone()), is_top_level, ), Average(AverageAggregation { field, missing, .. }) | Max(MaxAggregation { field, missing, .. }) | Min(MinAggregation { field, missing, .. }) | Stats(StatsAggregation { field, missing, .. }) | ExtendedStats(ExtendedStatsAggregation { field, missing, .. }) | Sum(SumAggregation { field, missing, .. }) | Count(CountAggregation { field, missing, .. }) => { let allowed_column_types = if matches!(&req.agg, Count(_)) { Some( &[ ColumnType::I64, ColumnType::U64, ColumnType::F64, ColumnType::Str, ColumnType::DateTime, ColumnType::Bool, ColumnType::IpAddr, ][..], ) } else { Some(get_numeric_or_date_column_types()) }; let collecting_for = match &req.agg { Average(_) => StatsType::Average, Max(_) => StatsType::Max, Min(_) => StatsType::Min, Stats(_) => StatsType::Stats, ExtendedStats(req) => StatsType::ExtendedStats(req.sigma), Sum(_) => StatsType::Sum, Count(_) => StatsType::Count, _ => { return Err(crate::TantivyError::InvalidArgument( "Internal error: unexpected aggregation type in metric aggregation \ handling." .to_string(), )) } }; let (accessor, field_type) = get_ff_reader(reader, field, allowed_column_types)?; let idx_in_req_data = data.push_metric_req_data(MetricAggReqData { accessor, field_type, name: agg_name.to_string(), collecting_for, missing: *missing, missing_u64: (*missing).and_then(|m| f64_to_fastfield_u64(m, &field_type)), is_number_or_date_type: matches!( field_type, ColumnType::I64 | ColumnType::U64 | ColumnType::F64 | ColumnType::DateTime ), }); let children = build_children(&req.sub_aggregation, reader, segment_ordinal, data)?; Ok(vec![AggRefNode { kind: AggKind::StatsKind(collecting_for), idx_in_req_data, children, }]) } // Percentiles handled as Metric as well AggregationVariants::Percentiles(percentiles_req) => { percentiles_req.validate()?; let (accessor, field_type) = get_ff_reader( reader, percentiles_req.field_name(), Some(get_numeric_or_date_column_types()), )?; let idx_in_req_data = data.push_metric_req_data(MetricAggReqData { accessor, field_type, name: agg_name.to_string(), collecting_for: StatsType::Percentiles, missing: percentiles_req.missing, missing_u64: percentiles_req .missing .and_then(|m| f64_to_fastfield_u64(m, &field_type)), is_number_or_date_type: matches!( field_type, ColumnType::I64 | ColumnType::U64 | ColumnType::F64 | ColumnType::DateTime ), }); let children = build_children(&req.sub_aggregation, reader, segment_ordinal, data)?; Ok(vec![AggRefNode { kind: AggKind::StatsKind(StatsType::Percentiles), idx_in_req_data, children, }]) } AggregationVariants::TopHits(top_hits_req) => { let mut top_hits = top_hits_req.clone(); top_hits.validate_and_resolve_field_names(reader.fast_fields().columnar())?; let accessors: Vec<(Column, ColumnType)> = top_hits .field_names() .iter() .map(|field| get_ff_reader(reader, field, Some(get_numeric_or_date_column_types()))) .collect::>()?; let value_accessors = top_hits .value_field_names() .iter() .map(|field_name| { Ok(( field_name.to_string(), get_dynamic_columns(reader, field_name)?, )) }) .collect::>()?; let idx_in_req_data = data.push_top_hits_req_data(TopHitsAggReqData { accessors, value_accessors, segment_ordinal, name: agg_name.to_string(), req: top_hits.clone(), }); let children = build_children(&req.sub_aggregation, reader, segment_ordinal, data)?; Ok(vec![AggRefNode { kind: AggKind::TopHits, idx_in_req_data, children, }]) } AggregationVariants::Composite(composite_req) => Ok(vec![build_composite_node( agg_name, reader, segment_ordinal, data, &req.sub_aggregation, composite_req, )?]), AggregationVariants::MultiTerms(multi_terms_req) => build_multi_terms_nodes( agg_name, reader, segment_ordinal, data, &req.sub_aggregation, multi_terms_req, is_top_level, ), AggregationVariants::Filter(filter_req) => { // Build the query and evaluator upfront let schema = reader.schema(); let tokenizers = &data.context.tokenizers; let query = filter_req.parse_query(schema, tokenizers)?; let evaluator = std::rc::Rc::new(crate::aggregation::bucket::DocumentQueryEvaluator::new( query, schema.clone(), reader, )?); let idx_in_req_data = data.push_filter_req_data(FilterAggReqData { name: agg_name.to_string(), req: filter_req.clone(), segment_reader: reader.clone(), evaluator, is_top_level, }); let children = build_children(&req.sub_aggregation, reader, segment_ordinal, data)?; Ok(vec![AggRefNode { kind: AggKind::Filter, idx_in_req_data, children, }]) } } } fn build_composite_node( agg_name: &str, reader: &SegmentReader, _segment_ordinal: SegmentOrdinal, data: &mut AggregationsSegmentCtx, sub_aggs: &Aggregations, req: &CompositeAggregation, ) -> crate::Result { let mut composite_accessors = Vec::with_capacity(req.sources.len()); for source in &req.sources { let source_after_key_opt = req.after.get(source.name()).map(|k| &k.0); let source_accessor = CompositeSourceAccessors::build_for_source(reader, source, source_after_key_opt)?; composite_accessors.push(source_accessor); } let agg = CompositeAggReqData { name: agg_name.to_string(), req: req.clone(), composite_accessors, }; let idx = data.push_composite_req_data(agg); let children = build_children(sub_aggs, reader, _segment_ordinal, data)?; Ok(AggRefNode { kind: AggKind::Composite, idx_in_req_data: idx, children, }) } fn build_multi_terms_nodes( agg_name: &str, reader: &SegmentReader, segment_ordinal: SegmentOrdinal, data: &mut AggregationsSegmentCtx, sub_aggs: &Aggregations, req: &MultiTermsAggregation, is_top_level: bool, ) -> crate::Result> { if req.terms.is_empty() { return Err(crate::TantivyError::InvalidArgument( "multi_terms aggregation requires at least one field".to_string(), )); } let mut accessors_by_field = Vec::with_capacity(req.terms.len()); for field_def in &req.terms { let field_name = &field_def.field; let str_dict_column = reader.fast_fields().str(field_name)?; let columns = get_term_agg_accessors(reader, field_name, &field_def.missing, true)?; if let Some((_, column_type)) = columns .iter() .find(|(_, column_type)| *column_type == ColumnType::Bytes) { return Err(crate::TantivyError::InvalidArgument(format!( "multi_terms aggregation is not supported for column type {:?} in field {}", column_type, field_name ))); } // Exactly one typed accessor choice carries missing handling. It checks all physical // columns before injecting the fallback, so a value in another type-specific collector is // not mistaken for a missing field and a genuinely missing document is not counted once // per type. let missing_accessor = prepare_multi_terms_missing( &columns, str_dict_column.as_ref(), field_def.missing.as_ref(), field_name, )?; let mut typed_accessors = Vec::with_capacity(columns.len()); for (column_idx, (column, column_type)) in columns.into_iter().enumerate() { let missing = match &missing_accessor { Some((missing_idx, missing)) if *missing_idx == column_idx => Some(missing.clone()), _ => None, }; typed_accessors.push(( MultiTermsFieldAccessor { column, column_type, str_dict_column: if column_type == ColumnType::Str { str_dict_column.clone() } else { None }, field: field_name.clone(), }, missing, )); } accessors_by_field.push(typed_accessors); } // Fan out one collector for every Cartesian product of physical column choices. Collectors // share the aggregation name, so their intermediate buckets are merged by // `IntermediateAggregationResults::push`. As with terms aggregation type fan-out, // `segment_size` is applied per physical combination and the merged error/count metadata is // therefore the sum of those independently cut-off results. let mut field_combinations: Vec< Vec<(MultiTermsFieldAccessor, Option)>, > = vec![Vec::with_capacity(req.terms.len())]; for typed_accessors in accessors_by_field { let mut next = Vec::new(); for field_choices in field_combinations { for typed_accessor in &typed_accessors { let mut combination = field_choices.clone(); combination.push(typed_accessor.clone()); next.push(combination); } } field_combinations = next; } let mut nodes = Vec::with_capacity(field_combinations.len()); for field_choices in field_combinations { let (fields, missing_accessors) = field_choices.into_iter().unzip(); let idx = data.push_multi_terms_req_data(MultiTermsAggReqData { name: agg_name.to_string(), req: req.clone(), fields, missing_accessors, sub_aggregations: sub_aggs.clone(), is_top_level, }); let children = build_children(sub_aggs, reader, segment_ordinal, data)?; nodes.push(AggRefNode { kind: AggKind::MultiTerms, idx_in_req_data: idx, children, }); } Ok(nodes) } fn prepare_multi_terms_missing( columns: &[(Column, ColumnType)], str_dict_column: Option<&StrColumn>, missing: Option<&Key>, field_name: &str, ) -> crate::Result> { let Some(missing) = missing else { return Ok(None); }; // Attach string fallbacks to the string column when one exists. A string fallback on any // other physical type is handled synthetically, just like the special terms missing // collector. let column_idx = if matches!(missing, Key::Str(_)) { columns .iter() .position(|(_, column_type)| *column_type == ColumnType::Str) .unwrap_or(0) } else { // Prefer an exact physical type for numeric missing values, then any numerical type, then // a string column (which accepts numeric fallbacks through a synthetic value). let preferred_type = match missing { Key::F64(_) => ColumnType::F64, Key::I64(_) => ColumnType::I64, Key::U64(_) => ColumnType::U64, Key::Str(_) => unreachable!("handled above"), }; columns .iter() .position(|(_, column_type)| *column_type == preferred_type) .or_else(|| { columns .iter() .position(|(_, column_type)| column_type.numerical_type().is_some()) }) .or_else(|| { columns .iter() .position(|(_, column_type)| *column_type == ColumnType::Str) }) .unwrap_or(0) }; let (column, column_type) = &columns[column_idx]; if !matches!(missing, Key::Str(_)) && *column_type != ColumnType::Str { // Validate the same lenient numeric coercions as a terms aggregation. get_missing_val_as_u64_lenient(*column_type, column.max_value(), missing, field_name)?; } // Reuse an existing term ordinal so real and missing values enter the same bucket before the // segment-level cutoff. Non-string columns and missing terms absent from the dictionary keep // using a collision-free sentinel. let existing_term_ord = match (missing, *column_type, str_dict_column) { (Key::Str(missing_str), ColumnType::Str, Some(str_dict_column)) => str_dict_column .dictionary() .term_ord(missing_str.as_bytes())?, _ => None, }; // A full physical column means every document has a value for this logical field, so no typed // collector branch can ever emit the configured missing value. if columns .iter() .any(|(column, _)| column.get_cardinality().is_full()) { return Ok(None); } let all_columns = Arc::from( columns .iter() .map(|(column, _)| column.clone()) .collect::>(), ); Ok(Some(( column_idx, MultiTermsMissingAccessor { all_columns, key: missing.clone(), missing_value: existing_term_ord.unwrap_or_else(|| find_missing_sentinel(column)), }, ))) } /// Returns a value that cannot collide with a value in `column`. /// /// Usually one of the column bounds leaves a free value. Only a column whose bounds span the /// entire `u64` domain requires the slower scan. fn find_missing_sentinel(column: &Column) -> u64 { if let Some(sentinel) = column.max_value().checked_add(1) { return sentinel; } if let Some(sentinel) = column.min_value().checked_sub(1) { return sentinel; } // TODO: This is an extreme edge case that would be better handled by a collector that does // not use sentinel missing values. For now, we just scan the column. let values: FxHashSet = column.values.iter().collect(); let mut sentinel = 1u64; while values.contains(&sentinel) { sentinel += 1; } sentinel } fn build_children( aggs: &Aggregations, reader: &SegmentReader, segment_ordinal: SegmentOrdinal, data: &mut AggregationsSegmentCtx, ) -> crate::Result> { let mut children = Vec::new(); for (name, agg) in aggs.iter() { children.extend(build_nodes( name, agg, reader, segment_ordinal, data, false, )?); } Ok(children) } fn get_term_agg_accessors( reader: &SegmentReader, field_name: &str, missing: &Option, include_bytes: bool, ) -> crate::Result, ColumnType)>> { // `terms` and `multi_terms` both explicitly reject `Bytes` columns downstream, which needs // to actually see them as a real column (rather than the empty shim below) to do so. // `cardinality` has no such rejection: it would hash raw `Bytes` term ordinals as if they // were comparable numeric values, but those ordinals are segment-local, so distinct byte // values in different segments could collide and be undercounted. Keep `Bytes` out of its // accessors entirely instead. let mut allowed_column_types = vec![ ColumnType::I64, ColumnType::U64, ColumnType::F64, ColumnType::Str, ColumnType::DateTime, ColumnType::Bool, ColumnType::IpAddr, ]; if include_bytes { allowed_column_types.push(ColumnType::Bytes); } // In case the column is empty we want the shim column to match the missing type let fallback_type = missing .as_ref() .map(|missing| match missing { Key::Str(_) => ColumnType::Str, Key::F64(_) => ColumnType::F64, Key::I64(_) => ColumnType::I64, Key::U64(_) => ColumnType::U64, }) .unwrap_or(ColumnType::U64); let column_and_types = get_all_ff_reader_or_empty( reader, field_name, Some(&allowed_column_types), fallback_type, )?; Ok(column_and_types) } enum TermsOrCardinalityRequest { Terms(TermsAggregation), Cardinality(CardinalityAggregationReq), } impl TermsOrCardinalityRequest { fn as_terms(&self) -> Option<&TermsAggregation> { match self { TermsOrCardinalityRequest::Terms(t) => Some(t), _ => None, } } } #[allow(clippy::too_many_arguments)] fn build_terms_or_cardinality_nodes( agg_name: &str, field_name: &str, missing: &Option, reader: &SegmentReader, segment_ordinal: SegmentOrdinal, data: &mut AggregationsSegmentCtx, sub_aggs: &Aggregations, req: TermsOrCardinalityRequest, is_top_level: bool, ) -> crate::Result> { let mut nodes = Vec::new(); let str_dict_column = reader.fast_fields().str(field_name)?; let include_bytes = matches!(req, TermsOrCardinalityRequest::Terms(_)); let column_and_types = get_term_agg_accessors(reader, field_name, missing, include_bytes)?; // Special handling when missing + multi column or incompatible type on text/date. let missing_and_more_than_one_col = column_and_types.len() > 1 && missing.is_some(); let text_on_non_text_col = column_and_types.len() == 1 && column_and_types[0].1 != ColumnType::Str && matches!(missing, Some(Key::Str(_))); let use_special_missing_agg = missing_and_more_than_one_col || text_on_non_text_col; // If special missing handling is required, build a MissingTerm node that carries all // accessors (across any column types) for existence checks. if use_special_missing_agg { let fallback_type = missing .as_ref() .map(|missing| match missing { Key::Str(_) => ColumnType::Str, Key::F64(_) => ColumnType::F64, Key::I64(_) => ColumnType::I64, Key::U64(_) => ColumnType::U64, }) .unwrap_or(ColumnType::U64); let all_accessors = get_all_ff_reader_or_empty(reader, field_name, None, fallback_type)? .into_iter() .collect::>(); // This case only happens when we have term aggregation, or we fail let req = req.as_terms().cloned().ok_or_else(|| { crate::TantivyError::InvalidArgument( "Cardinality aggregation with missing on non-text/number field is not supported." .to_string(), ) })?; let children = build_children(sub_aggs, reader, segment_ordinal, data)?; let idx_in_req_data = data.push_missing_term_req_data(MissingTermAggReqData { accessors: all_accessors, name: agg_name.to_string(), req, }); nodes.push(AggRefNode { kind: AggKind::MissingTerm, idx_in_req_data, children, }); } // Add one node per accessor for (accessor, column_type) in column_and_types { let missing_value_for_accessor = if use_special_missing_agg { None } else if let Some(m) = missing.as_ref() { get_missing_val_as_u64_lenient(column_type, accessor.max_value(), m, field_name)? } else { None }; let children = build_children(sub_aggs, reader, segment_ordinal, data)?; let (idx, kind) = match req { TermsOrCardinalityRequest::Terms(ref req) => { let mut allowed_term_ids = None; if req.include.is_some() || req.exclude.is_some() { if column_type != ColumnType::Str { // Skip non-string columns entirely when filtering is requested. // When excluding, the behavior could be to include non-string values continue; } let str_col = str_dict_column .as_ref() .expect("str_dict_column must exist for string column"); allowed_term_ids = build_allowed_term_ids_for_str( str_col, &req.include, &req.exclude, missing.is_some(), )?; }; let idx_in_req_data = data.push_term_req_data(TermsAggReqData { accessor, column_type, str_dict_column: str_dict_column.clone(), missing_value_for_accessor, name: agg_name.to_string(), req: TermsAggregationInternal::from_req(req), sug_aggregations: sub_aggs.clone(), allowed_term_ids, is_top_level, }); (idx_in_req_data, AggKind::Terms) } TermsOrCardinalityRequest::Cardinality(ref req) => { // `str_dict_column` is computed once per field; for JSON paths // with mixed types it's `Some` even on the numeric req_data. // Cardinality only consults it for the str column path, so // gate by column_type to avoid driving non-str collectors // through the coupon-cache path. let str_dict_column_for_req = if column_type == ColumnType::Str { str_dict_column.clone() } else { None }; let idx_in_req_data = data.push_cardinality_req_data(CardinalityAggReqData { accessor, column_type, str_dict_column: str_dict_column_for_req, missing_value_for_accessor, name: agg_name.to_string(), req: req.clone(), }); (idx_in_req_data, AggKind::Cardinality) } }; nodes.push(AggRefNode { kind, idx_in_req_data: idx, children, }); } Ok(nodes) } /// Builds a single BitSet of allowed term ordinals for a string dictionary column according to /// include/exclude parameters. /// /// When `reserve_missing_sentinel` is true, the bitset will have 1 additional slot for the missing /// term ordinal fn build_allowed_term_ids_for_str( str_col: &StrColumn, include: &Option, exclude: &Option, reserve_missing_sentinel: bool, ) -> crate::Result> { let mut allowed: Option = None; let missing_sentinel_adjustment = if reserve_missing_sentinel { 1 } else { 0 }; let allowed_capacity = str_col.dictionary().num_terms() as u32 + missing_sentinel_adjustment; if let Some(include) = include { // add matches allowed = Some(BitSet::with_max_value(allowed_capacity)); let allowed = allowed.as_mut().unwrap(); for_each_matching_term_ord(str_col, include, |ord| { let _ = allowed.insert(ord); })?; }; if let Some(exclude) = exclude { if allowed.is_none() { // Start with all terms allowed allowed = Some(BitSet::with_max_value_and_full(allowed_capacity)); } let allowed = allowed.as_mut().unwrap(); for_each_matching_term_ord(str_col, exclude, |ord| allowed.remove(ord))?; } Ok(allowed) } /// Apply a callback to each matching term ordinal for the given include/exclude parameter. fn for_each_matching_term_ord( str_col: &StrColumn, param: &IncludeExcludeParam, mut cb: impl FnMut(u32), ) -> crate::Result<()> { match param { IncludeExcludeParam::Regex(pattern) => { let re = Regex::new(pattern).map_err(|e| { crate::TantivyError::InvalidArgument(format!("Invalid regex `{}`: {}", pattern, e)) })?; // TODO: we can handle patterns like `^prefix.*` more efficiently let mut stream = str_col.dictionary().search(re).into_stream()?; while stream.advance() { cb(stream.term_ord() as u32); } } IncludeExcludeParam::Values(values) => { let set: FxHashSet<&str> = values.iter().map(|s| s.as_str()).collect(); let mut stream = str_col.dictionary().stream()?; while stream.advance() { if let Ok(key_str) = std::str::from_utf8(stream.key()) { if set.contains(key_str) { cb(stream.term_ord() as u32); } } } } } Ok(()) } /// Convert the aggregation tree to something serializable and easy to read. #[derive(Serialize, Debug, Clone, PartialEq, Eq)] pub struct AggTreeViewNode { pub name: String, pub kind: String, #[serde(skip_serializing_if = "Vec::is_empty", default)] pub children: Vec, } #[cfg(test)] mod tests { use super::*; use crate::aggregation::agg_req::Aggregations; use crate::aggregation::tests::get_test_index_2_segments; fn agg_from_json(val: serde_json::Value) -> crate::aggregation::agg_req::Aggregation { serde_json::from_value(val).unwrap() } #[test] fn test_multi_terms_expands_physical_column_cartesian_product() -> crate::Result<()> { let mut schema_builder = crate::schema::Schema::builder(); let attrs = schema_builder.add_json_field("attrs", crate::schema::FAST); let score = schema_builder .add_u64_field("score", crate::schema::NumericOptions::default().set_fast()); let index = crate::Index::create_in_ram(schema_builder.build()); let mut writer = index.writer_for_tests()?; writer.add_document( crate::doc!(attrs => json!({"left": "x", "right": "y"}), score => 1u64), )?; writer.add_document( crate::doc!(attrs => json!({"left": 10.5, "right": true}), score => 2u64), )?; writer.commit()?; let agg = agg_from_json(json!({ "multi_terms": { "terms": [ {"field": "attrs.left"}, {"field": "attrs.right"} ] }, "aggs": { "sum_score": {"sum": {"field": "score"}} } })); let aggs: Aggregations = vec![("mt".to_string(), agg)].into_iter().collect(); let searcher = index.reader()?.searcher(); let data = build_aggregations_data_from_req( &aggs, searcher.segment_reader(0), 0, Default::default(), )?; assert_eq!(data.per_request.agg_tree.len(), 4); assert_eq!(data.per_request.multi_terms_req_data.len(), 4); assert!(data .per_request .agg_tree .iter() .all(|node| node.children.len() == 1)); let actual_types: FxHashSet> = data .per_request .multi_terms_req_data .iter() .map(|req_data| { req_data .fields .iter() .map(|field| field.column_type) .collect() }) .collect(); let expected_types: FxHashSet> = [ vec![ColumnType::F64, ColumnType::Bool], vec![ColumnType::F64, ColumnType::Str], vec![ColumnType::Str, ColumnType::Bool], vec![ColumnType::Str, ColumnType::Str], ] .into_iter() .collect(); assert_eq!(actual_types, expected_types); assert!(data .per_request .multi_terms_req_data .iter() .flat_map(|req_data| &req_data.fields) .all(|field| { field.str_dict_column.is_some() == (field.column_type == ColumnType::Str) })); Ok(()) } #[test] fn test_multi_terms_skips_missing_when_any_physical_column_is_full() -> crate::Result<()> { let mut schema_builder = crate::schema::Schema::builder(); let attrs = schema_builder.add_json_field("attrs", crate::schema::FAST); let index = crate::Index::create_in_ram(schema_builder.build()); let mut writer = index.writer_for_tests()?; writer.add_document(crate::doc!(attrs => json!({"value": ["a", 10.5]})))?; writer.add_document(crate::doc!(attrs => json!({"value": "b"})))?; writer.commit()?; let agg = agg_from_json(json!({ "multi_terms": { "terms": [{"field": "attrs.value", "missing": "MISSING"}] } })); let aggs: Aggregations = vec![("mt".to_string(), agg)].into_iter().collect(); let searcher = index.reader()?.searcher(); let data = build_aggregations_data_from_req( &aggs, searcher.segment_reader(0), 0, Default::default(), )?; assert_eq!(data.per_request.multi_terms_req_data.len(), 2); assert!(data .per_request .multi_terms_req_data .iter() .any(|req_data| req_data.fields[0].column.get_cardinality().is_full())); assert!(data .per_request .multi_terms_req_data .iter() .all(|req_data| req_data.missing_accessors.iter().all(Option::is_none))); Ok(()) } #[test] fn test_tree_roots_and_expansion_terms_missing_on_numeric() -> crate::Result<()> { let index = get_test_index_2_segments(true)?; let reader = index.reader()?; let searcher = reader.searcher(); let seg_reader = searcher.segment_reader(0u32); // Build request with: // 1) Terms on numeric field with missing as string => expands to MissingTerm + Terms // 2) Avg metric // 3) Terms on string with child histogram let terms_score_missing = agg_from_json(json!({ "terms": {"field": "score", "missing": "NA"} })); let avg_score = agg_from_json(json!({ "avg": {"field": "score"} })); let terms_string_with_child = agg_from_json(json!({ "terms": {"field": "string_id"}, "aggs": { "histo": {"histogram": {"field": "score", "interval": 10.0}} } })); let aggs: Aggregations = vec![ ("t_score_missing_str".to_string(), terms_score_missing), ("avg_score".to_string(), avg_score), ("terms_string".to_string(), terms_string_with_child), ] .into_iter() .collect(); let data = build_aggregations_data_from_req(&aggs, seg_reader, 0u32, Default::default())?; let printed_nodes = data.per_request.get_view_tree(); let printed = serde_json::to_value(&printed_nodes).unwrap(); let expected = json!([ {"name": "avg_score", "kind": "Metric"}, {"name": "t_score_missing_str", "kind": "MissingTerm"}, {"name": "t_score_missing_str", "kind": "Terms"}, {"name": "terms_string", "kind": "Terms", "children": [ {"name": "histo", "kind": "Histogram"} ]} ]); assert_eq!( printed, expected, "tree json:\n{}", serde_json::to_string_pretty(&printed).unwrap() ); Ok(()) } }