repo_description: Tantivy is a full-text search engine library inspired by Apache Lucene and written in Rust
repo_topics: rust search-engine
code: use columnar::{Column, ColumnBlockAccessor, 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_range_collector, CompositeAggReqData,
      CompositeAggregation, CompositeSourceAccessors, FilterAggReqData, HistogramAggReqData,
      HistogramBounds, IncludeExcludeParam, MissingTermAggReqData, RangeAggReqData,
      SegmentHistogramCollector, 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, TopHitsAggReqData,
      TopHitsSegmentCollector,
  };
  use crate::aggregation::segment_agg_result::{
      GenericSegmentAggregationResultsCollector, SegmentAggregationCollector,
  };
  use crate::aggregation::{f64_to_fastfield_u64, AggContextParams, 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 column_block_accessor: ColumnBlockAccessor<u64>,
  }
  
  impl AggregationsSegmentCtx {
      pub(crate) fn push_term_req_data(&mut self, data: TermsAggReqData) -> usize {
          self.per_request.term_req_data.push(Some(Box::new(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(Some(Box::new(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(Some(Box::new(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(Some(Box::new(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(Some(Box::new(data)));
          self.per_request.composite_req_data.len() - 1
      }
  
      #[inline]
      pub(crate) fn get_term_req_data(&self, idx: usize) -> &TermsAggReqData {
          self.per_request.term_req_data[idx]
              .as_deref()
              .expect("term_req_data slot is empty (taken)")
      }
      #[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]
      }
      #[inline]
      pub(crate) fn get_histogram_req_data(&self, idx: usize) -> &HistogramAggReqData {
          self.per_request.histogram_req_data[idx]
              .as_deref()
              .expect("histogram_req_data slot is empty (taken)")
      }
      #[inline]
      pub(crate) fn get_range_req_data(&self, idx: usize) -> &RangeAggReqData {
          self.per_request.range_req_data[idx]
              .as_deref()
              .expect("range_req_data slot is empty (taken)")
      }
      #[inline]
      pub(crate) fn get_composite_req_data(&self, idx: usize) -> &CompositeAggReqData {
          self.per_request.composite_req_data[idx]
              .as_deref()
              .expect("composite_req_data slot is empty (taken)")
      }
  
      // ---------- mutable getters ----------
  
      #[inline]
      pub(crate) fn get_metric_req_data_mut(&mut self, idx: usize) -> &mut MetricAggReqData {
          &mut self.per_request.stats_metric_req_data[idx]
      }
  
      #[inline]
      pub(crate) fn get_cardinality_req_data_mut(
          &mut self,
          idx: usize,
      ) -> &mut CardinalityAggReqData {
          &mut self.per_request.cardinality_req_data[idx]
      }
  
      #[inline]
      pub(crate) fn get_histogram_req_data_mut(&mut self, idx: usize) -> &mut HistogramAggReqData {
          self.per_request.histogram_req_data[idx]
              .as_deref_mut()
              .expect("histogram_req_data slot is empty (taken)")
      }
  
      // ---------- take / put (terms, histogram, range) ----------
  
      /// Move out the boxed Histogram request at `idx`, leaving `None`.
      #[inline]
      pub(crate) fn take_histogram_req_data(&mut self, idx: usize) -> Box<HistogramAggReqData> {
          self.per_request.histogram_req_data[idx]
              .take()
              .expect("histogram_req_data slot is empty (taken)")
      }
  
      /// Put back a Histogram request into an empty slot at `idx`.
      #[inline]
      pub(crate) fn put_back_histogram_req_data(
          &mut self,
          idx: usize,
          value: Box<HistogramAggReqData>,
      ) {
          debug_assert!(self.per_request.histogram_req_data[idx].is_none());
          self.per_request.histogram_req_data[idx] = Some(value);
      }
  
      /// Move out the boxed Range request at `idx`, leaving `None`.
      #[inline]
      pub(crate) fn take_range_req_data(&mut self, idx: usize) -> Box<RangeAggReqData> {
          self.per_request.range_req_data[idx]
              .take()
              .expect("range_req_data slot is empty (taken)")
      }
  
      /// Put back a Range request into an empty slot at `idx`.
      #[inline]
      pub(crate) fn put_back_range_req_data(&mut self, idx: usize, value: Box<RangeAggReqData>) {
          debug_assert!(self.per_request.range_req_data[idx].is_none());
          self.per_request.range_req_data[idx] = Some(value);
      }
  
      /// Move out the boxed Filter request at `idx`, leaving `None`.
      #[inline]
      pub(crate) fn take_filter_req_data(&mut self, idx: usize) -> Box<FilterAggReqData> {
          self.per_request.filter_req_data[idx]
              .take()
              .expect("filter_req_data slot is empty (taken)")
      }
  
      /// Put back a Filter request into an empty slot at `idx`.
      #[inline]
      pub(crate) fn put_back_filter_req_data(&mut self, idx: usize, value: Box<FilterAggReqData>) {
          debug_assert!(self.per_request.filter_req_data[idx].is_none());
          self.per_request.filter_req_data[idx] = Some(value);
      }
  
      /// Move out the Composite request at `idx`.
      #[inline]
      pub(crate) fn take_composite_req_data(&mut self, idx: usize) -> Box<CompositeAggReqData> {
          self.per_request.composite_req_data[idx]
              .take()
              .expect("composite_req_data slot is empty (taken)")
      }
  
      /// Put back a Composite request into an empty slot at `idx`.
      #[inline]
      pub(crate) fn put_back_composite_req_data(
          &mut self,
          idx: usize,
          value: Box<CompositeAggReqData>,
      ) {
          debug_assert!(self.per_request.composite_req_data[idx].is_none());
          self.per_request.composite_req_data[idx] = Some(value);
      }
  }
  
  /// 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 {
      // Box for cheap take/put - Only necessary for bucket aggs that have sub-aggregations
      /// TermsAggReqData contains the request data for a terms aggregation.
      pub term_req_data: Vec<Option<Box<TermsAggReqData>>>,
      /// HistogramAggReqData contains the request data for a histogram aggregation.
      pub histogram_req_data: Vec<Option<Box<HistogramAggReqData>>>,
      /// RangeAggReqData contains the request data for a range aggregation.
      pub range_req_data: Vec<Option<Box<RangeAggReqData>>>,
      /// FilterAggReqData contains the request data for a filter aggregation.
      pub filter_req_data: Vec<Option<Box<FilterAggReqData>>>,
      /// Shared by avg, min, max, sum, stats, extended_stats, count
      pub stats_metric_req_data: Vec<MetricAggReqData>,
      /// CardinalityAggReqData contains the request data for a cardinality aggregation.
      pub cardinality_req_data: Vec<CardinalityAggReqData>,
      /// TopHitsAggReqData contains the request data for a top_hits aggregation.
      pub top_hits_req_data: Vec<TopHitsAggReqData>,
      /// MissingTermAggReqData contains the request data for a missing term aggregation.
      pub missing_term_req_data: Vec<MissingTermAggReqData>,
      /// CompositeAggReqData contains the request data for a composite aggregation.
      pub composite_req_data: Vec<Option<Box<CompositeAggReqData>>>,
  
      /// Request tree used to build collectors.
      pub agg_tree: Vec<AggRefNode>,
  }
  
  impl PerRequestAggSegCtx {
      /// Estimate the memory consumption of this struct in bytes.
      fn get_memory_consumption(&self) -> usize {
          self.term_req_data
              .iter()
              .map(|b| b.as_ref().unwrap().get_memory_consumption())
              .sum::<usize>()
              + self
                  .histogram_req_data
                  .iter()
                  .map(|b| b.as_ref().unwrap().get_memory_consumption())
                  .sum::<usize>()
              + self
                  .range_req_data
                  .iter()
                  .map(|b| b.as_ref().unwrap().get_memory_consumption())
                  .sum::<usize>()
              + self
                  .filter_req_data
                  .iter()
                  .map(|b| b.as_ref().unwrap().get_memory_consumption())
                  .sum::<usize>()
              + self
                  .stats_metric_req_data
                  .iter()
                  .map(|t| t.get_memory_consumption())
                  .sum::<usize>()
              + self
                  .cardinality_req_data
                  .iter()
                  .map(|t| t.get_memory_consumption())
                  .sum::<usize>()
              + self
                  .top_hits_req_data
                  .iter()
                  .map(|t| t.get_memory_consumption())
                  .sum::<usize>()
              + self
                  .missing_term_req_data
                  .iter()
                  .map(|t| t.get_memory_consumption())
                  .sum::<usize>()
              + self
                  .composite_req_data
                  .iter()
                  .map(|b| b.as_ref().map(|d| d.get_memory_consumption()).unwrap_or(0))
                  .sum::<usize>()
              + self.agg_tree.len() * std::mem::size_of::<AggRefNode>()
      }
  
      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]
                  .as_deref()
                  .expect("term_req_data slot is empty (taken)")
                  .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]
                  .as_deref()
                  .expect("histogram_req_data slot is empty (taken)")
                  .name
                  .as_str(),
              AggKind::DateHistogram => self.histogram_req_data[idx]
                  .as_deref()
                  .expect("histogram_req_data slot is empty (taken)")
                  .name
                  .as_str(),
              AggKind::Range => self.range_req_data[idx]
                  .as_deref()
                  .expect("range_req_data slot is empty (taken)")
                  .name
                  .as_str(),
              AggKind::Filter => self.filter_req_data[idx]
                  .as_deref()
                  .expect("filter_req_data slot is empty (taken)")
                  .name
                  .as_str(),
              AggKind::Composite => self.composite_req_data[idx]
                  .as_deref()
                  .expect("composite_req_data slot is empty (taken)")
                  .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<AggTreeViewNode> {
          fn node_to_view(node: &AggRefNode, pr: &PerRequestAggSegCtx) -> AggTreeViewNode {
              let mut children: Vec<AggTreeViewNode> =
                  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<AggTreeViewNode> = 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<Box<dyn SegmentAggregationCollector>> {
      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<Box<dyn SegmentAggregationCollector>> {
      build_segment_agg_collectors_generic(req, nodes)
  }
  
  fn build_segment_agg_collectors_generic(
      req: &mut AggregationsSegmentCtx,
      nodes: &[AggRefNode],
  ) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
      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<Box<dyn SegmentAggregationCollector>> {
      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 = &mut req.get_cardinality_req_data_mut(node.idx_in_req_data);
              Ok(Box::new(SegmentCardinalityCollector::from_req(
                  req_data.column_type,
                  node.idx_in_req_data,
                  req_data.accessor.clone(),
                  req_data.missing_value_for_accessor,
              )))
          }
          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_mut(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 => Ok(Box::new(SegmentHistogramCollector::from_req_and_validate(
              req, node,
          )?)),
          AggKind::DateHistogram => Ok(Box::new(SegmentHistogramCollector::from_req_and_validate(
              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,
              )?,
          )),
      }
  }
  
  /// See [PerRequestAggSegCtx]
  #[derive(Debug, Clone)]
  pub struct AggRefNode {
      pub kind: AggKind,
      pub idx_in_req_data: usize,
      pub children: Vec<AggRefNode>,
  }
  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,
  }
  
  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",
          }
      }
  }
  
  /// 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<AggregationsSegmentCtx> {
      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<Vec<AggRefNode>> {
      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<u64>, ColumnType)> = top_hits
                  .field_names()
                  .iter()
                  .map(|field| get_ff_reader(reader, field, Some(get_numeric_or_date_column_types())))
                  .collect::<crate::Result<_>>()?;
  
              let value_accessors = top_hits
                  .value_field_names()
                  .iter()
                  .map(|field_name| {
                      Ok((
                          field_name.to_string(),
                          get_dynamic_columns(reader, field_name)?,
                      ))
                  })
                  .collect::<crate::Result<_>>()?;
  
              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::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 = crate::aggregation::bucket::DocumentQueryEvaluator::new(
                  query,
                  schema.clone(),
                  reader,
              )?;
  
              // Pre-allocate buffer for batch filtering
              let max_doc = reader.max_doc();
              let buffer_capacity = crate::docset::COLLECT_BLOCK_BUFFER_LEN.min(max_doc as usize);
              let matching_docs_buffer = Vec::with_capacity(buffer_capacity);
  
              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,
                  matching_docs_buffer,
                  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<AggRefNode> {
      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_children(
      aggs: &Aggregations,
      reader: &SegmentReader,
      segment_ordinal: SegmentOrdinal,
      data: &mut AggregationsSegmentCtx,
  ) -> crate::Result<Vec<AggRefNode>> {
      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<Key>,
  ) -> crate::Result<Vec<(Column<u64>, ColumnType)>> {
      let allowed_column_types = [
          ColumnType::I64,
          ColumnType::U64,
          ColumnType::F64,
          ColumnType::Str,
          ColumnType::DateTime,
          ColumnType::Bool,
          ColumnType::IpAddr,
      ];
  
      // 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<Key>,
      reader: &SegmentReader,
      segment_ordinal: SegmentOrdinal,
      data: &mut AggregationsSegmentCtx,
      sub_aggs: &Aggregations,
      req: TermsOrCardinalityRequest,
      is_top_level: bool,
  ) -> crate::Result<Vec<AggRefNode>> {
      let mut nodes = Vec::new();
  
      let str_dict_column = reader.fast_fields().str(field_name)?;
  
      let column_and_types = get_term_agg_accessors(reader, field_name, missing)?;
  
      // 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::<Vec<_>>();
          // 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)?;
                  };
                  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) => {
                  let idx_in_req_data = data.push_cardinality_req_data(CardinalityAggReqData {
                      accessor,
                      column_type,
                      str_dict_column: str_dict_column.clone(),
                      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.
  fn build_allowed_term_ids_for_str(
      str_col: &StrColumn,
      include: &Option<IncludeExcludeParam>,
      exclude: &Option<IncludeExcludeParam>,
  ) -> crate::Result<Option<BitSet>> {
      let mut allowed: Option<BitSet> = None;
      let num_terms = str_col.dictionary().num_terms() as u32;
      if let Some(include) = include {
          // add matches
          allowed = Some(BitSet::with_max_value(num_terms));
          let allowed = allowed.as_mut().unwrap();
          for_each_matching_term_ord(str_col, include, |ord| 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(num_terms));
          }
          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<AggTreeViewNode>,
  }
  
  #[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_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(())
      }
  }
  