diff --git a/src/aggregation/accessor_helpers.rs b/src/aggregation/accessor_helpers.rs index 9d050d1b9..7668cba55 100644 --- a/src/aggregation/accessor_helpers.rs +++ b/src/aggregation/accessor_helpers.rs @@ -1,7 +1,6 @@ //! This will enhance the request tree with access to the fastfield and metadata. use std::io; -use std::sync::Arc; use columnar::{Column, ColumnType, DynamicColumn, DynamicColumnHandle}; @@ -62,7 +61,7 @@ fn resolve_registered_source( value_sources: &ValueSourceRegistry, field_name: &str, allowed_column_types_opt: Option<&[ColumnType]>, -) -> crate::Result>> { +) -> crate::Result>> { let Some(provider) = value_sources.get(field_name) else { return Ok(None); }; @@ -81,7 +80,7 @@ pub(crate) fn get_value_source( value_sources: &ValueSourceRegistry, field_name: &str, allowed_column_types: Option<&[ColumnType]>, -) -> crate::Result> { +) -> crate::Result> { if let Some(registered) = resolve_registered_source(reader, value_sources, field_name, allowed_column_types)? { @@ -98,7 +97,7 @@ pub(crate) fn get_value_source( }); // The empty-column shim stays physical on purpose: several fast paths check // `as_column()` and would otherwise degrade for a merely absent field. - Ok(Arc::new((column, column_type))) + Ok(Box::new((column, column_type))) } pub(crate) fn get_dynamic_columns( @@ -124,7 +123,7 @@ pub(crate) fn get_all_value_sources( field_name: &str, allowed_column_types: Option<&[ColumnType]>, fallback_type: ColumnType, -) -> crate::Result>> { +) -> crate::Result>> { // A registered source shadows the physical type fan-out entirely. if let Some(registered) = resolve_registered_source(reader, value_sources, field_name, allowed_column_types)? @@ -140,7 +139,7 @@ pub(crate) fn get_all_value_sources( Ok(ff_field_with_type .into_iter() .map(|(column, column_type)| { - let source: Arc = Arc::new((column, column_type)); + let source: Box = Box::new((column, column_type)); source }) .collect()) diff --git a/src/aggregation/agg_data.rs b/src/aggregation/agg_data.rs index 0213e002e..6dae08c26 100644 --- a/src/aggregation/agg_data.rs +++ b/src/aggregation/agg_data.rs @@ -15,9 +15,9 @@ 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, + IncludeExcludeParam, MissingTermAggReqData, MultiTermsAggReqData, MultiTermsAggregation, + MultiTermsFieldAccessor, MultiTermsMissingAccessor, RangeAggReqData, TermMissingAgg, + TermsAggReqData, TermsAggregation, TermsAggregationInternal, }; use crate::aggregation::metric::{ build_segment_stats_collector, AverageAggregation, CardinalityAggReqData, @@ -34,236 +34,152 @@ use crate::aggregation::{ }; 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. +/// Shared state passed to the collectors during collection on a segment. pub struct AggregationsSegmentCtx { - /// Request data for each aggregation type. - pub(crate) per_request: PerRequestAggSegCtx, pub(crate) context: AggContextParams, + /// Scratch buffers shared by all of the collectors of the tree, to load a block of values. 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] + pub(crate) fn new(context: AggContextParams) -> Self { + AggregationsSegmentCtx { + context, + column_block_accessor: ColumnBlockAccessor::default(), + } } } -/// 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. +/// A node of the per-segment aggregation request tree. /// -/// 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(crate) term_req_data: Vec, - /// HistogramAggReqData contains the request data for a histogram aggregation. - pub(crate) histogram_req_data: Vec, - /// RangeAggReqData contains the request data for a range aggregation. - pub(crate) range_req_data: Vec, - /// FilterAggReqData contains the request data for a filter aggregation. - pub(crate) filter_req_data: Vec, - /// Shared by avg, min, max, sum, stats, extended_stats, count - pub(crate) stats_metric_req_data: Vec, - /// CardinalityAggReqData contains the request data for a cardinality aggregation. - pub(crate) cardinality_req_data: Vec, - /// TopHitsAggReqData contains the request data for a top_hits aggregation. - pub(crate) top_hits_req_data: Vec, - /// MissingTermAggReqData contains the request data for a missing term aggregation. - pub(crate) missing_term_req_data: Vec, - /// CompositeAggReqData contains the request data for a composite aggregation. - pub(crate) composite_req_data: Vec, - /// MultiTermsAggReqData contains the request data for a multi_terms aggregation. - pub(crate) multi_terms_req_data: Vec, - - /// Request tree used to build collectors. - pub(crate) agg_tree: Vec, +/// A node owns the request data of its aggregation, including its value source. Building the +/// collectors consumes the tree: each node is turned into exactly one collector, which takes +/// ownership of the node's request data. +/// +/// A single aggregation of the request can expand into several sibling nodes (e.g. a terms +/// aggregation over a JSON field with several column types). +pub(crate) struct AggNode { + pub(crate) data: AggNodeData, + pub(crate) children: 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::() +/// The request data of an [`AggNode`]. Each type of aggregation has its own request data +/// struct. +pub(crate) enum AggNodeData { + Terms(TermsAggReqData), + MissingTerm(MissingTermAggReqData), + Cardinality(CardinalityAggReqData), + /// Shared by avg, min, max, sum, stats, extended_stats, count and percentiles. + Metric(MetricAggReqData), + TopHits(Box), + /// Shared by histogram and date_histogram. + Histogram(HistogramAggReqData), + Range(RangeAggReqData), + Filter(Box), + Composite(CompositeAggReqData), + MultiTerms(MultiTermsAggReqData), +} + +impl AggNode { + fn new(data: AggNodeData, children: Vec) -> Self { + AggNode { data, children } } - 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(), + /// Name of the aggregation, as given in the request. + pub(crate) fn name(&self) -> &str { + match &self.data { + AggNodeData::Terms(req_data) => &req_data.name, + AggNodeData::MissingTerm(req_data) => &req_data.name, + AggNodeData::Cardinality(req_data) => &req_data.name, + AggNodeData::Metric(req_data) => &req_data.name, + AggNodeData::TopHits(req_data) => &req_data.name, + AggNodeData::Histogram(req_data) => &req_data.name, + AggNodeData::Range(req_data) => &req_data.name, + AggNodeData::Filter(req_data) => &req_data.name, + AggNodeData::Composite(req_data) => &req_data.name, + AggNodeData::MultiTerms(req_data) => &req_data.name, } } - /// 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, - } + #[cfg_attr(not(test), allow(dead_code))] + fn kind_name(&self) -> &'static str { + match &self.data { + AggNodeData::Terms(_) => "Terms", + AggNodeData::MissingTerm(_) => "MissingTerm", + AggNodeData::Cardinality(_) => "Cardinality", + AggNodeData::Metric(_) => "Metric", + AggNodeData::TopHits(_) => "TopHits", + AggNodeData::Histogram(req_data) if req_data.is_date_histogram => "DateHistogram", + AggNodeData::Histogram(_) => "Histogram", + AggNodeData::Range(_) => "Range", + AggNodeData::Filter(_) => "Filter", + AggNodeData::Composite(_) => "Composite", + AggNodeData::MultiTerms(_) => "MultiTerms", } + } - 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 + /// Estimate the memory consumption of this node's request data in bytes, excluding its + /// children. + /// + /// The node itself is not counted: it is consumed when building its collector, which only + /// keeps the request data. + pub(crate) fn get_memory_consumption(&self) -> usize { + match &self.data { + AggNodeData::Terms(req_data) => req_data.get_memory_consumption(), + AggNodeData::MissingTerm(req_data) => req_data.get_memory_consumption(), + AggNodeData::Cardinality(req_data) => req_data.get_memory_consumption(), + AggNodeData::Metric(req_data) => req_data.get_memory_consumption(), + AggNodeData::TopHits(req_data) => req_data.get_memory_consumption(), + AggNodeData::Histogram(req_data) => req_data.get_memory_consumption(), + AggNodeData::Range(req_data) => req_data.get_memory_consumption(), + AggNodeData::Filter(req_data) => req_data.get_memory_consumption(), + AggNodeData::Composite(req_data) => req_data.get_memory_consumption(), + AggNodeData::MultiTerms(req_data) => req_data.get_memory_consumption(), + } } } -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()) +/// Returns the child aggregation named `name`, if any. +pub(crate) fn find_sub_agg<'a>(children: &'a [AggNode], name: &str) -> Option<&'a AggNode> { + children.iter().find(|child| child.name() == name) } +/// Convert the aggregation tree into a serializable struct representation. +/// Each node contains: { name, kind, children }. +#[cfg_attr(not(test), allow(dead_code))] +pub(crate) fn get_view_tree(nodes: &[AggNode]) -> Vec { + let mut views: Vec = nodes + .iter() + .map(|node| AggTreeViewNode { + name: node.name().to_string(), + kind: node.kind_name().to_string(), + children: get_view_tree(&node.children), + }) + .collect(); + views.sort_by_key(|view| serde_json::to_string(view).unwrap()); + views +} + +/// Builds the collectors for `nodes`, consuming them. +/// +/// This is where the request data of `nodes` is charged to the memory limits. Hidden contract: +/// collector builders must not charge their request data again, and a builder that consumes a +/// child node without calling this function (e.g. the flattened terms×histogram collector) has to +/// charge that child itself. pub(crate) fn build_segment_agg_collectors( - req: &mut AggregationsSegmentCtx, - nodes: &[AggRefNode], + ctx: &mut AggregationsSegmentCtx, + nodes: Vec, ) -> 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)?); + // The request data is moved into the collectors, so we need to measure it beforehand. + let memory_consumption: usize = nodes.iter().map(AggNode::get_memory_consumption).sum(); + let mut collectors = Vec::with_capacity(nodes.len()); + for node in nodes { + collectors.push(build_segment_agg_collector(ctx, node)?); } - req.context + ctx.context .limits - .add_memory_consumed(req.per_request.get_memory_consumption() as u64)?; + .add_memory_consumed(memory_consumption as u64)?; // Single collector special case if collectors.len() == 1 { return Ok(collectors.pop().unwrap()); @@ -272,137 +188,92 @@ fn build_segment_agg_collectors_generic( Ok(Box::new(agg)) } +/// Builds the sub-aggregation collectors of a bucket aggregation, consuming `children`. +/// +/// Returns `None` when there are no children: collectors rely on `None` to skip doc buffering. +pub(crate) fn build_sub_agg_collectors( + ctx: &mut AggregationsSegmentCtx, + children: Vec, +) -> crate::Result>> { + if children.is_empty() { + return Ok(None); + } + Ok(Some(build_segment_agg_collectors(ctx, children)?)) +} + +/// Builds the collector for `node`, consuming it. pub(crate) fn build_segment_agg_collector( - req: &mut AggregationsSegmentCtx, - node: &AggRefNode, + ctx: &mut AggregationsSegmentCtx, + node: AggNode, ) -> 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]; + let AggNode { data, children } = node; + match data { + AggNodeData::Terms(req_data) => { + crate::aggregation::bucket::build_segment_term_collector(ctx, req_data, children) + } + AggNodeData::MissingTerm(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)?)) + let sub_agg = build_sub_agg_collectors(ctx, children)?; + Ok(Box::new(TermMissingAgg::new(req_data, sub_agg))) } - AggKind::Cardinality => { - crate::aggregation::metric::build_segment_cardinality_collector(req, node) + AggNodeData::Cardinality(req_data) => { + crate::aggregation::metric::build_segment_cardinality_collector(req_data) } - 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.missing_u64, - req_data.accessor.clone(), - node.idx_in_req_data, - ), - )) - } - } + AggNodeData::Metric(req_data) => match req_data.collecting_for { + 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 => Ok(Box::new( + SegmentPercentilesCollector::from_req_and_validate(req_data), + )), + }, + AggNodeData::TopHits(req_data) => { + Ok(Box::new(TopHitsSegmentCollector::from_req(*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, - ))) + AggNodeData::Histogram(req_data) => { + build_segment_histogram_collector(ctx, req_data, children) } - 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", + AggNodeData::Range(req_data) => build_segment_range_collector(ctx, req_data, children), + AggNodeData::Filter(req_data) => build_segment_filter_collector(ctx, *req_data, children), + AggNodeData::Composite(req_data) => { + let sub_agg = build_sub_agg_collectors(ctx, children)?; + Ok(Box::new( + crate::aggregation::bucket::SegmentCompositeCollector::from_req_and_validate( + req_data, sub_agg, + )?, + )) + } + AggNodeData::MultiTerms(req_data) => { + build_segment_multi_terms_collector(ctx, req_data, children) } } } -/// Build AggregationsData by walking the request tree. +/// Builds the aggregation request tree for a segment. +/// +/// This resolves the value sources of every aggregation and validates field types. 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(), - }; - + context: &AggContextParams, +) -> crate::Result> { + let mut agg_tree = Vec::with_capacity(aggs.len()); 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); + let nodes = build_nodes(name, agg, reader, segment_ordinal, context, true)?; + agg_tree.extend(nodes); } - Ok(data) + Ok(agg_tree) } /// Resolves the substitute value used for documents that have none. @@ -450,11 +321,11 @@ fn build_nodes( req: &Aggregation, reader: &SegmentReader, segment_ordinal: SegmentOrdinal, - data: &mut AggregationsSegmentCtx, + context: &AggContextParams, is_top_level: bool, -) -> crate::Result> { +) -> crate::Result> { use AggregationVariants::*; - let value_sources = &data.context.value_sources; + let value_sources = &context.value_sources; match &req.agg { Range(range_req) => { let accessor = get_value_source( @@ -463,18 +334,14 @@ fn build_nodes( &range_req.field, Some(get_numeric_or_date_column_types()), )?; - let idx_in_req_data = data.push_range_req_data(RangeAggReqData { + let node_data = AggNodeData::Range(RangeAggReqData { accessor, 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, - }]) + let children = build_children(&req.sub_aggregation, reader, segment_ordinal, context)?; + Ok(vec![AggNode::new(node_data, children)]) } Histogram(histo_req) => { let accessor = get_value_source( @@ -483,23 +350,11 @@ fn build_nodes( &histo_req.field, Some(get_numeric_or_date_column_types()), )?; - let idx_in_req_data = data.push_histogram_req_data(HistogramAggReqData { - accessor, - 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, - }]) + let req_data = + HistogramAggReqData::new(accessor, agg_name.to_string(), histo_req.clone(), false)?; + let node_data = AggNodeData::Histogram(req_data); + let children = build_children(&req.sub_aggregation, reader, segment_ordinal, context)?; + Ok(vec![AggNode::new(node_data, children)]) } DateHistogram(date_req) => { let accessor = get_value_source( @@ -511,23 +366,11 @@ fn build_nodes( // 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, - 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, - }]) + let req_data = + HistogramAggReqData::new(accessor, agg_name.to_string(), histo_req, true)?; + let node_data = AggNodeData::Histogram(req_data); + let children = build_children(&req.sub_aggregation, reader, segment_ordinal, context)?; + Ok(vec![AggNode::new(node_data, children)]) } Terms(terms_req) => build_terms_or_cardinality_nodes( agg_name, @@ -535,7 +378,7 @@ fn build_nodes( &terms_req.missing, reader, segment_ordinal, - data, + context, &req.sub_aggregation, TermsOrCardinalityRequest::Terms(terms_req.clone()), is_top_level, @@ -546,7 +389,7 @@ fn build_nodes( &card_req.missing, reader, segment_ordinal, - data, + context, &req.sub_aggregation, TermsOrCardinalityRequest::Cardinality(card_req.clone()), is_top_level, @@ -591,7 +434,7 @@ fn build_nodes( }; let accessor = get_value_source(reader, value_sources, field, allowed_column_types)?; let field_type = accessor.column_type(); - let idx_in_req_data = data.push_metric_req_data(MetricAggReqData { + let node_data = AggNodeData::Metric(MetricAggReqData { accessor, name: agg_name.to_string(), collecting_for, @@ -602,12 +445,8 @@ fn build_nodes( 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, - }]) + let children = build_children(&req.sub_aggregation, reader, segment_ordinal, context)?; + Ok(vec![AggNode::new(node_data, children)]) } // Percentiles handled as Metric as well AggregationVariants::Percentiles(percentiles_req) => { @@ -619,7 +458,7 @@ fn build_nodes( Some(get_numeric_or_date_column_types()), )?; let field_type = accessor.column_type(); - let idx_in_req_data = data.push_metric_req_data(MetricAggReqData { + let node_data = AggNodeData::Metric(MetricAggReqData { accessor, name: agg_name.to_string(), collecting_for: StatsType::Percentiles, @@ -632,12 +471,8 @@ fn build_nodes( 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, - }]) + let children = build_children(&req.sub_aggregation, reader, segment_ordinal, context)?; + Ok(vec![AggNode::new(node_data, children)]) } AggregationVariants::TopHits(top_hits_req) => { let mut top_hits = top_hits_req.clone(); @@ -670,25 +505,21 @@ fn build_nodes( }) .collect::>()?; - let idx_in_req_data = data.push_top_hits_req_data(TopHitsAggReqData { + let node_data = AggNodeData::TopHits(Box::new(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, - }]) + })); + let children = build_children(&req.sub_aggregation, reader, segment_ordinal, context)?; + Ok(vec![AggNode::new(node_data, children)]) } AggregationVariants::Composite(composite_req) => Ok(vec![build_composite_node( agg_name, reader, segment_ordinal, - data, + context, &req.sub_aggregation, composite_req, )?]), @@ -696,7 +527,7 @@ fn build_nodes( agg_name, reader, segment_ordinal, - data, + context, &req.sub_aggregation, multi_terms_req, is_top_level, @@ -704,27 +535,22 @@ fn build_nodes( AggregationVariants::Filter(filter_req) => { // Build the query and evaluator upfront let schema = reader.schema(); - let tokenizers = &data.context.tokenizers; + let tokenizers = &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 evaluator = crate::aggregation::bucket::DocumentQueryEvaluator::new( + query, + schema.clone(), + reader, + )?; - let idx_in_req_data = data.push_filter_req_data(FilterAggReqData { + let node_data = AggNodeData::Filter(Box::new(FilterAggReqData { name: agg_name.to_string(), 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, - }]) + })); + let children = build_children(&req.sub_aggregation, reader, segment_ordinal, context)?; + Ok(vec![AggNode::new(node_data, children)]) } } } @@ -733,10 +559,10 @@ fn build_composite_node( agg_name: &str, reader: &SegmentReader, _segment_ordinal: SegmentOrdinal, - data: &mut AggregationsSegmentCtx, + context: &AggContextParams, sub_aggs: &Aggregations, req: &CompositeAggregation, -) -> crate::Result { +) -> 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); @@ -749,31 +575,26 @@ fn build_composite_node( 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, - }) + let children = build_children(sub_aggs, reader, _segment_ordinal, context)?; + Ok(AggNode::new(AggNodeData::Composite(agg), children)) } fn build_multi_terms_nodes( agg_name: &str, reader: &SegmentReader, segment_ordinal: SegmentOrdinal, - data: &mut AggregationsSegmentCtx, + context: &AggContextParams, sub_aggs: &Aggregations, req: &MultiTermsAggregation, is_top_level: bool, -) -> crate::Result> { +) -> crate::Result> { if req.terms.is_empty() { return Err(crate::TantivyError::InvalidArgument( "multi_terms aggregation requires at least one field".to_string(), )); } - let value_sources = data.context.value_sources.clone(); + let value_sources = &context.value_sources; let mut accessors_by_field = Vec::with_capacity(req.terms.len()); for field_def in &req.terms { let field_name = &field_def.field; @@ -781,7 +602,7 @@ fn build_multi_terms_nodes( // multi_terms resolves missing values through `ColumnIndex::has_value` per document, and // exposes its columns on a public struct, so it stays physical-only. let columns = - get_term_agg_accessors(reader, &value_sources, field_name, &field_def.missing, true)? + get_term_agg_accessors(reader, value_sources, field_name, &field_def.missing, true)? .into_iter() .map(|source| { let column = require_physical_column(&*source, field_name, "multi_terms")?; @@ -856,19 +677,15 @@ fn build_multi_terms_nodes( 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 { + let node_data = AggNodeData::MultiTerms(MultiTermsAggReqData { name: agg_name.to_string(), req: req.clone(), fields, missing_accessors, 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, - }); + let children = build_children(sub_aggs, reader, segment_ordinal, context)?; + nodes.push(AggNode::new(node_data, children)); } Ok(nodes) } @@ -983,8 +800,8 @@ fn build_children( aggs: &Aggregations, reader: &SegmentReader, segment_ordinal: SegmentOrdinal, - data: &mut AggregationsSegmentCtx, -) -> crate::Result> { + context: &AggContextParams, +) -> crate::Result> { let mut children = Vec::new(); for (name, agg) in aggs.iter() { children.extend(build_nodes( @@ -992,7 +809,7 @@ fn build_children( agg, reader, segment_ordinal, - data, + context, false, )?); } @@ -1005,7 +822,7 @@ fn get_term_agg_accessors( field_name: &str, missing: &Option, include_bytes: bool, -) -> crate::Result>> { +) -> crate::Result>> { // `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 @@ -1067,19 +884,19 @@ fn build_terms_or_cardinality_nodes( missing: &Option, reader: &SegmentReader, segment_ordinal: SegmentOrdinal, - data: &mut AggregationsSegmentCtx, + context: &AggContextParams, sub_aggs: &Aggregations, req: TermsOrCardinalityRequest, is_top_level: bool, -) -> crate::Result> { +) -> crate::Result> { let mut nodes = Vec::new(); let str_dict_column = reader.fast_fields().str(field_name)?; - let value_sources = data.context.value_sources.clone(); + let value_sources = &context.value_sources; let include_bytes = matches!(req, TermsOrCardinalityRequest::Terms(_)); let sources = - get_term_agg_accessors(reader, &value_sources, field_name, missing, include_bytes)?; + get_term_agg_accessors(reader, value_sources, field_name, missing, include_bytes)?; // Special handling when missing + multi column or incompatible type on text/date. let missing_and_more_than_one_col = sources.len() > 1 && missing.is_some(); @@ -1105,7 +922,7 @@ fn build_terms_or_cardinality_nodes( // documents are missing across several typed columns. There is no way to ask that through // `ValueSource`, so it stays physical-only. let all_accessors = - get_all_value_sources(reader, &value_sources, field_name, None, fallback_type)? + get_all_value_sources(reader, value_sources, field_name, None, fallback_type)? .into_iter() .map(|source| { let column = require_physical_column( @@ -1124,17 +941,13 @@ fn build_terms_or_cardinality_nodes( ) })?; - let children = build_children(sub_aggs, reader, segment_ordinal, data)?; - let idx_in_req_data = data.push_missing_term_req_data(MissingTermAggReqData { + let children = build_children(sub_aggs, reader, segment_ordinal, context)?; + let node_data = AggNodeData::MissingTerm(MissingTermAggReqData { accessors: all_accessors, name: agg_name.to_string(), req, }); - nodes.push(AggRefNode { - kind: AggKind::MissingTerm, - idx_in_req_data, - children, - }); + nodes.push(AggNode::new(node_data, children)); } // Add one node per accessor @@ -1148,8 +961,8 @@ fn build_terms_or_cardinality_nodes( None }; - let children = build_children(sub_aggs, reader, segment_ordinal, data)?; - let (idx, kind) = match req { + let children = build_children(sub_aggs, reader, segment_ordinal, context)?; + let node_data = match req { TermsOrCardinalityRequest::Terms(ref req) => { let mut allowed_term_ids = None; if req.include.is_some() || req.exclude.is_some() { @@ -1168,7 +981,7 @@ fn build_terms_or_cardinality_nodes( missing.is_some(), )?; }; - let idx_in_req_data = data.push_term_req_data(TermsAggReqData { + AggNodeData::Terms(TermsAggReqData { accessor, str_dict_column: str_dict_column.clone(), missing_value_for_accessor, @@ -1177,8 +990,7 @@ fn build_terms_or_cardinality_nodes( sub_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 @@ -1191,21 +1003,16 @@ fn build_terms_or_cardinality_nodes( } else { None }; - let idx_in_req_data = data.push_cardinality_req_data(CardinalityAggReqData { + AggNodeData::Cardinality(CardinalityAggReqData { accessor, 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, - }); + nodes.push(AggNode::new(node_data, children)); } Ok(nodes) @@ -1297,6 +1104,16 @@ mod tests { serde_json::from_value(val).unwrap() } + fn multi_terms_req_data(agg_tree: &[AggNode]) -> Vec<&MultiTermsAggReqData> { + agg_tree + .iter() + .filter_map(|node| match &node.data { + AggNodeData::MultiTerms(req_data) => Some(req_data), + _ => None, + }) + .collect() + } + #[test] fn test_multi_terms_expands_physical_column_cartesian_product() -> crate::Result<()> { let mut schema_builder = crate::schema::Schema::builder(); @@ -1326,24 +1143,19 @@ mod tests { })); let aggs: Aggregations = vec![("mt".to_string(), agg)].into_iter().collect(); let searcher = index.reader()?.searcher(); - let data = build_aggregations_data_from_req( + let agg_tree = build_aggregations_data_from_req( &aggs, searcher.segment_reader(0), 0, - Default::default(), + &Default::default(), )?; + let multi_terms_req_data = multi_terms_req_data(&agg_tree); - 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)); + assert_eq!(agg_tree.len(), 4); + assert_eq!(multi_terms_req_data.len(), 4); + assert!(agg_tree.iter().all(|node| node.children.len() == 1)); - let actual_types: FxHashSet> = data - .per_request - .multi_terms_req_data + let actual_types: FxHashSet> = multi_terms_req_data .iter() .map(|req_data| { req_data @@ -1362,9 +1174,7 @@ mod tests { .into_iter() .collect(); assert_eq!(actual_types, expected_types); - assert!(data - .per_request - .multi_terms_req_data + assert!(multi_terms_req_data .iter() .flat_map(|req_data| &req_data.fields) .all(|field| { @@ -1391,22 +1201,19 @@ mod tests { })); let aggs: Aggregations = vec![("mt".to_string(), agg)].into_iter().collect(); let searcher = index.reader()?.searcher(); - let data = build_aggregations_data_from_req( + let agg_tree = build_aggregations_data_from_req( &aggs, searcher.segment_reader(0), 0, - Default::default(), + &Default::default(), )?; + let multi_terms_req_data = multi_terms_req_data(&agg_tree); - assert_eq!(data.per_request.multi_terms_req_data.len(), 2); - assert!(data - .per_request - .multi_terms_req_data + assert_eq!(multi_terms_req_data.len(), 2); + assert!(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 + assert!(multi_terms_req_data .iter() .all(|req_data| req_data.missing_accessors.iter().all(Option::is_none))); @@ -1445,8 +1252,9 @@ mod tests { .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 agg_tree = + build_aggregations_data_from_req(&aggs, seg_reader, 0u32, &Default::default())?; + let printed_nodes = get_view_tree(&agg_tree); let printed = serde_json::to_value(&printed_nodes).unwrap(); let expected = json!([ diff --git a/src/aggregation/agg_limits.rs b/src/aggregation/agg_limits.rs index 8e023e09b..f2d31c625 100644 --- a/src/aggregation/agg_limits.rs +++ b/src/aggregation/agg_limits.rs @@ -2,7 +2,7 @@ use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; -use common::ByteCount; +use common::{BitSet, ByteCount}; use super::collector::DEFAULT_MEMORY_LIMIT; use super::{AggregationError, DEFAULT_BUCKET_LIMIT}; @@ -12,6 +12,14 @@ pub trait MemoryConsumption { fn memory_consumption(&self) -> usize; } +impl MemoryConsumption for BitSet { + /// Heap bytes of the bitset. It allocates one `u64` word per 64 values up to `max_value`, + /// regardless of how many values are set. + fn memory_consumption(&self) -> usize { + self.max_value().div_ceil(64) as usize * std::mem::size_of::() + } +} + impl MemoryConsumption for HashMap { fn memory_consumption(&self) -> usize { let capacity = self.capacity(); diff --git a/src/aggregation/bucket/composite/accessors.rs b/src/aggregation/bucket/composite/accessors.rs index c21e3a065..d95418101 100644 --- a/src/aggregation/bucket/composite/accessors.rs +++ b/src/aggregation/bucket/composite/accessors.rs @@ -16,7 +16,7 @@ use crate::{SegmentReader, TantivyError}; /// Contains all information required by the SegmentCompositeCollector to perform the /// composite aggregation on a segment. -#[derive(Debug, Clone)] +#[derive(Debug)] pub(crate) struct CompositeAggReqData { /// The name of the aggregation. pub(crate) name: String, diff --git a/src/aggregation/bucket/composite/collector.rs b/src/aggregation/bucket/composite/collector.rs index 171c2c1da..91645899e 100644 --- a/src/aggregation/bucket/composite/collector.rs +++ b/src/aggregation/bucket/composite/collector.rs @@ -10,9 +10,7 @@ use columnar::{ use rustc_hash::FxHashMap; use smallvec::SmallVec; -use crate::aggregation::agg_data::{ - build_segment_agg_collectors, AggRefNode, AggregationsSegmentCtx, -}; +use crate::aggregation::agg_data::AggregationsSegmentCtx; use crate::aggregation::bucket::composite::accessors::{ CompositeAccessor, CompositeAggReqData, PrecomputedDateInterval, }; @@ -213,24 +211,11 @@ impl SegmentCompositeCollector { } pub(crate) fn from_req_and_validate( - req_data: &mut AggregationsSegmentCtx, - node: &AggRefNode, + composite_req_data: CompositeAggReqData, + sub_agg: Option>, ) -> crate::Result { - let composite_req_data = - req_data.per_request.composite_req_data[node.idx_in_req_data].clone(); validate_req(&composite_req_data)?; - req_data - .context - .limits - .add_memory_consumed(composite_req_data.get_memory_consumption() as u64)?; - - let has_sub_aggregations = !node.children.is_empty(); - let sub_agg = if has_sub_aggregations { - let sub_agg_collector = build_segment_agg_collectors(req_data, &node.children)?; - Some(BufferedSubAggs::new(sub_agg_collector)) - } else { - None - }; + let sub_agg = sub_agg.map(BufferedSubAggs::new); let num_sources = composite_req_data.req.sources.len(); diff --git a/src/aggregation/bucket/filter.rs b/src/aggregation/bucket/filter.rs index 0c08d08c7..8d2797f62 100644 --- a/src/aggregation/bucket/filter.rs +++ b/src/aggregation/bucket/filter.rs @@ -1,12 +1,10 @@ use std::fmt::Debug; -use std::rc::Rc; use common::BitSet; use serde::{Deserialize, Deserializer, Serialize, Serializer}; -use crate::aggregation::agg_data::{ - build_segment_agg_collectors, AggRefNode, AggregationsSegmentCtx, -}; +use crate::aggregation::agg_data::{build_sub_agg_collectors, AggNode, AggregationsSegmentCtx}; +use crate::aggregation::agg_limits::MemoryConsumption; use crate::aggregation::buffered_sub_aggs::{ BufferedSubAggs, HighCardSubAggBuffer, LowCardSubAggBuffer, SubAggBuffer, }; @@ -397,27 +395,21 @@ impl PartialEq for FilterAggregation { /// Request data for filter aggregation /// This struct holds the per-segment data needed to execute a filter aggregation -#[derive(Clone)] pub(crate) struct FilterAggReqData { /// The name of the filter aggregation pub(crate) name: String, /// The segment reader pub(crate) segment_reader: SegmentReader, /// Document evaluator for the filter query (precomputed BitSet). - /// Wrapped in `Rc` so cloning the request data does not duplicate the (potentially large) - /// underlying BitSet. - pub(crate) evaluator: Rc, + pub(crate) evaluator: DocumentQueryEvaluator, /// True if this filter aggregation is at the top level of the aggregation tree (not nested). pub(crate) is_top_level: bool, } impl FilterAggReqData { pub(crate) fn get_memory_consumption(&self) -> usize { - // Estimate: name + segment reader reference + bitset - self.name.len() - + std::mem::size_of::() - + self.evaluator.bitset.len() / 8 // BitSet memory (bits to bytes) - + std::mem::size_of::() + // The struct itself (including the segment reader handle), the name and the bitset. + std::mem::size_of::() + self.name.len() + self.evaluator.bitset.memory_consumption() } } @@ -516,52 +508,39 @@ pub struct SegmentFilterCollector { impl SegmentFilterCollector { /// Create a new filter segment collector following the new agg_data pattern pub(crate) fn from_req_and_validate( - req: &mut AggregationsSegmentCtx, - node: &AggRefNode, req_data: FilterAggReqData, - ) -> crate::Result { - // Build sub-aggregation collectors if any - let sub_agg_collector = if !node.children.is_empty() { - Some(build_segment_agg_collectors(req, &node.children)?) - } else { - None - }; - let sub_agg_collector = sub_agg_collector.map(BufferedSubAggs::new); - + sub_agg: Option>, + ) -> Self { let max_doc = req_data.segment_reader.max_doc(); let buffer_capacity = crate::docset::COLLECT_BLOCK_BUFFER_LEN.min(max_doc as usize); - Ok(SegmentFilterCollector { + SegmentFilterCollector { parent_buckets: Vec::new(), - sub_aggregations: sub_agg_collector, + sub_aggregations: sub_agg.map(BufferedSubAggs::new), req_data, bucket_id_provider: BucketIdProvider::default(), matching_docs_buffer: Vec::with_capacity(buffer_capacity), - }) + } } } pub(crate) fn build_segment_filter_collector( req: &mut AggregationsSegmentCtx, - node: &AggRefNode, + req_data: FilterAggReqData, + children: Vec, ) -> crate::Result> { - let req_data = req.per_request.filter_req_data[node.idx_in_req_data].clone(); - req.context - .limits - .add_memory_consumed(req_data.get_memory_consumption() as u64)?; let is_top_level = req_data.is_top_level; + let sub_agg = build_sub_agg_collectors(req, children)?; if is_top_level { Ok(Box::new( - SegmentFilterCollector::::from_req_and_validate( - req, node, req_data, - )?, + SegmentFilterCollector::::from_req_and_validate(req_data, sub_agg), )) } else { Ok(Box::new( SegmentFilterCollector::::from_req_and_validate( - req, node, req_data, - )?, + req_data, sub_agg, + ), )) } } diff --git a/src/aggregation/bucket/histogram/histogram.rs b/src/aggregation/bucket/histogram/histogram.rs index 80079503d..4142ff935 100644 --- a/src/aggregation/bucket/histogram/histogram.rs +++ b/src/aggregation/bucket/histogram/histogram.rs @@ -1,14 +1,11 @@ use std::cmp::Ordering; -use std::sync::Arc; use columnar::ColumnType; use rustc_hash::FxHashMap; use serde::{Deserialize, Serialize}; use tantivy_bitpacker::minmax; -use crate::aggregation::agg_data::{ - build_segment_agg_collectors, AggRefNode, AggregationsSegmentCtx, -}; +use crate::aggregation::agg_data::{build_sub_agg_collectors, AggNode, AggregationsSegmentCtx}; use crate::aggregation::agg_req::Aggregations; use crate::aggregation::agg_result::BucketEntry; use crate::aggregation::buffered_sub_aggs::{BufferedSubAggs, HighCardBufferedSubAggs}; @@ -22,10 +19,10 @@ use crate::TantivyError; /// Contains all information required by the SegmentHistogramCollector to perform the /// histogram or date_histogram aggregation on a segment. -#[derive(Debug, Clone)] +#[derive(Debug)] pub(crate) struct HistogramAggReqData { /// The column accessor to access the fast field values. - pub(crate) accessor: Arc, + pub(crate) accessor: Box, /// The name of the aggregation. pub(crate) name: String, /// The histogram aggregation request. @@ -37,6 +34,55 @@ pub(crate) struct HistogramAggReqData { /// The offset used to calculate the bucket position. pub(crate) offset: f64, } + +impl HistogramAggReqData { + /// Creates the request data of a histogram aggregation, validating and normalizing `req`. + pub(crate) fn new( + accessor: Box, + name: String, + mut req: HistogramAggregation, + is_date_histogram: bool, + ) -> crate::Result { + req.validate()?; + let field_type = accessor.column_type(); + // Date values are stored in ns, but the request uses ms. + if field_type == ColumnType::DateTime { + req.normalize_date_time(); + } + + let bounds = req + .hard_bounds + .as_ref() + .filter(|hard_bounds| { + // Drop `hard_bounds` that can't exclude any value (the column's range already sits + // inside them): the per-doc `bounds.contains` check is then a no-op, so collapsing + // to the unbounded sentinel lets the histogram hot loop skip it and the flattened + // term×histogram path derive per-term counts from the grid. Only this + // collect-time filter is touched — empty-bucket emission reads `req.hard_bounds` + // directly (see `get_req_min_max`), and `hard_bounds` only ever clips that range, + // so a wider-than-data bound leaves the result unchanged. + let Some((min_value, max_value)) = accessor.bounds() else { + return true; + }; + f64_from_fastfield_u64(min_value, field_type) < hard_bounds.min + || f64_from_fastfield_u64(max_value, field_type) > hard_bounds.max + }) + .cloned() + .unwrap_or(HistogramBounds { + min: f64::MIN, + max: f64::MAX, + }); + let offset = req.offset.unwrap_or(0.0); + Ok(HistogramAggReqData { + accessor, + name, + req, + is_date_histogram, + bounds, + offset, + }) + } +} impl HistogramAggReqData { /// Estimate the memory consumption of this struct in bytes. pub fn get_memory_consumption(&self) -> usize { @@ -483,15 +529,14 @@ impl SegmentAggregationCollector for SegmentHistogramCollector< // Upgrade to dense storage before processing the block if the buckets are dense enough. store.maybe_densify(dense_range); - let req = &self.req_data; - let bounds = req.bounds; - let interval = req.req.interval; - let offset = req.offset; + let bounds = self.req_data.bounds; + let interval = self.req_data.req.interval; + let offset = self.req_data.offset; let get_bucket_pos = |val| get_bucket_pos_f64(val, interval, offset) as i64; agg_data .column_block_accessor - .fetch_block(docs, &*req.accessor); + .fetch_block(docs, &mut *self.req_data.accessor); // special path for nested buckets if let Some(sub_agg) = &mut self.sub_agg { for (doc, val) in agg_data.column_block_accessor.iter_docid_vals(docs) { @@ -594,20 +639,9 @@ impl SegmentHistogramCollector { } pub(crate) fn from_req_and_validate( - agg_data: &mut AggregationsSegmentCtx, - node: &AggRefNode, - ) -> crate::Result { - let sub_agg = if !node.children.is_empty() { - Some(build_segment_agg_collectors(agg_data, &node.children)?) - } else { - None - }; - let mut req_data = agg_data.per_request.histogram_req_data[node.idx_in_req_data].clone(); - normalize_histogram_req(&mut req_data)?; - agg_data - .context - .limits - .add_memory_consumed(req_data.get_memory_consumption() as u64)?; + req_data: HistogramAggReqData, + sub_agg: Option>, + ) -> Self { let dense_range = compute_dense_range( &*req_data.accessor, req_data.req.interval, @@ -616,33 +650,47 @@ impl SegmentHistogramCollector { ); let sub_agg = sub_agg.map(BufferedSubAggs::new); let column_type = req_data.accessor.column_type(); - Ok(Self { + Self { parent_buckets: Default::default(), sub_agg, req_data, column_type, bucket_id_provider: BucketIdProvider::default(), dense_range, - }) + } } } impl SegmentHistogramCollector<()> { - /// Builds a histogram collector whose parent `t` is a dense histogram filled from - /// `counts[t * num_time_buckets .. (t + 1) * num_time_buckets]` (row-major), consolidating each - /// cell's count lanes. Used by the flattened terms×histogram collector to turn its flat 2D - /// counters into the regular intermediate result, so cross-segment merging is shared with the + /// Builds an empty histogram collector without sub aggregations, used by the flattened + /// terms×histogram collector to turn its flat 2D counters into the regular intermediate + /// result (see [`Self::fill_from_dense_rows`]), so cross-segment merging is shared with the /// general path. - pub(crate) fn from_dense_rows( - req_data: HistogramAggReqData, + pub(crate) fn from_dense_rows(req_data: HistogramAggReqData) -> Self { + let column_type = req_data.accessor.column_type(); + Self { + parent_buckets: Vec::new(), + sub_agg: None, + req_data, + column_type, + bucket_id_provider: BucketIdProvider::default(), + dense_range: None, + } + } + + /// Replaces the buckets of this collector: parent `t` becomes a dense histogram filled from + /// `counts[t * num_time_buckets .. (t + 1) * num_time_buckets]` (row-major), consolidating + /// each cell's count lanes. + pub(crate) fn fill_from_dense_rows( + &mut self, base_pos: i64, num_time_buckets: usize, counts: &[[u32; LANES]], - ) -> Self { - let interval = req_data.req.interval; - let offset = req_data.offset; + ) { + let interval = self.req_data.req.interval; + let offset = self.req_data.offset; let num_parents = counts.len().checked_div(num_time_buckets).unwrap_or(0); - let parent_buckets = (0..num_parents) + self.parent_buckets = (0..num_parents) .map(|t| { let row = &counts[t * num_time_buckets..(t + 1) * num_time_buckets]; let buckets = row @@ -661,68 +709,18 @@ impl SegmentHistogramCollector<()> { HistogramBuckets::Dense { base_pos, buckets } }) .collect(); - let column_type = req_data.accessor.column_type(); - Self { - parent_buckets, - sub_agg: None, - req_data, - column_type, - bucket_id_provider: BucketIdProvider::default(), - dense_range: None, - } } } -/// Validates and normalizes a histogram request in place: applies date ns-normalization (for a -/// `histogram` on a date column) and resolves `bounds`/`offset` from the request. -fn normalize_histogram_req(req_data: &mut HistogramAggReqData) -> crate::Result<()> { - req_data.req.validate()?; - let field_type = req_data.accessor.column_type(); - if field_type == ColumnType::DateTime && !req_data.is_date_histogram { - req_data.req.normalize_date_time(); - } - req_data.bounds = req_data.req.hard_bounds.unwrap_or(HistogramBounds { - min: f64::MIN, - max: f64::MAX, - }); - req_data.offset = req_data.req.offset.unwrap_or(0.0); - // Drop `hard_bounds` that can't exclude any value (the column's range already sits inside - // them): the per-doc `bounds.contains` check is then a no-op, so collapsing to the unbounded - // sentinel lets the histogram hot loop skip it and the flattened term×histogram path derive - // per-term counts from the grid. Only this collect-time filter is touched — empty-bucket - // emission reads `req.hard_bounds` directly (see `get_req_min_max`), and `hard_bounds` only - // ever clips that range, so a wider-than-data bound leaves the result unchanged. - if req_data.req.hard_bounds.is_some() { - if let Some((min_value, max_value)) = req_data.accessor.bounds() { - let col_min = f64_from_fastfield_u64(min_value, field_type); - let col_max = f64_from_fastfield_u64(max_value, field_type); - if col_min >= req_data.bounds.min && col_max <= req_data.bounds.max { - req_data.bounds = HistogramBounds { - min: f64::MIN, - max: f64::MAX, - }; - } - } - } - Ok(()) -} - -/// Clones and normalizes (resolving interval/offset/bounds) the histogram request at `node`, and -/// returns it together with its dense bucket range — or `None` if the column has no usable range. -/// Used by the flattened terms×histogram collector, which then owns the normalized request. -pub(crate) fn prepare_histogram_dense_range( - agg_data: &AggregationsSegmentCtx, - node: &AggRefNode, -) -> crate::Result> { - let mut req_data = agg_data.per_request.histogram_req_data[node.idx_in_req_data].clone(); - normalize_histogram_req(&mut req_data)?; - let dense_range = compute_dense_range( +/// Returns the dense bucket range of a normalized histogram request, or `None` if the column has +/// no usable range. Used by the flattened terms×histogram collector. +pub(crate) fn histogram_dense_range(req_data: &HistogramAggReqData) -> Option { + compute_dense_range( &*req_data.accessor, req_data.req.interval, req_data.offset, req_data.bounds, - ); - Ok(dense_range.map(|range| (req_data, range))) + ) } /// Builds a boxed histogram (or date histogram) segment collector, picking the bucket-id storage @@ -730,15 +728,17 @@ pub(crate) fn prepare_histogram_dense_range( /// [`BucketId`]. pub(crate) fn build_segment_histogram_collector( agg_data: &mut AggregationsSegmentCtx, - node: &AggRefNode, + req_data: HistogramAggReqData, + children: Vec, ) -> crate::Result> { - if node.children.is_empty() { + let sub_agg = build_sub_agg_collectors(agg_data, children)?; + if sub_agg.is_some() { Ok(Box::new( - SegmentHistogramCollector::<()>::from_req_and_validate(agg_data, node)?, + SegmentHistogramCollector::::from_req_and_validate(req_data, sub_agg), )) } else { Ok(Box::new( - SegmentHistogramCollector::::from_req_and_validate(agg_data, node)?, + SegmentHistogramCollector::<()>::from_req_and_validate(req_data, sub_agg), )) } } diff --git a/src/aggregation/bucket/multi_terms/mod.rs b/src/aggregation/bucket/multi_terms/mod.rs index 0b2d53c23..f0a4aaaa5 100644 --- a/src/aggregation/bucket/multi_terms/mod.rs +++ b/src/aggregation/bucket/multi_terms/mod.rs @@ -14,7 +14,7 @@ use smallvec::SmallVec; use super::{BucketIdSlot, CustomOrder, Order, OrderTarget}; use crate::aggregation::agg_data::{ - build_segment_agg_collectors, AggRefNode, AggregationsSegmentCtx, PerRequestAggSegCtx, + build_sub_agg_collectors, find_sub_agg, AggNode, AggregationsSegmentCtx, }; use crate::aggregation::agg_req::Aggregations; use crate::aggregation::bucket::term_agg::{ @@ -158,7 +158,7 @@ impl MultiTermsFieldAccessor { } /// Per-request data bundle passed to the segment collector. -#[derive(Debug, Clone)] +#[derive(Debug)] pub(crate) struct MultiTermsAggReqData { /// Aggregation name used to look up this entry in the result tree. pub(crate) name: String, @@ -229,7 +229,7 @@ fn fetch_field_block( let missing_value = block_missing_value(missing); block_accessor.fetch_block_with_missing_unique_per_doc( docs, - &(&field.column, field.column_type), + &mut (&field.column, field.column_type), missing_value, true, ); @@ -376,8 +376,7 @@ where /// exists. fn validate_multi_terms( req_data: &MultiTermsAggReqData, - node: &AggRefNode, - per_request: &PerRequestAggSegCtx, + children: &[AggNode], ) -> crate::Result<()> { if req_data.fields.is_empty() { return Err(TantivyError::InvalidArgument( @@ -405,7 +404,7 @@ fn validate_multi_terms( let req_agg = MultiTermsAggregationInternal::from_req(&req_data.req); if let OrderTarget::SubAggregation(sub_agg_name) = &req_agg.order.target { let (agg_name, _) = get_agg_name_and_property(sub_agg_name); - node.get_sub_agg(agg_name, per_request).ok_or_else(|| { + find_sub_agg(children, agg_name).ok_or_else(|| { TantivyError::InvalidArgument(format!( "could not find aggregation with name {agg_name} in metric sub_aggregations" )) @@ -519,16 +518,17 @@ fn expand_partial_combinations_for_field( /// uses [`MultiTermsKey`]. pub(crate) fn build_segment_multi_terms_collector( req: &mut AggregationsSegmentCtx, - node: &AggRefNode, + req_data: MultiTermsAggReqData, + children: Vec, ) -> crate::Result> { - let req_data = req.per_request.multi_terms_req_data[node.idx_in_req_data].clone(); - validate_multi_terms(&req_data, node, &req.per_request)?; + validate_multi_terms(&req_data, &children)?; + let sub_agg = build_sub_agg_collectors(req, children)?; let packed_layout = compute_packed_u64_layout(&req_data.fields, &req_data.missing_accessors); - if node.children.is_empty() { - build_multi_terms_collector::<()>(req, node, req_data, packed_layout) + if sub_agg.is_none() { + build_multi_terms_collector::<()>(req, sub_agg, req_data, packed_layout) } else { - build_multi_terms_collector::(req, node, req_data, packed_layout) + build_multi_terms_collector::(req, sub_agg, req_data, packed_layout) } } @@ -974,7 +974,7 @@ struct PackedU64KeyPacking { /// Selects the key packing and bucket storage, then boxes the concrete collector. fn build_multi_terms_collector( req: &mut AggregationsSegmentCtx, - node: &AggRefNode, + sub_agg: Option>, req_data: MultiTermsAggReqData, packed_layout: Option>, ) -> crate::Result> { @@ -984,12 +984,9 @@ fn build_multi_terms_collector( let num_terms = max_packed.saturating_add(1); let packing = PackedU64KeyPacking { packs }; - if is_top_level - && !node.children.is_empty() - && max_packed < MAX_NUM_TERMS_FOR_LOWCARD_SUBAGG - { + if is_top_level && sub_agg.is_some() && max_packed < MAX_NUM_TERMS_FOR_LOWCARD_SUBAGG { return box_multi_terms_collector::>( - req, node, req_data, packing, num_terms, true, + req, sub_agg, req_data, packing, num_terms, true, ); } if is_top_level && max_packed < MAX_NUM_TERMS_FOR_VEC { @@ -998,19 +995,19 @@ fn build_multi_terms_collector( HighCardSubAggBuffer, _, VecTermBuckets, - >(req, node, req_data, packing, num_terms, true); + >(req, sub_agg, req_data, packing, num_terms, true); } return box_multi_terms_collector::>( - req, node, req_data, packing, num_terms, true, + req, sub_agg, req_data, packing, num_terms, true, ); } if is_top_level && max_packed < MAX_NUM_TERMS_FOR_PAGED_MAP { return box_multi_terms_collector::>( - req, node, req_data, packing, max_packed, false, + req, sub_agg, req_data, packing, max_packed, false, ); } return box_multi_terms_collector::>( - req, node, req_data, packing, 0, false, + req, sub_agg, req_data, packing, 0, false, ); } @@ -1021,12 +1018,12 @@ fn build_multi_terms_collector( HighCardSubAggBuffer, _, HashMapTermBuckets, - >(req, node, req_data, packing, 0, false) + >(req, sub_agg, req_data, packing, 0, false) } fn box_multi_terms_collector( req: &mut AggregationsSegmentCtx, - node: &AggRefNode, + sub_agg: Option>, req_data: MultiTermsAggReqData, packing: Packing, map_init_value: u64, @@ -1044,13 +1041,7 @@ where .limits .add_memory_consumed(buckets.get_memory_consumption() as u64)?; } - let sub_agg = if node.children.is_empty() { - None - } else { - Some(BufferedSubAggs::::new( - build_segment_agg_collectors(req, &node.children)?, - )) - }; + let sub_agg = sub_agg.map(BufferedSubAggs::::new); let all_fields_full = req_data .fields .iter() diff --git a/src/aggregation/bucket/range.rs b/src/aggregation/bucket/range.rs index e2b581367..6904ef98f 100644 --- a/src/aggregation/bucket/range.rs +++ b/src/aggregation/bucket/range.rs @@ -1,14 +1,11 @@ use std::fmt::Debug; use std::ops::Range; -use std::sync::Arc; use columnar::ColumnType; use rustc_hash::FxHashMap; use serde::{Deserialize, Serialize}; -use crate::aggregation::agg_data::{ - build_segment_agg_collectors, AggRefNode, AggregationsSegmentCtx, -}; +use crate::aggregation::agg_data::{build_sub_agg_collectors, AggNode, AggregationsSegmentCtx}; use crate::aggregation::agg_limits::AggregationLimitsGuard; use crate::aggregation::buffered_sub_aggs::{ BufferedSubAggs, HighCardSubAggBuffer, LowCardBufferedSubAggs, LowCardSubAggBuffer, @@ -25,10 +22,10 @@ use crate::TantivyError; /// Contains all information required by the SegmentRangeCollector to perform the /// range aggregation on a segment. -#[derive(Debug, Clone)] +#[derive(Debug)] pub(crate) struct RangeAggReqData { /// The column accessor to access the fast field values. - pub(crate) accessor: Arc, + pub(crate) accessor: Box, /// The range aggregation request. pub(crate) req: RangeAggregation, /// The name of the aggregation. @@ -280,7 +277,7 @@ impl SegmentAggregationCollector for SegmentRangeCollector { ) -> crate::Result<()> { agg_data .column_block_accessor - .fetch_block(docs, &*self.req_data.accessor); + .fetch_block(docs, &mut *self.req_data.accessor); let buckets = &mut self.parent_buckets[parent_bucket_id as usize]; @@ -335,24 +332,15 @@ impl SegmentAggregationCollector for SegmentRangeCollector { /// bucket storage, depending on the column type and aggregation level. pub(crate) fn build_segment_range_collector( agg_data: &mut AggregationsSegmentCtx, - node: &AggRefNode, + req_data: RangeAggReqData, + children: Vec, ) -> crate::Result> { - let req_data = agg_data.per_request.range_req_data[node.idx_in_req_data].clone(); - agg_data - .context - .limits - .add_memory_consumed(req_data.get_memory_consumption() as u64)?; - // TODO: A better metric instead of is_top_level would be the number of buckets expected. // E.g. If range agg is not top level, but the parent is a bucket agg with less than 10 buckets, // we can are still in low cardinality territory. let is_low_card = req_data.is_top_level && req_data.req.ranges.len() <= 64; - let sub_agg = if !node.children.is_empty() { - Some(build_segment_agg_collectors(agg_data, &node.children)?) - } else { - None - }; + let sub_agg = build_sub_agg_collectors(agg_data, children)?; if is_low_card { Ok(Box::new(SegmentRangeCollector:: { diff --git a/src/aggregation/bucket/term_agg/flattened_term_histogram.rs b/src/aggregation/bucket/term_agg/flattened_term_histogram.rs index 3f741539b..783fd02d8 100644 --- a/src/aggregation/bucket/term_agg/flattened_term_histogram.rs +++ b/src/aggregation/bucket/term_agg/flattened_term_histogram.rs @@ -2,7 +2,7 @@ //! `histogram`/`date_histogram` sub-aggregation with nothing nested below it. //! //! See [`FlattenedTermHistogramCollector`] for the approach and -//! [`maybe_build_flattened_collector`] for the conditions under which it is used. +//! [`plan_flattened_collector`] for the conditions under which it is used. use std::fmt::Debug; use std::sync::Arc; @@ -13,9 +13,9 @@ use super::{ Bucket, SegmentTermCollector, TermsAggReqData, VecTermBuckets, MAX_NUM_BUCKETS_FOR_COUNT_LANES, NUM_LOW_CARD_COUNT_LANES, }; -use crate::aggregation::agg_data::{AggKind, AggRefNode, AggregationsSegmentCtx}; +use crate::aggregation::agg_data::{AggNode, AggNodeData, AggregationsSegmentCtx}; use crate::aggregation::bucket::{ - get_bucket_pos_f64, prepare_histogram_dense_range, DenseRange, HistogramAggReqData, + get_bucket_pos_f64, histogram_dense_range, DenseRange, HistogramAggReqData, SegmentHistogramCollector, }; use crate::aggregation::buffered_sub_aggs::LowCardSubAggBuffer; @@ -367,7 +367,9 @@ struct FlattenedTermHistogramCollector { terms_req_data: TermsAggReqData, /// The terms full column's values terms_values: Arc, - hist_req_data: HistogramAggReqData, + /// Histogram collector without any bucket. It is only used to produce the histogram + /// intermediate results, once filled from `counts`. + histogram: SegmentHistogramCollector<()>, /// Private term block accessor. The bucket resolver owns a histogram block accessor when it /// needs one; the single-bucket resolver deliberately does not. term_block: ColumnBlockAccessor, @@ -423,16 +425,12 @@ impl SegmentAggregationCollector }) .collect(), }; - let mut histogram = SegmentHistogramCollector::<()>::from_dense_rows( - self.hist_req_data.clone(), - self.base_pos, - num_time_buckets, - &self.counts, - ); + self.histogram + .fill_from_dense_rows(self.base_pos, num_time_buckets, &self.counts); let name = self.terms_req_data.name.clone(); let bucket = SegmentTermCollector::, LowCardSubAggBuffer>::into_intermediate_bucket_result( &self.terms_req_data, - Some(&mut histogram as &mut dyn SegmentAggregationCollector), + Some(&mut self.histogram as &mut dyn SegmentAggregationCollector), term_buckets, agg_data, )?; @@ -499,19 +497,31 @@ impl SegmentAggregationCollector } } -/// Builds the flattened terms×histogram collector for a single top-level parent, when the shape is -/// eligible. Returns `Ok(None)` to fall back to the general buffered terms path. +/// The resolved inputs of the flattened terms×histogram collector, computed by +/// [`plan_flattened_collector`]. +pub(super) struct FlattenedPlan { + terms_values: Arc, + hist_values: Arc, + num_terms: usize, + range: DenseRange, + use_count_lanes: bool, +} + +/// Checks whether the flattened terms×histogram collector can be used for a single top-level +/// parent. Returns `None` to fall back to the general buffered terms path. +/// +/// This only inspects the request data: the decision is taken before the terms and histogram +/// request data are moved into a collector. /// /// Eligibility: top-level, low-cardinality terms over a full column with no missing/include-exclude /// handling; a single `histogram`/`date_histogram` leaf (no nesting below it) over a full column; /// and a physical counter grid no larger than [`MAX_FLATTENED_GRID_COUNTERS`]. -pub(super) fn maybe_build_flattened_collector( - agg_data: &mut AggregationsSegmentCtx, - node: &AggRefNode, +pub(super) fn plan_flattened_collector( terms_req_data: &TermsAggReqData, + children: &[AggNode], col_max_val: u64, is_top_level: bool, -) -> crate::Result>> { +) -> Option { // Both columns must be full (one value per doc) so their values align positionally with `docs` // and we can zip them. Requiring full columns also makes the terms agg's `missing` config a // no-op (`fetch_block_with_missing` early-returns on full columns), so we needn't check for it. @@ -524,40 +534,37 @@ pub(super) fn maybe_build_flattened_collector( // are less likely to get enough docs for the preallocation to be worth and there's a risk of // using too much memory. We could check the maximum theoretical buckets up-front and pass // them down. - let fuseable = is_top_level + if !is_top_level // TODO: We can easily support this - && terms_req_data.allowed_term_ids.is_none() - && node.children.len() == 1 - && matches!( - node.children[0].kind, - AggKind::Histogram | AggKind::DateHistogram - ) - && node.children[0].children.is_empty(); - if !fuseable { - return Ok(None); + || terms_req_data.allowed_term_ids.is_some() + { + return None; + } + let [AggNode { + data: AggNodeData::Histogram(hist_req_data), + children: hist_children, + }] = children + else { + return None; + }; + if !hist_children.is_empty() { + return None; } // Check fullness once, here, and hand the bare values array downstream. - let Some(terms_values) = try_get_column_full_values(&*terms_req_data.accessor) else { - return Ok(None); - }; + let terms_values = try_get_column_full_values(&*terms_req_data.accessor)?; // The flat counters are `u32`, bumped once per value, so no count can exceed the // column's value count. (Essentially always true here: the column is full, so its // value count equals the doc count, and `DocId` is `u32`.) if terms_values.num_vals() == u32::MAX { - return Ok(None); + return None; } - // Clone + normalize the histogram request and get its dense bucket range; only take the + // Get the dense bucket range of the (already normalized) histogram request; only take the // flattened path when the physical counter grid is small enough. Very small logical grids use // multiple counters per cell; larger grids retain scalar cells to avoid paying for lanes when // writes are already spread across many locations. - let Some((hist_req_data, range)) = prepare_histogram_dense_range(agg_data, &node.children[0])? - else { - return Ok(None); - }; - let Some(hist_values) = try_get_column_full_values(&*hist_req_data.accessor) else { - return Ok(None); - }; + let range = histogram_dense_range(hist_req_data)?; + let hist_values = try_get_column_full_values(&*hist_req_data.accessor)?; let num_terms = col_max_val.saturating_add(1) as usize; let num_grid_cells = num_terms.saturating_mul(range.len); @@ -568,10 +575,53 @@ pub(super) fn maybe_build_flattened_collector( SINGLE_COUNT_LANE }; if num_grid_cells.saturating_mul(num_count_lanes) > MAX_FLATTENED_GRID_COUNTERS { - return Ok(None); + return None; } + Some(FlattenedPlan { + terms_values, + hist_values, + num_terms, + range, + use_count_lanes, + }) +} - let collector = if use_count_lanes { +/// Builds the flattened terms×histogram collector, consuming the terms request data and its single +/// histogram child. +/// +/// Hidden contract: `plan` must have been returned by [`plan_flattened_collector`] for these +/// `terms_req_data` and `children`. +pub(super) fn build_flattened_collector_from_plan( + agg_data: &mut AggregationsSegmentCtx, + plan: FlattenedPlan, + terms_req_data: TermsAggReqData, + mut children: Vec, +) -> crate::Result> { + // The histogram child is consumed here rather than by `build_segment_agg_collectors`, so its + // request data has to be charged here. + let children_memory_consumption: usize = + children.iter().map(AggNode::get_memory_consumption).sum(); + agg_data + .context + .limits + .add_memory_consumed(children_memory_consumption as u64)?; + let Some(AggNode { + data: AggNodeData::Histogram(hist_req_data), + .. + }) = children.pop() + else { + return Err(crate::TantivyError::InternalError( + "the flattened terms×histogram collector requires a single histogram child".to_string(), + )); + }; + let FlattenedPlan { + terms_values, + hist_values, + num_terms, + range, + use_count_lanes, + } = plan; + if use_count_lanes { build_flattened_collector::( agg_data, terms_req_data, @@ -580,7 +630,7 @@ pub(super) fn maybe_build_flattened_collector( hist_values, num_terms, range, - )? + ) } else { build_flattened_collector::( agg_data, @@ -590,9 +640,8 @@ pub(super) fn maybe_build_flattened_collector( hist_values, num_terms, range, - )? - }; - Ok(Some(collector)) + ) + } } fn try_get_column_full_values(value_source: &dyn ValueSource) -> Option> { @@ -606,7 +655,7 @@ fn try_get_column_full_values(value_source: &dyn ValueSource) -> Option( agg_data: &mut AggregationsSegmentCtx, - terms_req_data: &TermsAggReqData, + terms_req_data: TermsAggReqData, hist_req_data: HistogramAggReqData, terms_values: Arc, hist_values: Arc, @@ -686,7 +735,7 @@ fn build_flattened_collector( fn build_flattened_collector_with_resolver( agg_data: &mut AggregationsSegmentCtx, - terms_req_data: &TermsAggReqData, + terms_req_data: TermsAggReqData, terms_values: Arc>, hist_req_data: HistogramAggReqData, num_terms: usize, @@ -716,9 +765,9 @@ fn build_flattened_collector_with_resolver, + pub(crate) accessor: Box, /// The string dictionary column if the field is of type text. pub(crate) str_dict_column: Option, /// The missing value as u64 value. @@ -59,11 +58,10 @@ impl TermsAggReqData { /// Estimate the memory consumption of this struct in bytes. pub fn get_memory_consumption(&self) -> usize { std::mem::size_of::() - + std::mem::size_of::() + self .allowed_term_ids .as_ref() - .map(|bs| bs.len() / 8) + .map(BitSet::memory_consumption) .unwrap_or(0) } } @@ -390,9 +388,9 @@ fn add_memory_consumption( /// bucket storage, depending on the column type and aggregation level. pub(crate) fn build_segment_term_collector( req_data: &mut AggregationsSegmentCtx, - node: &AggRefNode, + terms_req_data: TermsAggReqData, + children: Vec, ) -> crate::Result> { - let terms_req_data = req_data.get_term_req_data(node.idx_in_req_data).clone(); let column_type = terms_req_data.accessor.column_type(); if column_type == ColumnType::Bytes { @@ -404,16 +402,15 @@ pub(crate) fn build_segment_term_collector( // Validate that the referenced sub-aggregation exists when ordering by one. if let OrderTarget::SubAggregation(sub_agg_name) = &terms_req_data.req.order.target { let (agg_name, _agg_property) = get_agg_name_and_property(sub_agg_name); - node.get_sub_agg(agg_name, &req_data.per_request) - .ok_or_else(|| { - TantivyError::InvalidArgument(format!( - "could not find aggregation with name {agg_name} in metric sub_aggregations" - )) - })?; + find_sub_agg(&children, agg_name).ok_or_else(|| { + TantivyError::InvalidArgument(format!( + "could not find aggregation with name {agg_name} in metric sub_aggregations" + )) + })?; } // Build sub-aggregation blueprint if there are children. - let has_sub_aggregations = !node.children.is_empty(); + let has_sub_aggregations = !children.is_empty(); // TODO: A better metric instead of is_top_level would be the number of buckets expected. // E.g. If term agg is not top level, but the parent is a bucket agg with less than 10 buckets, @@ -433,21 +430,21 @@ pub(crate) fn build_segment_term_collector( // Flattened fast path: low-cardinality terms × a single `histogram`/`date_histogram` leaf over // full columns with a small enough bucket grid. Anything else falls through to the general // path. - if let Some(collector) = flattened_term_histogram::maybe_build_flattened_collector( - req_data, - node, + if let Some(plan) = flattened_term_histogram::plan_flattened_collector( &terms_req_data, + &children, max_column_val, is_top_level, - )? { - return Ok(collector); + ) { + return flattened_term_histogram::build_flattened_collector_from_plan( + req_data, + plan, + terms_req_data, + children, + ); } - let sub_agg_collector = if has_sub_aggregations { - Some(build_segment_agg_collectors(req_data, &node.children)?) - } else { - None - }; + let sub_agg_collector = build_sub_agg_collectors(req_data, children)?; let mut bucket_id_provider = BucketIdProvider::default(); // Decide which bucket storage is best suited for this aggregation. @@ -1067,7 +1064,7 @@ impl SegmentAggregationCollector .column_block_accessor .fetch_block_with_missing_unique_per_doc( docs, - &*req_data.accessor, + &mut *req_data.accessor, req_data.missing_value_for_accessor, false, ); diff --git a/src/aggregation/bucket/term_missing_agg.rs b/src/aggregation/bucket/term_missing_agg.rs index 1da26d856..21355f975 100644 --- a/src/aggregation/bucket/term_missing_agg.rs +++ b/src/aggregation/bucket/term_missing_agg.rs @@ -1,8 +1,6 @@ use columnar::{Column, ColumnType}; -use crate::aggregation::agg_data::{ - build_segment_agg_collectors, AggRefNode, AggregationsSegmentCtx, -}; +use crate::aggregation::agg_data::AggregationsSegmentCtx; use crate::aggregation::bucket::term_agg::TermsAggregation; use crate::aggregation::buffered_sub_aggs::{BufferedSubAggs, HighCardBufferedSubAggs}; use crate::aggregation::intermediate_agg_result::{ @@ -19,7 +17,7 @@ use crate::aggregation::BucketId; /// - The field is multi-valued and we therefore have multiple columns /// - The field is not text and missing is provided as string (we cannot use the numeric missing /// value optimization) -#[derive(Default)] +#[derive(Default, Debug)] pub(crate) struct MissingTermAggReqData { /// The accessors to check for existence of a value. pub(crate) accessors: Vec<(Column, ColumnType)>, @@ -45,7 +43,7 @@ struct MissingCount { /// The specialized missing term aggregation. #[derive(Default, Debug)] pub struct TermMissingAgg { - accessor_idx: usize, + req_data: MissingTermAggReqData, sub_agg: Option, /// Idx = parent bucket id, Value = missing count for that bucket missing_count_per_bucket: Vec, @@ -53,27 +51,15 @@ pub struct TermMissingAgg { } impl TermMissingAgg { pub(crate) fn new( - agg_data: &mut AggregationsSegmentCtx, - node: &AggRefNode, - ) -> crate::Result { - let has_sub_aggregations = !node.children.is_empty(); - let accessor_idx = node.idx_in_req_data; - let sub_agg = if has_sub_aggregations { - let sub_aggregation = build_segment_agg_collectors(agg_data, &node.children)?; - Some(sub_aggregation) - } else { - None - }; - - let sub_agg = sub_agg.map(BufferedSubAggs::new); - let bucket_id_provider = BucketIdProvider::default(); - - Ok(Self { - accessor_idx, - sub_agg, + req_data: MissingTermAggReqData, + sub_agg: Option>, + ) -> Self { + Self { + req_data, + sub_agg: sub_agg.map(BufferedSubAggs::new), missing_count_per_bucket: Vec::new(), - bucket_id_provider, - }) + bucket_id_provider: BucketIdProvider::default(), + } } } @@ -85,7 +71,7 @@ impl SegmentAggregationCollector for TermMissingAgg { parent_bucket_id: BucketId, ) -> crate::Result<()> { self.prepare_max_bucket(parent_bucket_id, agg_data)?; - let req_data = agg_data.get_missing_term_req_data(self.accessor_idx); + let req_data = &self.req_data; let term_agg = &req_data.req; let missing = term_agg .missing @@ -130,7 +116,7 @@ impl SegmentAggregationCollector for TermMissingAgg { agg_data: &mut AggregationsSegmentCtx, ) -> crate::Result<()> { let bucket = &mut self.missing_count_per_bucket[parent_bucket_id as usize]; - let req_data = agg_data.get_missing_term_req_data(self.accessor_idx); + let req_data = &self.req_data; for doc in docs { let doc = *doc; diff --git a/src/aggregation/collector.rs b/src/aggregation/collector.rs index c1a825786..b3bdef927 100644 --- a/src/aggregation/collector.rs +++ b/src/aggregation/collector.rs @@ -5,7 +5,7 @@ use super::intermediate_agg_result::IntermediateAggregationResults; use super::AggContextParams; // group buffering strategy is chosen explicitly by callers; no need to hash-group on the fly. use crate::aggregation::agg_data::{ - build_aggregations_data_from_req, build_segment_agg_collectors_root, AggregationsSegmentCtx, + build_aggregations_data_from_req, build_segment_agg_collectors, AggregationsSegmentCtx, }; use crate::collector::{Collector, SegmentCollector}; use crate::index::SegmentReader; @@ -149,10 +149,10 @@ impl AggregationSegmentCollector { segment_ordinal: SegmentOrdinal, context: &AggContextParams, ) -> crate::Result { - let mut agg_data = - build_aggregations_data_from_req(agg, reader, segment_ordinal, context.clone())?; + let agg_tree = build_aggregations_data_from_req(agg, reader, segment_ordinal, context)?; + let mut agg_data = AggregationsSegmentCtx::new(context.clone()); let mut result = - LowCardBufferedSubAggs::new(build_segment_agg_collectors_root(&mut agg_data)?); + LowCardBufferedSubAggs::new(build_segment_agg_collectors(&mut agg_data, agg_tree)?); result .get_sub_agg_collector() .prepare_max_bucket(0, &agg_data)?; // prepare for bucket zero diff --git a/src/aggregation/metric/cardinality/mod.rs b/src/aggregation/metric/cardinality/mod.rs index 8440cd948..c164ad388 100644 --- a/src/aggregation/metric/cardinality/mod.rs +++ b/src/aggregation/metric/cardinality/mod.rs @@ -16,7 +16,6 @@ mod str_collector; mod term_ord_accumulator; use std::hash::Hash; -use std::sync::Arc; use columnar::{ColumnType, StrColumn}; use common::BitSet; @@ -26,7 +25,6 @@ use serde::{Deserialize, Deserializer, Serialize, Serializer}; pub(crate) use str_collector::SegmentStrCardinalityCollector; pub(crate) use term_ord_accumulator::{TermOrdSet, BITSET_MAX_TERM_ORD}; -use crate::aggregation::agg_data::{AggRefNode, AggregationsSegmentCtx}; use crate::aggregation::segment_agg_result::SegmentAggregationCollector; use crate::aggregation::value_source::ValueSource; use crate::aggregation::*; @@ -102,7 +100,7 @@ pub struct CardinalityAggregationReq { /// cardinality aggregation on a segment. pub(crate) struct CardinalityAggReqData { /// The column accessor to access the fast field values. - pub(crate) accessor: Arc, + pub(crate) accessor: Box, /// The string dictionary column if the field is of type string. pub(crate) str_dict_column: Option, /// The missing value normalized to the internal u64 representation of the field type. @@ -220,15 +218,11 @@ impl CardinalityCollector { /// finalization, non-str feeds the HLL sketch directly. Both produce the same /// [`IntermediateMetricResult::Cardinality`], so they merge uniformly. pub(crate) fn build_segment_cardinality_collector( - req: &mut AggregationsSegmentCtx, - node: &AggRefNode, + req_data: CardinalityAggReqData, ) -> crate::Result> { - let req_data = req.get_cardinality_req_data(node.idx_in_req_data); if req_data.accessor.column_type() != ColumnType::Str { return Ok(Box::new(SegmentNumericCardinalityCollector::from_req( - node.idx_in_req_data, - req_data.accessor.clone(), - req_data.missing_value_for_accessor, + req_data, )?)); } // For str columns, we need to collect the set of term ordinals encounterred. @@ -244,19 +238,12 @@ pub(crate) fn build_segment_cardinality_collector( let max_term_ord_inclusive = column.max_value(); if max_term_ord_inclusive < BITSET_MAX_TERM_ORD { Ok(Box::new( - SegmentStrCardinalityCollector::::from_req( - node.idx_in_req_data, - req_data.accessor.clone(), - req_data.missing_value_for_accessor, - max_term_ord_inclusive, - ), + SegmentStrCardinalityCollector::::from_req(req_data, max_term_ord_inclusive), )) } else { Ok(Box::new( SegmentStrCardinalityCollector::::from_req( - node.idx_in_req_data, - req_data.accessor.clone(), - req_data.missing_value_for_accessor, + req_data, max_term_ord_inclusive, ), )) diff --git a/src/aggregation/metric/cardinality/numeric_collector.rs b/src/aggregation/metric/cardinality/numeric_collector.rs index d6b3fa7bc..d88871653 100644 --- a/src/aggregation/metric/cardinality/numeric_collector.rs +++ b/src/aggregation/metric/cardinality/numeric_collector.rs @@ -12,7 +12,7 @@ use std::sync::Arc; use columnar::column_values::CompactSpaceU64Accessor; use columnar::ColumnType; -use super::CardinalityCollector; +use super::{CardinalityAggReqData, CardinalityCollector}; use crate::aggregation::agg_data::AggregationsSegmentCtx; use crate::aggregation::intermediate_agg_result::{ IntermediateAggregationResult, IntermediateAggregationResults, IntermediateMetricResult, @@ -32,9 +32,10 @@ pub(crate) struct SegmentNumericCardinalityCollector { /// Buckets are Some(_) until they get consumed by /// `add_intermediate_aggregation_result`. buckets: Vec>, - accessor_idx: usize, + /// The name of the aggregation. + name: String, /// The column accessor to access the fast field values. - accessor: Arc, + accessor: Box, /// The column_type of the field. column_type: ColumnType, /// Set iff `column_type == ColumnType::IpAddr`. Resolved once at @@ -58,11 +59,13 @@ impl Debug for SegmentNumericCardinalityCollector { } impl SegmentNumericCardinalityCollector { - pub fn from_req( - accessor_idx: usize, - accessor: Arc, - missing_value_for_accessor: Option, - ) -> crate::Result { + pub(crate) fn from_req(req_data: CardinalityAggReqData) -> crate::Result { + let CardinalityAggReqData { + accessor, + missing_value_for_accessor, + name, + .. + } = req_data; let column_type = accessor.column_type(); assert_ne!(column_type, ColumnType::Str); let compact_space_accessor = if column_type == ColumnType::IpAddr { @@ -92,7 +95,7 @@ impl SegmentNumericCardinalityCollector { }; Ok(Self { buckets: Vec::new(), - accessor_idx, + name, accessor, column_type, compact_space_accessor, @@ -109,10 +112,7 @@ impl SegmentAggregationCollector for SegmentNumericCardinalityCollector { bucket_id: BucketId, ) -> crate::Result<()> { self.prepare_max_bucket(bucket_id, agg_data)?; - let name = agg_data - .get_cardinality_req_data(self.accessor_idx) - .name - .to_string(); + let name = self.name.clone(); // take the bucket in buckets and replace it with a new empty one let Some(cardinality) = self.buckets[bucket_id as usize].take() else { return Err(crate::TantivyError::InternalError( @@ -136,7 +136,7 @@ impl SegmentAggregationCollector for SegmentNumericCardinalityCollector { ) -> crate::Result<()> { agg_data.column_block_accessor.fetch_block_with_missing( docs, - &*self.accessor, + &mut *self.accessor, self.missing_value_for_accessor, ); let cardinality = self.buckets[parent_bucket_id as usize] @@ -179,10 +179,9 @@ impl SegmentAggregationCollector for SegmentNumericCardinalityCollector { bucket_id: BucketId, sub_agg_name: &str, sub_agg_property: &str, - agg_data: &AggregationsSegmentCtx, + _agg_data: &AggregationsSegmentCtx, ) -> Option { - let req_data = &agg_data.get_cardinality_req_data(self.accessor_idx); - if req_data.name != sub_agg_name || !sub_agg_property.is_empty() { + if self.name != sub_agg_name || !sub_agg_property.is_empty() { return None; } let cardinality = self.buckets.get(bucket_id as usize)?.as_ref()?; diff --git a/src/aggregation/metric/cardinality/str_collector.rs b/src/aggregation/metric/cardinality/str_collector.rs index b2ec7d218..4bfa61d6c 100644 --- a/src/aggregation/metric/cardinality/str_collector.rs +++ b/src/aggregation/metric/cardinality/str_collector.rs @@ -10,20 +10,18 @@ use std::fmt::Debug; use std::io; -use std::sync::Arc; use columnar::{ColumnType, Dictionary}; use datasketches::hll::Coupon; use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet}; use super::term_ord_accumulator::TermOrdAccumulator; -use super::CardinalityCollector; +use super::{CardinalityAggReqData, CardinalityCollector}; use crate::aggregation::agg_data::AggregationsSegmentCtx; use crate::aggregation::intermediate_agg_result::{ IntermediateAggregationResult, IntermediateAggregationResults, IntermediateMetricResult, }; use crate::aggregation::segment_agg_result::SegmentAggregationCollector; -use crate::aggregation::value_source::ValueSource; use crate::aggregation::*; /// A CouponCache is here to cache the mapping term ordinal -> coupon (see above). @@ -99,11 +97,8 @@ pub(crate) struct SegmentStrCardinalityCollector { /// Buckets are Some(_) until they get consumed by /// `add_intermediate_aggregation_result`. buckets: Vec>, - accessor_idx: usize, - /// The column accessor to access the fast field values (term ordinals). - accessor: Arc, - /// The missing value normalized to the internal u64 representation of the field type. - missing_value_for_accessor: Option, + /// The request data. Its accessor gives access to the term ordinals. + req_data: CardinalityAggReqData, /// Lazily built at finalization time, shared by every bucket. coupon_cache: Option, /// Largest term_ord that may be inserted into a bucket, i.e. @@ -117,7 +112,7 @@ impl Debug for SegmentStrCardinalityCollector { .field("num_buckets", &self.buckets.len()) .field( "missing_value_for_accessor", - &self.missing_value_for_accessor, + &self.req_data.missing_value_for_accessor, ) .finish() } @@ -209,17 +204,10 @@ fn append_to_sketch( } impl SegmentStrCardinalityCollector { - pub fn from_req( - accessor_idx: usize, - accessor: Arc, - missing_value_for_accessor: Option, - max_term_ord_inclusive: u64, - ) -> Self { + pub(crate) fn from_req(req_data: CardinalityAggReqData, max_term_ord_inclusive: u64) -> Self { Self { buckets: Vec::new(), - accessor_idx, - accessor, - missing_value_for_accessor, + req_data, coupon_cache: None, max_term_ord_inclusive, } @@ -236,7 +224,7 @@ impl SegmentAggregationCollector bucket_id: BucketId, ) -> crate::Result<()> { self.prepare_max_bucket(bucket_id, agg_data)?; - let req_data = &agg_data.get_cardinality_req_data(self.accessor_idx); + let req_data = &self.req_data; let Some(str_dict_column) = &req_data.str_dict_column else { return Err(crate::TantivyError::InternalError( "a str cardinality collector requires a str dictionary column".to_string(), @@ -284,8 +272,8 @@ impl SegmentAggregationCollector ) -> crate::Result<()> { agg_data.column_block_accessor.fetch_block_with_missing( docs, - &*self.accessor, - self.missing_value_for_accessor, + &mut *self.req_data.accessor, + self.req_data.missing_value_for_accessor, ); let Some(term_ords) = self.buckets[parent_bucket_id as usize].as_mut() else { return Err(crate::TantivyError::InternalError( @@ -322,10 +310,9 @@ impl SegmentAggregationCollector bucket_id: BucketId, sub_agg_name: &str, sub_agg_property: &str, - agg_data: &AggregationsSegmentCtx, + _agg_data: &AggregationsSegmentCtx, ) -> Option { - let req_data = &agg_data.get_cardinality_req_data(self.accessor_idx); - if req_data.name != sub_agg_name || !sub_agg_property.is_empty() { + if self.req_data.name != sub_agg_name || !sub_agg_property.is_empty() { return None; } // The sketch isn't built until finalization; the term_ord set's len is diff --git a/src/aggregation/metric/extended_stats.rs b/src/aggregation/metric/extended_stats.rs index 3747a0921..78e28c8ee 100644 --- a/src/aggregation/metric/extended_stats.rs +++ b/src/aggregation/metric/extended_stats.rs @@ -1,6 +1,5 @@ use std::fmt::Debug; use std::mem; -use std::sync::Arc; use serde::{Deserialize, Serialize}; @@ -318,23 +317,23 @@ impl IntermediateExtendedStats { } } -#[derive(Clone, Debug)] +#[derive(Debug)] pub(crate) struct SegmentExtendedStatsCollector { name: String, missing: Option, - accessor: Arc, + accessor: Box, buckets: Vec, sigma: Option, } impl SegmentExtendedStatsCollector { - pub fn from_req(req: &MetricAggReqData, sigma: Option) -> Self { + pub(crate) fn from_req(req: MetricAggReqData, sigma: Option) -> Self { let missing = req .missing .and_then(|val| f64_to_fastfield_u64(val, &req.accessor.column_type())); Self { - name: req.name.clone(), - accessor: req.accessor.clone(), + name: req.name, + accessor: req.accessor, missing, buckets: vec![IntermediateExtendedStats::with_sigma(sigma); 16], sigma, @@ -374,7 +373,7 @@ impl SegmentAggregationCollector for SegmentExtendedStatsCollector { agg_data.column_block_accessor.fetch_block_with_missing( docs, - &*self.accessor, + &mut *self.accessor, self.missing, ); let field_type = self.accessor.column_type(); diff --git a/src/aggregation/metric/mod.rs b/src/aggregation/metric/mod.rs index 16197a7a2..732ee5e2f 100644 --- a/src/aggregation/metric/mod.rs +++ b/src/aggregation/metric/mod.rs @@ -28,7 +28,6 @@ mod sum; mod top_hits; use std::collections::HashMap; -use std::sync::Arc; pub use average::*; pub use cardinality::*; @@ -55,7 +54,7 @@ pub(crate) struct MetricAggReqData { /// The missing value normalized to the internal u64 representation of the field type. pub(crate) missing_u64: Option, /// The column accessor to access the fast field values. - pub(crate) accessor: Arc, + pub(crate) accessor: Box, /// Used when converting to intermediate result pub(crate) collecting_for: StatsType, /// The missing value diff --git a/src/aggregation/metric/percentiles.rs b/src/aggregation/metric/percentiles.rs index e2817617e..67b865222 100644 --- a/src/aggregation/metric/percentiles.rs +++ b/src/aggregation/metric/percentiles.rs @@ -1,5 +1,4 @@ use std::fmt::Debug; -use std::sync::Arc; use serde::{Deserialize, Serialize}; @@ -131,14 +130,15 @@ impl PercentilesAggregationReq { } } -#[derive(Clone, Debug)] +#[derive(Debug)] pub(crate) struct SegmentPercentilesCollector { pub(crate) buckets: Vec, - pub(crate) accessor_idx: usize, + /// The name of the aggregation. + pub(crate) name: String, /// The missing value normalized to the internal u64 representation of the field type. pub missing_u64: Option, /// The column accessor to access the fast field values. - pub(crate) accessor: Arc, + pub(crate) accessor: Box, } #[derive(Clone, Serialize, Deserialize)] @@ -248,16 +248,12 @@ impl PercentilesCollector { } impl SegmentPercentilesCollector { - pub fn from_req_and_validate( - missing_u64: Option, - accessor: Arc, - accessor_idx: usize, - ) -> Self { + pub(crate) fn from_req_and_validate(req_data: MetricAggReqData) -> Self { Self { buckets: Vec::with_capacity(64), - missing_u64, - accessor, - accessor_idx, + name: req_data.name, + missing_u64: req_data.missing_u64, + accessor: req_data.accessor, } } } @@ -270,7 +266,7 @@ impl SegmentAggregationCollector for SegmentPercentilesCollector { results: &mut IntermediateAggregationResults, parent_bucket_id: BucketId, ) -> crate::Result<()> { - let name = agg_data.get_metric_req_data(self.accessor_idx).name.clone(); + let name = self.name.clone(); self.prepare_max_bucket(parent_bucket_id, agg_data)?; // Swap collector with an empty one to avoid cloning let percentiles_collector = std::mem::take(&mut self.buckets[parent_bucket_id as usize]); @@ -296,7 +292,7 @@ impl SegmentAggregationCollector for SegmentPercentilesCollector { let percentiles = &mut self.buckets[parent_bucket_id as usize]; agg_data.column_block_accessor.fetch_block_with_missing( docs, - &*self.accessor, + &mut *self.accessor, self.missing_u64, ); @@ -325,9 +321,9 @@ impl SegmentAggregationCollector for SegmentPercentilesCollector { bucket_id: BucketId, sub_agg_name: &str, sub_agg_property: &str, - agg_data: &AggregationsSegmentCtx, + _agg_data: &AggregationsSegmentCtx, ) -> Option { - if agg_data.get_metric_req_data(self.accessor_idx).name != sub_agg_name { + if self.name != sub_agg_name { return None; } let percentile: f64 = sub_agg_property.parse().ok()?; diff --git a/src/aggregation/metric/stats.rs b/src/aggregation/metric/stats.rs index 30ec18e33..b6aa87ecf 100644 --- a/src/aggregation/metric/stats.rs +++ b/src/aggregation/metric/stats.rs @@ -1,5 +1,4 @@ use std::fmt::Debug; -use std::sync::Arc; use columnar::{Column, ColumnType}; use serde::{Deserialize, Serialize}; @@ -199,25 +198,26 @@ pub enum StatsType { } fn create_collector( - req: &MetricAggReqData, + req: MetricAggReqData, ) -> Box { + let column_opt = req + .accessor + .as_column() + .map(|column| (column.clone(), req.accessor.column_type())); Box::new(SegmentStatsCollector:: { - name: req.name.clone(), + name: req.name, collecting_for: req.collecting_for, is_number_or_date_type: req.is_number_or_date_type, missing_u64: req.missing_u64, - column_opt: req - .accessor - .as_column() - .map(|column| (column.clone(), req.accessor.column_type())), - accessor: req.accessor.clone(), + column_opt, + accessor: req.accessor, buckets: vec![IntermediateStats::default()], }) } /// Build a concrete `SegmentStatsCollector` depending on the column type. pub(crate) fn build_segment_stats_collector( - req: &MetricAggReqData, + req: MetricAggReqData, ) -> crate::Result> { match req.accessor.column_type() { ColumnType::I64 => Ok(create_collector::<{ ColumnType::I64 as u8 }>(req)), @@ -232,10 +232,10 @@ pub(crate) fn build_segment_stats_collector( } #[repr(C)] -#[derive(Clone, Debug)] +#[derive(Debug)] pub(crate) struct SegmentStatsCollector { pub(crate) missing_u64: Option, - pub(crate) accessor: Arc, + pub(crate) accessor: Box, /// The physical column backing `accessor`, if any, resolved once at construction. /// /// `collect` is called once per bucket, often with a single doc (e.g. under a @@ -305,7 +305,7 @@ impl SegmentAggregationCollector // access, which a computed source cannot offer. Those fall through to the block path // below, which is semantically identical. // TODO: remove once we fetch all values for all bucket ids in one go - if let Some(column_source) = &self.column_opt { + if let Some(column_source) = &mut self.column_opt { if docs.len() == 1 && self.missing_u64.is_none() { collect_stats::( &mut self.buckets[parent_bucket_id as usize], @@ -322,7 +322,7 @@ impl SegmentAggregationCollector } else { agg_data.column_block_accessor.fetch_block_with_missing( docs, - &*self.accessor, + &mut *self.accessor, self.missing_u64, ); } diff --git a/src/aggregation/metric/top_hits.rs b/src/aggregation/metric/top_hits.rs index 4cf1cf4eb..e5f289055 100644 --- a/src/aggregation/metric/top_hits.rs +++ b/src/aggregation/metric/top_hits.rs @@ -516,34 +516,35 @@ impl TopHitsTopNComputer { } } -#[derive(Clone, Debug)] pub(crate) struct TopHitsSegmentCollector { - segment_ordinal: SegmentOrdinal, - accessor_idx: usize, + req_data: TopHitsAggReqData, buckets: Vec, DocAddress, ReverseComparator>>, num_hits: usize, } +impl std::fmt::Debug for TopHitsSegmentCollector { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TopHitsSegmentCollector") + .field("name", &self.req_data.name) + .field("segment_ordinal", &self.req_data.segment_ordinal) + .field("num_buckets", &self.buckets.len()) + .field("num_hits", &self.num_hits) + .finish() + } +} + impl TopHitsSegmentCollector { - pub fn from_req( - req: &TopHitsAggregationReq, - accessor_idx: usize, - segment_ordinal: SegmentOrdinal, - ) -> Self { - let num_hits = req.size + req.from.unwrap_or(0); + pub(crate) fn from_req(req_data: TopHitsAggReqData) -> Self { + let num_hits = req_data.req.size + req_data.req.from.unwrap_or(0); Self { num_hits, - segment_ordinal, - accessor_idx, + req_data, buckets: vec![TopNComputer::new_with_comparator(num_hits, ReverseComparator); 1], } } - fn get_top_hits_computer( - &mut self, - parent_bucket_id: BucketId, - value_accessors: &HashMap>, - req: &TopHitsAggregationReq, - ) -> TopHitsTopNComputer { + fn get_top_hits_computer(&mut self, parent_bucket_id: BucketId) -> TopHitsTopNComputer { + let req = &self.req_data.req; + let value_accessors = &self.req_data.value_accessors; if parent_bucket_id as usize >= self.buckets.len() { return TopHitsTopNComputer::new(req); } @@ -572,21 +573,14 @@ impl TopHitsSegmentCollector { impl SegmentAggregationCollector for TopHitsSegmentCollector { fn add_intermediate_aggregation_result( &mut self, - agg_data: &AggregationsSegmentCtx, + _agg_data: &AggregationsSegmentCtx, results: &mut crate::aggregation::intermediate_agg_result::IntermediateAggregationResults, parent_bucket_id: BucketId, ) -> crate::Result<()> { - let req_data = agg_data.get_top_hits_req_data(self.accessor_idx); - - let value_accessors = &req_data.value_accessors; - - let intermediate_result = IntermediateMetricResult::TopHits(self.get_top_hits_computer( - parent_bucket_id, - value_accessors, - &req_data.req, - )); + let intermediate_result = + IntermediateMetricResult::TopHits(self.get_top_hits_computer(parent_bucket_id)); results.push( - req_data.name.to_string(), + self.req_data.name.to_string(), IntermediateAggregationResult::Metric(intermediate_result), ) } @@ -596,10 +590,10 @@ impl SegmentAggregationCollector for TopHitsSegmentCollector { &mut self, parent_bucket_id: BucketId, docs: &[crate::DocId], - agg_data: &mut AggregationsSegmentCtx, + _agg_data: &mut AggregationsSegmentCtx, ) -> crate::Result<()> { let top_n = &mut self.buckets[parent_bucket_id as usize]; - let req_data = agg_data.get_top_hits_req_data(self.accessor_idx); + let req_data = &self.req_data; let req = &req_data.req; let accessors = &req_data.accessors; for &doc_id in docs { @@ -625,7 +619,7 @@ impl SegmentAggregationCollector for TopHitsSegmentCollector { top_n.push( sorts, DocAddress { - segment_ord: self.segment_ordinal, + segment_ord: req_data.segment_ordinal, doc_id, }, ); diff --git a/src/aggregation/value_source/block_accessor.rs b/src/aggregation/value_source/block_accessor.rs index ccc6b9951..a77790254 100644 --- a/src/aggregation/value_source/block_accessor.rs +++ b/src/aggregation/value_source/block_accessor.rs @@ -40,7 +40,7 @@ pub(crate) struct ColumnBlockAccessor { impl ColumnBlockAccessor { #[inline] - pub(crate) fn fetch_block(&mut self, docs: &[DocId], source: &S) { + pub(crate) fn fetch_block(&mut self, docs: &[DocId], source: &mut S) { self.cardinality = source.load_block( docs, &mut self.val_cache, @@ -72,7 +72,7 @@ impl ColumnBlockAccessor { pub(crate) fn fetch_block_with_missing( &mut self, docs: &[DocId], - source: &S, + source: &mut S, missing_opt: Option, ) { self.fetch_block_with_missing_ordered(docs, source, missing_opt, false) @@ -85,7 +85,7 @@ impl ColumnBlockAccessor { pub(crate) fn fetch_block_with_missing_ordered( &mut self, docs: &[DocId], - source: &S, + source: &mut S, missing_opt: Option, ordered: bool, ) { @@ -155,7 +155,7 @@ impl ColumnBlockAccessor { pub(crate) fn fetch_block_with_missing_unique_per_doc( &mut self, docs: &[DocId], - source: &dyn ValueSource, + source: &mut dyn ValueSource, missing: Option, ordered: bool, ) { @@ -304,7 +304,6 @@ fn find_missing_docs(docs: &[u32], hits: &[u32], output: &mut Vec) { #[cfg(test)] #[allow(clippy::field_reassign_with_default)] mod tests { - use std::sync::Arc; use columnar::{Column, ColumnType}; @@ -322,7 +321,7 @@ mod tests { } fn load_block( - &self, + &mut self, docs: &[DocId], values: &mut Vec, docids: &mut Vec, @@ -345,11 +344,11 @@ mod tests { #[test] fn test_fetch_block_accepts_trait_object() { let docs = [2, 4, 8]; - let source = TestValueSource { + let mut source = TestValueSource { cardinality: Cardinality::Full, entries: vec![(2, 20), (4, 40), (8, 80)], }; - let dyn_source: &dyn ValueSource = &source; + let dyn_source: &mut dyn ValueSource = &mut source; let mut accessor = ColumnBlockAccessor::default(); accessor.fetch_block(&docs, dyn_source); @@ -374,11 +373,11 @@ mod tests { #[test] fn test_as_column_distinguishes_the_two_kinds() { - let column: Arc = Arc::new((full_column(&[5, 6, 7]), ColumnType::U64)); + let column: Box = Box::new((full_column(&[5, 6, 7]), ColumnType::U64)); assert!(column.as_column().is_some()); assert_eq!(column.bounds(), Some((5, 7))); - let computed: Arc = Arc::new(TestValueSource { + let computed: Box = Box::new(TestValueSource { cardinality: Cardinality::Full, entries: vec![(0, 1)], }); @@ -392,20 +391,20 @@ mod tests { // A computed source reports what it produced: docs 0 and 2 have no value, so they are // absent from `docids` and the source is `Optional`. let docs = [0, 1, 2, 3]; - let computed: Arc = Arc::new(TestValueSource { + let mut computed: Box = Box::new(TestValueSource { cardinality: Cardinality::Optional, entries: vec![(1, 11), (3, 33)], }); let mut accessor = ColumnBlockAccessor::default(); - accessor.fetch_block(&docs, &*computed); + accessor.fetch_block(&docs, &mut *computed); assert!(!accessor.has_one_value_per_doc(&docs)); assert_eq!( accessor.iter_docid_vals(&docs).collect::>(), [(1, 11), (3, 33)] ); - accessor.fetch_block_with_missing(&docs, &*computed, Some(99)); + accessor.fetch_block_with_missing(&docs, &mut *computed, Some(99)); let mut pairs = accessor.iter_docid_vals(&docs).collect::>(); pairs.sort_unstable(); assert_eq!(pairs, [(0, 99), (1, 11), (2, 99), (3, 33)]); @@ -442,12 +441,12 @@ mod tests { #[test] fn test_source_neutral_full_block_alignment() { let docs = [2, 4, 8]; - let source = TestValueSource { + let mut source = TestValueSource { cardinality: Cardinality::Full, entries: vec![(2, 20), (4, 40), (8, 80)], }; let mut accessor = ColumnBlockAccessor::default(); - accessor.fetch_block(&docs, &source); + accessor.fetch_block(&docs, &mut source); assert!(accessor.has_one_value_per_doc(&docs)); assert_eq!( accessor.iter_docid_vals(&docs).collect::>(), @@ -458,13 +457,13 @@ mod tests { #[test] fn test_source_neutral_optional_block_with_missing() { let docs = [0, 1, 2, 4]; - let source = TestValueSource { + let mut source = TestValueSource { cardinality: Cardinality::Optional, entries: vec![(1, 10), (4, 40)], }; let mut accessor = ColumnBlockAccessor::default(); - accessor.fetch_block_with_missing_ordered(&docs, &source, Some(99), true); + accessor.fetch_block_with_missing_ordered(&docs, &mut source, Some(99), true); assert!(accessor.has_one_value_per_doc(&docs)); assert_eq!( @@ -476,13 +475,13 @@ mod tests { #[test] fn test_source_neutral_multivalue_block_deduplication() { let docs = [0, 1]; - let source = TestValueSource { + let mut source = TestValueSource { cardinality: Cardinality::Multivalued, entries: vec![(0, 3), (0, 1), (0, 3), (1, 5), (1, 5)], }; let mut accessor = ColumnBlockAccessor::default(); - accessor.fetch_block_with_missing_unique_per_doc(&docs, &source, None, false); + accessor.fetch_block_with_missing_unique_per_doc(&docs, &mut source, None, false); assert!(!accessor.has_one_value_per_doc(&docs)); assert_eq!( @@ -510,7 +509,7 @@ mod tests { accessor.fetch_block_with_missing_ordered( &docs, - &(&column, ColumnType::U64), + &mut (&column, ColumnType::U64), Some(99), true, ); @@ -593,7 +592,7 @@ mod tests { }; let check = |accessor: &mut ColumnBlockAccessor, docs: &[u32]| { - accessor.fetch_block(docs, &(&column, ColumnType::U64)); + accessor.fetch_block(docs, &mut (&column, ColumnType::U64)); let got: Vec<(u32, u64)> = accessor.iter_docid_vals(docs).collect(); let expected: Vec<(u32, u64)> = docs.iter().map(|&d| (d, vals[d as usize])).collect(); assert_eq!(got, expected); diff --git a/src/aggregation/value_source/mod.rs b/src/aggregation/value_source/mod.rs index 12fec6a41..8b937e684 100644 --- a/src/aggregation/value_source/mod.rs +++ b/src/aggregation/value_source/mod.rs @@ -35,8 +35,11 @@ pub trait ValueSource: std::fmt::Debug { /// `docids[i]`. `docids` only contains docs from `docs`. A doc is repeated once per value. /// /// `row_ids` is scratch the implementation may use freely. + /// + /// Takes `&mut self` so that implementations can keep per-segment state (caches, scratch + /// buffers) across blocks. Each source is owned by a single collector. fn load_block( - &self, + &mut self, docs: &[DocId], values: &mut Vec, docids: &mut Vec, @@ -64,7 +67,7 @@ impl> + std::fmt::Debug> ValueSource for (ColumnRe #[inline] fn load_block( - &self, + &mut self, docs: &[DocId], values: &mut Vec, docids: &mut Vec, diff --git a/src/aggregation/value_source/tests.rs b/src/aggregation/value_source/tests.rs index b2a9c3b4f..63589e138 100644 --- a/src/aggregation/value_source/tests.rs +++ b/src/aggregation/value_source/tests.rs @@ -14,7 +14,7 @@ impl ValueSource for Constant { } fn load_block( - &self, + &mut self, docs: &[DocId], values: &mut Vec, _docids: &mut Vec, @@ -29,8 +29,8 @@ impl ValueSource for Constant { pub(crate) struct ConstantProvider(pub u64); impl ValueSourceProvider for ConstantProvider { - fn for_segment(&self, _reader: &SegmentReader) -> crate::Result> { - Ok(Arc::new(Constant(self.0))) + fn for_segment(&self, _reader: &SegmentReader) -> crate::Result> { + Ok(Box::new(Constant(self.0))) } } fn index_with_scores(scores: &[u64]) -> crate::Index { diff --git a/src/aggregation/value_source/value_source_registry.rs b/src/aggregation/value_source/value_source_registry.rs index f9888a4e6..89d6cde88 100644 --- a/src/aggregation/value_source/value_source_registry.rs +++ b/src/aggregation/value_source/value_source_registry.rs @@ -9,7 +9,7 @@ use crate::SegmentReader; /// Creates a value source for each segment. pub trait ValueSourceProvider: Send + Sync + 'static { /// Binds this definition to a single segment. - fn for_segment(&self, reader: &SegmentReader) -> crate::Result>; + fn for_segment(&self, reader: &SegmentReader) -> crate::Result>; } /// Named computed sources available to an aggregation request. @@ -58,7 +58,7 @@ mod tests { writer.commit().unwrap(); let searcher = index.reader().unwrap().searcher(); let value_source_provider = registry.get("computed").unwrap(); - let value_source = value_source_provider + let mut value_source = value_source_provider .for_segment(searcher.segment_reader(0u32)) .unwrap(); let mut values = Vec::new();