rust/tantivy
1
0
Fork 0
mirror of https://github.com/quickwit-oss/tantivy.git synced 2026-10-23 04:02:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

Make ValueSource::load_block take &mut self

Browse Source 
Value sources were held as `Arc<dyn ValueSource>` 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<dyn ValueSource>`.
- 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.
This commit is contained in:
Paul Masurel
2026-10-02 18:37:49 +02:00
parent 4abc20f924
commit 39198bcf1f
25 changed files with 699 additions and 946 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/aggregation/accessor_helpers.rs
+5 -6
View File
@@ -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<Option<Arc<dyn ValueSource>>> {
) -> crate::Result<Option<Box<dyn ValueSource>>> {
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<Arc<dyn ValueSource>> {
) -> crate::Result<Box<dyn ValueSource>> {
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<Vec<Arc<dyn ValueSource>>> {
) -> crate::Result<Vec<Box<dyn ValueSource>>> {
// 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<dyn ValueSource> = Arc::new((column, column_type));
let source: Box<dyn ValueSource> = Box::new((column, column_type));
source
})
.collect())
src/aggregation/agg_data.rs
+282 -474
View File
File diff suppressed because it is too large Load Diff
src/aggregation/agg_limits.rs
+9 -1
View File
@@ -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::<u64>()
}
}
impl<K, V, S> MemoryConsumption for HashMap<K, V, S> {
fn memory_consumption(&self) -> usize {
let capacity = self.capacity();
src/aggregation/bucket/composite/accessors.rs
+1 -1
View File
@@ -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,
src/aggregation/bucket/composite/collector.rs
+4 -19
View File
@@ -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<Box<dyn SegmentAggregationCollector>>,
) -> crate::Result<Self> {
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();
src/aggregation/bucket/filter.rs
+16 -37
View File
@@ -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<DocumentQueryEvaluator>,
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::<SegmentReader>()
+ self.evaluator.bitset.len() / 8 // BitSet memory (bits to bytes)
+ std::mem::size_of::<bool>()
// The struct itself (including the segment reader handle), the name and the bitset.
std::mem::size_of::<Self>() + self.name.len() + self.evaluator.bitset.memory_consumption()
}
}
@@ -516,52 +508,39 @@ pub struct SegmentFilterCollector<B: SubAggBuffer> {
impl<B: SubAggBuffer> SegmentFilterCollector<B> {
/// 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<Self> {
// 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<Box<dyn SegmentAggregationCollector>>,
) -> 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<AggNode>,
) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
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::<LowCardSubAggBuffer>::from_req_and_validate(
req, node, req_data,
)?,
SegmentFilterCollector::<LowCardSubAggBuffer>::from_req_and_validate(req_data, sub_agg),
))
} else {
Ok(Box::new(
SegmentFilterCollector::<HighCardSubAggBuffer>::from_req_and_validate(
req, node, req_data,
)?,
req_data, sub_agg,
),
))
}
}
src/aggregation/bucket/histogram/histogram.rs
+96 -96
View File
@@ -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<dyn ValueSource>,
pub(crate) accessor: Box<dyn ValueSource>,
/// 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<dyn ValueSource>,
name: String,
mut req: HistogramAggregation,
is_date_histogram: bool,
) -> crate::Result<Self> {
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<B: BucketIdSlot> 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<B: BucketIdSlot> SegmentHistogramCollector<B> {
}
pub(crate) fn from_req_and_validate(
agg_data: &mut AggregationsSegmentCtx,
node: &AggRefNode,
) -> crate::Result<Self> {
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<Box<dyn SegmentAggregationCollector>>,
) -> Self {
let dense_range = compute_dense_range(
&*req_data.accessor,
req_data.req.interval,
@@ -616,33 +650,47 @@ impl<B: BucketIdSlot> SegmentHistogramCollector<B> {
);
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<const LANES: usize>(
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<const LANES: usize>(
&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<Option<(HistogramAggReqData, DenseRange)>> {
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<DenseRange> {
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<AggNode>,
) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
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::<BucketId>::from_req_and_validate(req_data, sub_agg),
))
} else {
Ok(Box::new(
SegmentHistogramCollector::<BucketId>::from_req_and_validate(agg_data, node)?,
SegmentHistogramCollector::<()>::from_req_and_validate(req_data, sub_agg),
))
}
}
src/aggregation/bucket/multi_terms/mod.rs
+22 -31
View File
@@ -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<Packing: MultiTermsPacking>(
/// uses [`MultiTermsKey`].
pub(crate) fn build_segment_multi_terms_collector(
req: &mut AggregationsSegmentCtx,
node: &AggRefNode,
req_data: MultiTermsAggReqData,
children: Vec<AggNode>,
) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
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::<BucketId>(req, node, req_data, packed_layout)
build_multi_terms_collector::<BucketId>(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<BucketSlot: BucketIdSlot>(
req: &mut AggregationsSegmentCtx,
node: &AggRefNode,
sub_agg: Option<Box<dyn SegmentAggregationCollector>>,
req_data: MultiTermsAggReqData,
packed_layout: Option<Vec<FieldPack>>,
) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
@@ -984,12 +984,9 @@ fn build_multi_terms_collector<BucketSlot: BucketIdSlot>(
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::<LowCardSubAggBuffer, _, VecTermBuckets<BucketSlot>>(
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<BucketSlot: BucketIdSlot>(
HighCardSubAggBuffer,
_,
VecTermBuckets<BucketSlot, true>,
>(req, node, req_data, packing, num_terms, true);
>(req, sub_agg, req_data, packing, num_terms, true);
}
return box_multi_terms_collector::<HighCardSubAggBuffer, _, VecTermBuckets<BucketSlot>>(
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::<HighCardSubAggBuffer, _, PagedTermMap<BucketSlot>>(
req, node, req_data, packing, max_packed, false,
req, sub_agg, req_data, packing, max_packed, false,
);
}
return box_multi_terms_collector::<HighCardSubAggBuffer, _, HashMapTermBuckets<BucketSlot>>(
req, node, req_data, packing, 0, false,
req, sub_agg, req_data, packing, 0, false,
);
}
@@ -1021,12 +1018,12 @@ fn build_multi_terms_collector<BucketSlot: BucketIdSlot>(
HighCardSubAggBuffer,
_,
HashMapTermBuckets<BucketSlot, MultiTermsKey>,
>(req, node, req_data, packing, 0, false)
>(req, sub_agg, req_data, packing, 0, false)
}
fn box_multi_terms_collector<Buffer, Packing, BucketMap>(
req: &mut AggregationsSegmentCtx,
node: &AggRefNode,
sub_agg: Option<Box<dyn SegmentAggregationCollector>>,
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::<Buffer>::new(
build_segment_agg_collectors(req, &node.children)?,
))
};
let sub_agg = sub_agg.map(BufferedSubAggs::<Buffer>::new);
let all_fields_full = req_data
.fields
.iter()
src/aggregation/bucket/range.rs
+7 -19
View File
@@ -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<dyn ValueSource>,
pub(crate) accessor: Box<dyn ValueSource>,
/// The range aggregation request.
pub(crate) req: RangeAggregation,
/// The name of the aggregation.
@@ -280,7 +277,7 @@ impl<B: SubAggBuffer> SegmentAggregationCollector for SegmentRangeCollector<B> {
) -> 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<B: SubAggBuffer> SegmentAggregationCollector for SegmentRangeCollector<B> {
/// 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<AggNode>,
) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
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::<LowCardSubAggBuffer> {
src/aggregation/bucket/term_agg/flattened_term_histogram.rs
+98 -49
View File
@@ -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<R: BucketResolver, const LANES: usize> {
terms_req_data: TermsAggReqData,
/// The terms full column's values
terms_values: Arc<dyn ColumnValues>,
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<R: BucketResolver, const LANES: usize> 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::<VecTermBuckets<BucketId>, 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<R: BucketResolver, const LANES: usize> 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<dyn ColumnValues>,
hist_values: Arc<dyn ColumnValues>,
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<Box<dyn SegmentAggregationCollector>>> {
) -> Option<FlattenedPlan> {
// 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<AggNode>,
) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
// 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::<NUM_LOW_CARD_COUNT_LANES>(
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::<SINGLE_COUNT_LANE>(
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<Arc<dyn ColumnValues>> {
@@ -606,7 +655,7 @@ fn try_get_column_full_values(value_source: &dyn ValueSource) -> Option<Arc<dyn
fn build_flattened_collector<const LANES: usize>(
agg_data: &mut AggregationsSegmentCtx,
terms_req_data: &TermsAggReqData,
terms_req_data: TermsAggReqData,
hist_req_data: HistogramAggReqData,
terms_values: Arc<dyn ColumnValues>,
hist_values: Arc<dyn ColumnValues>,
@@ -686,7 +735,7 @@ fn build_flattened_collector<const LANES: usize>(
fn build_flattened_collector_with_resolver<R: BucketResolver, const LANES: usize>(
agg_data: &mut AggregationsSegmentCtx,
terms_req_data: &TermsAggReqData,
terms_req_data: TermsAggReqData,
terms_values: Arc<dyn ColumnValues<u64>>,
hist_req_data: HistogramAggReqData,
num_terms: usize,
@@ -716,9 +765,9 @@ fn build_flattened_collector_with_resolver<R: BucketResolver, const LANES: usize
term_counts,
counts,
base_pos,
terms_req_data: terms_req_data.clone(),
terms_req_data,
terms_values,
hist_req_data,
histogram: SegmentHistogramCollector::from_dense_rows(hist_req_data),
term_block: ColumnBlockAccessor::default(),
bucket_resolver,
all_docs_in_bounds,
src/aggregation/bucket/term_agg/mod.rs
+23 -26
View File
@@ -1,7 +1,6 @@
use std::fmt::Debug;
use std::hash::Hash;
use std::net::Ipv6Addr;
use std::sync::Arc;
use columnar::column_values::CompactSpaceU64Accessor;
use columnar::{
@@ -14,7 +13,7 @@ use serde::{Deserialize, Serialize};
use super::{BucketIdSlot, CustomOrder, Order, OrderTarget};
use crate::aggregation::agg_data::{
build_segment_agg_collectors, AggRefNode, AggregationsSegmentCtx,
build_sub_agg_collectors, find_sub_agg, AggNode, AggregationsSegmentCtx,
};
use crate::aggregation::agg_limits::MemoryConsumption;
use crate::aggregation::agg_req::Aggregations;
@@ -35,10 +34,10 @@ mod flattened_term_histogram;
/// Contains all information required by the SegmentTermCollector to perform the
/// terms aggregation on a segment.
#[derive(Debug, Clone)]
#[derive(Debug)]
pub(crate) struct TermsAggReqData {
/// The column accessor to access the fast field values.
pub(crate) accessor: Arc<dyn ValueSource>,
pub(crate) accessor: Box<dyn ValueSource>,
/// The string dictionary column if the field is of type text.
pub(crate) str_dict_column: Option<StrColumn>,
/// 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::<Self>()
+ std::mem::size_of::<TermsAggregationInternal>()
+ 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<M: TermAggregationMap>(
/// 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<AggNode>,
) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
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<TermMap: TermAggregationMap, B: SubAggBuffer> 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,
);
src/aggregation/bucket/term_missing_agg.rs
+13 -27
View File
@@ -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<u64>, 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<HighCardBufferedSubAggs>,
/// Idx = parent bucket id, Value = missing count for that bucket
missing_count_per_bucket: Vec<MissingCount>,
@@ -53,27 +51,15 @@ pub struct TermMissingAgg {
}
impl TermMissingAgg {
pub(crate) fn new(
agg_data: &mut AggregationsSegmentCtx,
node: &AggRefNode,
) -> crate::Result<Self> {
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<Box<dyn SegmentAggregationCollector>>,
) -> 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;
src/aggregation/collector.rs
+4 -4
View File
@@ -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<Self> {
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
src/aggregation/metric/cardinality/mod.rs
+5 -18
View File
@@ -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<dyn ValueSource>,
pub(crate) accessor: Box<dyn ValueSource>,
/// The string dictionary column if the field is of type string.
pub(crate) str_dict_column: Option<StrColumn>,
/// 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<Box<dyn SegmentAggregationCollector>> {
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::<BitSet>::from_req(
node.idx_in_req_data,
req_data.accessor.clone(),
req_data.missing_value_for_accessor,
max_term_ord_inclusive,
),
SegmentStrCardinalityCollector::<BitSet>::from_req(req_data, max_term_ord_inclusive),
))
} else {
Ok(Box::new(
SegmentStrCardinalityCollector::<TermOrdSet>::from_req(
node.idx_in_req_data,
req_data.accessor.clone(),
req_data.missing_value_for_accessor,
req_data,
max_term_ord_inclusive,
),
))
src/aggregation/metric/cardinality/numeric_collector.rs
+16 -17
View File
@@ -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<Option<CardinalityCollector>>,
accessor_idx: usize,
/// The name of the aggregation.
name: String,
/// The column accessor to access the fast field values.
accessor: Arc<dyn ValueSource>,
accessor: Box<dyn ValueSource>,
/// 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<dyn ValueSource>,
missing_value_for_accessor: Option<u64>,
) -> crate::Result<Self> {
pub(crate) fn from_req(req_data: CardinalityAggReqData) -> crate::Result<Self> {
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<f64> {
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()?;
src/aggregation/metric/cardinality/str_collector.rs
+11 -24
View File
@@ -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<S: TermOrdAccumulator> {
/// Buckets are Some(_) until they get consumed by
/// `add_intermediate_aggregation_result`.
buckets: Vec<Option<S>>,
accessor_idx: usize,
/// The column accessor to access the fast field values (term ordinals).
accessor: Arc<dyn ValueSource>,
/// The missing value normalized to the internal u64 representation of the field type.
missing_value_for_accessor: Option<u64>,
/// 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<CouponCache>,
/// Largest term_ord that may be inserted into a bucket, i.e.
@@ -117,7 +112,7 @@ impl<S: TermOrdAccumulator> Debug for SegmentStrCardinalityCollector<S> {
.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<S: TermOrdAccumulator> SegmentStrCardinalityCollector<S> {
pub fn from_req(
accessor_idx: usize,
accessor: Arc<dyn ValueSource>,
missing_value_for_accessor: Option<u64>,
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<S: TermOrdAccumulator + 'static> 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<S: TermOrdAccumulator + 'static> 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<S: TermOrdAccumulator + 'static> SegmentAggregationCollector
bucket_id: BucketId,
sub_agg_name: &str,
sub_agg_property: &str,
agg_data: &AggregationsSegmentCtx,
_agg_data: &AggregationsSegmentCtx,
) -> Option<f64> {
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
src/aggregation/metric/extended_stats.rs
+6 -7
View File
@@ -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<u64>,
accessor: Arc<dyn ValueSource>,
accessor: Box<dyn ValueSource>,
buckets: Vec<IntermediateExtendedStats>,
sigma: Option<f64>,
}
impl SegmentExtendedStatsCollector {
pub fn from_req(req: &MetricAggReqData, sigma: Option<f64>) -> Self {
pub(crate) fn from_req(req: MetricAggReqData, sigma: Option<f64>) -> 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();
src/aggregation/metric/mod.rs
+1 -2
View File
@@ -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<u64>,
/// The column accessor to access the fast field values.
pub(crate) accessor: Arc<dyn ValueSource>,
pub(crate) accessor: Box<dyn ValueSource>,
/// Used when converting to intermediate result
pub(crate) collecting_for: StatsType,
/// The missing value
src/aggregation/metric/percentiles.rs
+12 -16
View File
@@ -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<PercentilesCollector>,
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<u64>,
/// The column accessor to access the fast field values.
pub(crate) accessor: Arc<dyn ValueSource>,
pub(crate) accessor: Box<dyn ValueSource>,
}
#[derive(Clone, Serialize, Deserialize)]
@@ -248,16 +248,12 @@ impl PercentilesCollector {
}
impl SegmentPercentilesCollector {
pub fn from_req_and_validate(
missing_u64: Option<u64>,
accessor: Arc<dyn ValueSource>,
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<f64> {
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()?;
src/aggregation/metric/stats.rs
+13 -13
View File
@@ -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<const TYPE_ID: u8>(
req: &MetricAggReqData,
req: MetricAggReqData,
) -> Box<dyn SegmentAggregationCollector> {
let column_opt = req
.accessor
.as_column()
.map(|column| (column.clone(), req.accessor.column_type()));
Box::new(SegmentStatsCollector::<TYPE_ID> {
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<Box<dyn SegmentAggregationCollector>> {
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<const COLUMN_TYPE_ID: u8> {
pub(crate) missing_u64: Option<u64>,
pub(crate) accessor: Arc<dyn ValueSource>,
pub(crate) accessor: Box<dyn ValueSource>,
/// 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<const COLUMN_TYPE_ID: u8> 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::<COLUMN_TYPE_ID>(
&mut self.buckets[parent_bucket_id as usize],
@@ -322,7 +322,7 @@ impl<const COLUMN_TYPE_ID: u8> SegmentAggregationCollector
} else {
agg_data.column_block_accessor.fetch_block_with_missing(
docs,
&*self.accessor,
&mut *self.accessor,
self.missing_u64,
);
}
src/aggregation/metric/top_hits.rs
+25 -31
View File
@@ -516,34 +516,35 @@ impl TopHitsTopNComputer {
}
}
#[derive(Clone, Debug)]
pub(crate) struct TopHitsSegmentCollector {
segment_ordinal: SegmentOrdinal,
accessor_idx: usize,
req_data: TopHitsAggReqData,
buckets: Vec<TopNComputer<Vec<DocValueAndOrder>, 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<String, Vec<DynamicColumn>>,
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,
},
);
src/aggregation/value_source/block_accessor.rs
+20 -21
View File
@@ -40,7 +40,7 @@ pub(crate) struct ColumnBlockAccessor {
impl ColumnBlockAccessor {
#[inline]
pub(crate) fn fetch_block<S: ValueSource + ?Sized>(&mut self, docs: &[DocId], source: &S) {
pub(crate) fn fetch_block<S: ValueSource + ?Sized>(&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<S: ValueSource + ?Sized>(
&mut self,
docs: &[DocId],
source: &S,
source: &mut S,
missing_opt: Option<u64>,
) {
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<S: ValueSource + ?Sized>(
&mut self,
docs: &[DocId],
source: &S,
source: &mut S,
missing_opt: Option<u64>,
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<u64>,
ordered: bool,
) {
@@ -304,7 +304,6 @@ fn find_missing_docs(docs: &[u32], hits: &[u32], output: &mut Vec<u32>) {
#[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<u64>,
docids: &mut Vec<DocId>,
@@ -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<dyn ValueSource> = Arc::new((full_column(&[5, 6, 7]), ColumnType::U64));
let column: Box<dyn ValueSource> = 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<dyn ValueSource> = Arc::new(TestValueSource {
let computed: Box<dyn ValueSource> = 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<dyn ValueSource> = Arc::new(TestValueSource {
let mut computed: Box<dyn ValueSource> = 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::<Vec<_>>(),
[(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::<Vec<_>>();
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::<Vec<_>>(),
@@ -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);
src/aggregation/value_source/mod.rs
+5 -2
View File
@@ -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<u64>,
docids: &mut Vec<DocId>,
@@ -64,7 +67,7 @@ impl<ColumnRef: Borrow<Column<u64>> + std::fmt::Debug> ValueSource for (ColumnRe
#[inline]
fn load_block(
&self,
&mut self,
docs: &[DocId],
values: &mut Vec<u64>,
docids: &mut Vec<DocId>,
src/aggregation/value_source/tests.rs
+3 -3
View File
@@ -14,7 +14,7 @@ impl ValueSource for Constant {
}
fn load_block(
&self,
&mut self,
docs: &[DocId],
values: &mut Vec<u64>,
_docids: &mut Vec<DocId>,
@@ -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<Arc<dyn ValueSource>> {
Ok(Arc::new(Constant(self.0)))
fn for_segment(&self, _reader: &SegmentReader) -> crate::Result<Box<dyn ValueSource>> {
Ok(Box::new(Constant(self.0)))
}
}
fn index_with_scores(scores: &[u64]) -> crate::Index {
src/aggregation/value_source/value_source_registry.rs
+2 -2
View File
@@ -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<Arc<dyn ValueSource>>;
fn for_segment(&self, reader: &SegmentReader) -> crate::Result<Box<dyn ValueSource>>;
}
/// 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();