From 39198bcf1f5fc323c6fd21089ff41c667ac8a590 Mon Sep 17 00:00:00 2001 From: Paul Masurel Date: Fri, 2 Oct 2026 11:37:39 +0200 Subject: [PATCH] Make ValueSource::load_block take &mut self Value sources were held as `Arc` and cloned into their collector, so `load_block` could only take `&self`. A computed source could not keep per-segment state (caches, scratch buffers) across blocks. Collectors were built in two phases: `build_nodes` pushed each `XxxAggReqData` into per-kind vectors of `AggregationsSegmentCtx`, and each collector builder cloned its request data out of those vectors by index. The ctx copy had to stay around for name lookups, memory accounting, and for the collectors that read their data back by index. Each request node is turned into exactly one collector, so nothing actually shares a value source. This change makes that ownership explicit: - `ValueSource::load_block` takes `&mut self`, `ValueSourceProvider::for_segment` returns `Box`. - The request tree (`AggNode { data: AggNodeData, children }`) owns the request data. Building collectors consumes the tree, moving the request data into the collectors instead of cloning it. - `AggregationsSegmentCtx` only keeps the shared collect-time state (context and block accessor). `PerRequestAggSegCtx`, `AggKind`, `idx_in_req_data` and the `push_*`/`get_*_req_data` helpers are gone. - Cardinality, percentiles, top_hits and missing_term collectors own their request data instead of reading it from the ctx by index. - Histogram requests are normalized once, when building the tree. The flattened terms x histogram path is split into a borrow-only plan step and a consuming build step, and no longer clones the histogram request at finalization. - Request data memory is charged exactly once per node. Previously it was charged at every nesting level, and range, histogram, filter and composite builders charged their request data a second time. - `FilterAggReqData::evaluator` no longer needs an `Rc`, and the filter, composite and multi_terms request data no longer derive Clone. --- src/aggregation/accessor_helpers.rs | 11 +- src/aggregation/agg_data.rs | 756 +++++++----------- src/aggregation/agg_limits.rs | 10 +- src/aggregation/bucket/composite/accessors.rs | 2 +- src/aggregation/bucket/composite/collector.rs | 23 +- src/aggregation/bucket/filter.rs | 53 +- src/aggregation/bucket/histogram/histogram.rs | 192 ++--- src/aggregation/bucket/multi_terms/mod.rs | 53 +- src/aggregation/bucket/range.rs | 26 +- .../term_agg/flattened_term_histogram.rs | 147 ++-- src/aggregation/bucket/term_agg/mod.rs | 49 +- src/aggregation/bucket/term_missing_agg.rs | 40 +- src/aggregation/collector.rs | 8 +- src/aggregation/metric/cardinality/mod.rs | 23 +- .../metric/cardinality/numeric_collector.rs | 33 +- .../metric/cardinality/str_collector.rs | 35 +- src/aggregation/metric/extended_stats.rs | 13 +- src/aggregation/metric/mod.rs | 3 +- src/aggregation/metric/percentiles.rs | 28 +- src/aggregation/metric/stats.rs | 26 +- src/aggregation/metric/top_hits.rs | 56 +- .../value_source/block_accessor.rs | 41 +- src/aggregation/value_source/mod.rs | 7 +- src/aggregation/value_source/tests.rs | 6 +- .../value_source/value_source_registry.rs | 4 +- 25 files changed, 699 insertions(+), 946 deletions(-) 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();