rust/tantivy
1
0
Fork 0
mirror of https://github.com/quickwit-oss/tantivy.git synced 2026-08-18 12:08:22 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
864a1d4b2dce76917ecef187da90d8bda2ca73f7
tantivy/src/aggregation/agg_data.rs
T
Pascal Seitz 864a1d4b2d improve missing performance
2026-08-04 13:23:19 +02:00

1367 lines
52 KiB
Rust
Raw Blame History

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