mirror of
https://github.com/quickwit-oss/tantivy.git
synced 2026-08-18 03:58:21 +00:00
Optimize fused term-histogram aggregation
Avoid histogram decoding for single-bucket ranges and use precomputed boundaries for small ranges. Split hot counters into lanes to reduce write dependencies on low-cardinality grids.
This commit is contained in:
@@ -632,14 +632,15 @@ impl<B: BucketIdSlot> SegmentHistogramCollector<B> {
|
||||
|
||||
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). Used by the fused
|
||||
/// terms×histogram collector to turn its flat 2D counters into the regular intermediate result,
|
||||
/// so cross-segment merging is shared with the general path.
|
||||
pub(crate) fn from_dense_rows(
|
||||
/// `counts[t * num_time_buckets .. (t + 1) * num_time_buckets]` (row-major), consolidating each
|
||||
/// cell's count lanes. Used by the fused terms×histogram collector to turn its flat 2D counters
|
||||
/// into the regular intermediate result, so cross-segment merging is shared with the general
|
||||
/// path.
|
||||
pub(crate) fn from_dense_rows<const LANES: usize>(
|
||||
req_data: HistogramAggReqData,
|
||||
base_pos: i64,
|
||||
num_time_buckets: usize,
|
||||
counts: &[u32],
|
||||
counts: &[[u32; LANES]],
|
||||
) -> Self {
|
||||
let interval = req_data.req.interval;
|
||||
let offset = req_data.offset;
|
||||
@@ -650,13 +651,13 @@ impl SegmentHistogramCollector<()> {
|
||||
let buckets = row
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(b, &doc_count)| SegmentHistogramBucketEntry {
|
||||
.map(|(b, count_lanes)| SegmentHistogramBucketEntry {
|
||||
key: get_bucket_key_from_pos(
|
||||
(base_pos + b as i64) as f64,
|
||||
interval,
|
||||
offset,
|
||||
),
|
||||
doc_count: doc_count as u64,
|
||||
doc_count: count_lanes.iter().map(|&count| u64::from(count)).sum(),
|
||||
bucket_id: (),
|
||||
})
|
||||
.collect();
|
||||
|
||||
@@ -337,19 +337,23 @@ impl TermsAggregationInternal {
|
||||
}
|
||||
|
||||
/// Threshold for using dense `Vec` storage ([`VecTermBuckets`] or
|
||||
/// [`LowCardVecTermBuckets`]) for term buckets. Below this many term ordinals the `Vec`
|
||||
/// [`VecTermBucketsWithLanes`]) for term buckets. Below this many term ordinals the `Vec`
|
||||
/// (direct-indexed, no hashing/paging) beats the paged/hashed storages. It preallocates
|
||||
/// `num_terms` slots, so the ceiling also bounds its memory.
|
||||
///
|
||||
/// TODO: Benchmark to validate the threshold
|
||||
pub const MAX_NUM_TERMS_FOR_VEC: u64 = 20_000;
|
||||
|
||||
/// Threshold below which a terms agg without sub-aggregations uses [`LowCardVecTermBuckets`], whose
|
||||
/// independent count lanes improve throughput when many consecutive values address the same few
|
||||
/// counters. With sub-aggregations, this threshold selects [`LowCardSubAggBuffer`] (which buffers
|
||||
/// docs in a per-bucket `Vec`) while retaining scalar term counters. Both specialized layouts only
|
||||
/// pay off for a handful of buckets, so this is far lower than [`MAX_NUM_TERMS_FOR_VEC`].
|
||||
pub const MAX_NUM_TERMS_FOR_LOWCARD_SUBAGG: u64 = 100;
|
||||
/// Maximum number of logical counters for which replicating count lanes is worthwhile.
|
||||
const MAX_NUM_BUCKETS_FOR_COUNT_LANES: usize = 100;
|
||||
|
||||
/// Threshold below which a terms agg without sub-aggregations uses
|
||||
/// [`VecTermBucketsWithLanes`], whose independent count lanes improve throughput when many
|
||||
/// consecutive values address the same few counters. With sub-aggregations, this threshold selects
|
||||
/// [`LowCardSubAggBuffer`] (which buffers docs in a per-bucket `Vec`) while retaining scalar term
|
||||
/// counters. Both specialized layouts only pay off for a handful of buckets, so this is far lower
|
||||
/// than [`MAX_NUM_TERMS_FOR_VEC`].
|
||||
pub const MAX_NUM_TERMS_FOR_LOWCARD_SUBAGG: u64 = MAX_NUM_BUCKETS_FOR_COUNT_LANES as u64;
|
||||
|
||||
/// Threshold for using [`PagedTermMap`] storage; larger term-id spaces use
|
||||
/// [`HashMapTermBuckets`] instead.
|
||||
@@ -845,16 +849,17 @@ const NUM_LOW_CARD_COUNT_LANES: usize = 8;
|
||||
/// A very-low-cardinality term bucket with split count storage and one shared sub-aggregation id.
|
||||
/// Only the counters are replicated: `bucket_id` remains unique per logical term.
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
struct TermBucketWithLanes<B> {
|
||||
count_lanes: [u32; NUM_LOW_CARD_COUNT_LANES],
|
||||
struct TermBucketWithLanes<B, const LANES: usize> {
|
||||
count_lanes: [u32; LANES],
|
||||
bucket_id: B,
|
||||
}
|
||||
|
||||
impl<B: BucketIdSlot> TermBucketWithLanes<B> {
|
||||
impl<B: BucketIdSlot, const LANES: usize> TermBucketWithLanes<B, LANES> {
|
||||
#[inline(always)]
|
||||
fn new(bucket_id_provider: &mut BucketIdProvider) -> Self {
|
||||
assert!(LANES > 0, "a term bucket needs at least one count lane");
|
||||
Self {
|
||||
count_lanes: [0; NUM_LOW_CARD_COUNT_LANES],
|
||||
count_lanes: [0; LANES],
|
||||
bucket_id: B::assign(bucket_id_provider),
|
||||
}
|
||||
}
|
||||
@@ -871,18 +876,18 @@ impl<B: BucketIdSlot> TermBucketWithLanes<B> {
|
||||
/// A dense term map for very low cardinality. Each logical bucket cycles writes over independent
|
||||
/// count lanes, which are consolidated when the map is converted into its result representation.
|
||||
#[derive(Clone, Debug)]
|
||||
struct VecTermBucketsWithLanes<B> {
|
||||
buckets: Vec<TermBucketWithLanes<B>>,
|
||||
struct VecTermBucketsWithLanes<B, const LANES: usize = NUM_LOW_CARD_COUNT_LANES> {
|
||||
buckets: Vec<TermBucketWithLanes<B, LANES>>,
|
||||
next_count_lane: usize,
|
||||
}
|
||||
|
||||
impl<B: BucketIdSlot> TermAggregationMap for VecTermBucketsWithLanes<B> {
|
||||
impl<B: BucketIdSlot, const LANES: usize> TermAggregationMap for VecTermBucketsWithLanes<B, LANES> {
|
||||
type Slot = B;
|
||||
|
||||
const SORTED_BY_KEY: bool = true;
|
||||
|
||||
fn get_memory_consumption(&self) -> usize {
|
||||
self.buckets.capacity() * std::mem::size_of::<TermBucketWithLanes<B>>()
|
||||
self.buckets.capacity() * std::mem::size_of::<TermBucketWithLanes<B, LANES>>()
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
@@ -890,12 +895,12 @@ impl<B: BucketIdSlot> TermAggregationMap for VecTermBucketsWithLanes<B> {
|
||||
let term_id_usize = term_id as usize;
|
||||
debug_assert!(
|
||||
term_id_usize < self.buckets.len(),
|
||||
"term_id {} out of bounds for LowCardVecTermBuckets (len={})",
|
||||
"term_id {} out of bounds for VecTermBucketsWithLanes (len={})",
|
||||
term_id,
|
||||
self.buckets.len()
|
||||
);
|
||||
|
||||
self.next_count_lane = (self.next_count_lane + 1) % NUM_LOW_CARD_COUNT_LANES;
|
||||
self.next_count_lane = (self.next_count_lane + 1) % LANES;
|
||||
let bucket = unsafe { self.buckets.get_unchecked_mut(term_id_usize) };
|
||||
bucket.count_lanes[self.next_count_lane] += 1;
|
||||
bucket.bucket_id
|
||||
@@ -913,9 +918,11 @@ impl<B: BucketIdSlot> TermAggregationMap for VecTermBucketsWithLanes<B> {
|
||||
}
|
||||
|
||||
fn new(num_terms: u64, bucket_id_provider: &mut BucketIdProvider) -> Self {
|
||||
let buckets = std::iter::repeat_with(|| TermBucketWithLanes::new(bucket_id_provider))
|
||||
.take(num_terms as usize)
|
||||
.collect();
|
||||
assert!(LANES > 0, "a term map needs at least one count lane");
|
||||
let buckets =
|
||||
std::iter::repeat_with(|| TermBucketWithLanes::<B, LANES>::new(bucket_id_provider))
|
||||
.take(num_terms as usize)
|
||||
.collect();
|
||||
Self {
|
||||
buckets,
|
||||
next_count_lane: 0,
|
||||
|
||||
@@ -4,9 +4,14 @@
|
||||
//! See [`SegmentTermHistogramCollector`] for the approach and [`maybe_build_collector`] for the
|
||||
//! conditions under which it is used.
|
||||
|
||||
use columnar::ColumnBlockAccessor;
|
||||
use std::fmt::Debug;
|
||||
|
||||
use super::{Bucket, SegmentTermCollector, TermsAggReqData, VecTermBuckets};
|
||||
use columnar::{Column, ColumnBlockAccessor, ColumnType};
|
||||
|
||||
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::bucket::{
|
||||
get_bucket_pos_f64, prepare_histogram_dense_range, HistogramAggReqData,
|
||||
@@ -19,19 +24,302 @@ use crate::aggregation::intermediate_agg_result::{
|
||||
use crate::aggregation::segment_agg_result::{BucketIdProvider, SegmentAggregationCollector};
|
||||
use crate::aggregation::{f64_from_fastfield_u64, BucketId};
|
||||
|
||||
/// Maximum number of cells (`num_terms × num_time_buckets`) in the fused flat 2D grid. Above this
|
||||
/// the grid would be too large/cache-unfriendly, so we fall back to the general buffered path.
|
||||
/// `1 << 14` cells = 128 KB of `u64` counters, comfortably L2-resident.
|
||||
/// Maximum number of physical counters in the fused flat grid. Above this the grid would be too
|
||||
/// large/cache-unfriendly, so we fall back to the general buffered path. Count lanes are included
|
||||
/// in this limit: `num_terms × num_time_buckets × LANES` may not exceed it.
|
||||
///
|
||||
/// Since we are only at the top-level, this won't be multiplied by any parent buckets.
|
||||
const MAX_FUSED_GRID_BUCKETS: usize = 16384;
|
||||
const MAX_FUSED_GRID_COUNTERS: usize = 16_384;
|
||||
|
||||
/// Scalar storage for grids whose term cardinality is too high for count lanes to pay off.
|
||||
const SINGLE_COUNT_LANE: usize = 1;
|
||||
|
||||
/// Fixed linear-scan sizes over encoded `u64` bucket starts. The collector picks the smallest
|
||||
/// specialization that fits; larger histograms use floating-point bucket computation.
|
||||
const NUM_SMALL_LINEAR_BUCKETS: usize = 4;
|
||||
const NUM_LARGE_LINEAR_BUCKETS: usize = 8;
|
||||
|
||||
trait BucketResolver: Debug + 'static {
|
||||
/// Fetches the histogram values needed for this block. Resolvers that do not inspect the
|
||||
/// histogram column (notably [`SingleBucketResolver`]) leave this as a no-op.
|
||||
fn prepare_block(&mut self, docs: &[crate::DocId]);
|
||||
|
||||
/// Number of logical histogram buckets produced by this resolver.
|
||||
fn num_buckets(&self) -> usize;
|
||||
|
||||
/// Resolves and counts one bucket per term when every histogram value is in bounds. Keeping the
|
||||
/// loop inside the resolver lets each implementation compile to its simplest block traversal.
|
||||
fn collect_block<const LANES: usize>(
|
||||
&mut self,
|
||||
term_ids: impl Iterator<Item = u64>,
|
||||
counts: &mut [[u32; LANES]],
|
||||
);
|
||||
|
||||
/// Bounded variant. Only the computed resolver is constructed for binding hard bounds; the
|
||||
/// default remains useful for the statically unbounded resolvers.
|
||||
fn collect_block_with_bounds<const LANES: usize>(
|
||||
&mut self,
|
||||
term_ids: impl Iterator<Item = u64>,
|
||||
counts: &mut [[u32; LANES]],
|
||||
_term_counts: &mut [[u32; LANES]],
|
||||
) {
|
||||
self.collect_block(term_ids, counts);
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn increment_grid_count<const LANES: usize>(
|
||||
counts: &mut [[u32; LANES]],
|
||||
term_id: u64,
|
||||
bucket: usize,
|
||||
num_time_buckets: usize,
|
||||
count_lane: usize,
|
||||
) {
|
||||
debug_assert!(
|
||||
bucket < num_time_buckets,
|
||||
"histogram bucket outside dense range"
|
||||
);
|
||||
counts[term_id as usize * num_time_buckets + bucket][count_lane] += 1;
|
||||
}
|
||||
|
||||
/// Resolver for a histogram whose entire value range maps to one bucket. It deliberately owns no
|
||||
/// block accessor: collecting this shape does not read or decode the histogram column at all.
|
||||
#[derive(Debug, Default)]
|
||||
struct SingleBucketResolver {
|
||||
next_count_lane: usize,
|
||||
}
|
||||
|
||||
impl BucketResolver for SingleBucketResolver {
|
||||
#[inline]
|
||||
fn prepare_block(&mut self, _docs: &[crate::DocId]) {}
|
||||
|
||||
#[inline]
|
||||
fn num_buckets(&self) -> usize {
|
||||
1
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn collect_block<const LANES: usize>(
|
||||
&mut self,
|
||||
term_ids: impl Iterator<Item = u64>,
|
||||
counts: &mut [[u32; LANES]],
|
||||
) {
|
||||
for term_id in term_ids {
|
||||
self.next_count_lane = (self.next_count_lane + 1) % LANES;
|
||||
increment_grid_count(counts, term_id, 0, 1, self.next_count_lane);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The general resolver. It preserves the existing field conversion and floating-point bucket
|
||||
/// calculation for histograms that do not use a specialized resolver.
|
||||
#[derive(Debug)]
|
||||
struct ComputedBucketResolver {
|
||||
hist_block: ColumnBlockAccessor<u64>,
|
||||
next_count_lane: usize,
|
||||
accessor: Column<u64>,
|
||||
is_full: bool,
|
||||
field_type: ColumnType,
|
||||
interval: f64,
|
||||
offset: f64,
|
||||
base_pos: i64,
|
||||
num_buckets: usize,
|
||||
bounds: crate::aggregation::bucket::HistogramBounds,
|
||||
}
|
||||
|
||||
impl ComputedBucketResolver {
|
||||
fn new(hist_req_data: &HistogramAggReqData, base_pos: i64, num_buckets: usize) -> Self {
|
||||
Self {
|
||||
hist_block: ColumnBlockAccessor::default(),
|
||||
next_count_lane: 0,
|
||||
accessor: hist_req_data.accessor.clone(),
|
||||
is_full: hist_req_data.accessor.get_cardinality().is_full(),
|
||||
field_type: hist_req_data.field_type,
|
||||
interval: hist_req_data.req.interval,
|
||||
offset: hist_req_data.offset,
|
||||
base_pos,
|
||||
num_buckets,
|
||||
bounds: hist_req_data.bounds,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl BucketResolver for ComputedBucketResolver {
|
||||
#[inline]
|
||||
fn prepare_block(&mut self, docs: &[crate::DocId]) {
|
||||
self.hist_block
|
||||
.fetch_block_with_is_full(docs, &self.accessor, self.is_full);
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn num_buckets(&self) -> usize {
|
||||
self.num_buckets
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn collect_block<const LANES: usize>(
|
||||
&mut self,
|
||||
term_ids: impl Iterator<Item = u64>,
|
||||
counts: &mut [[u32; LANES]],
|
||||
) {
|
||||
let field_type = self.field_type;
|
||||
let interval = self.interval;
|
||||
let offset = self.offset;
|
||||
let base_pos = self.base_pos;
|
||||
let num_buckets = self.num_buckets;
|
||||
for (term_id, hist_raw) in term_ids.zip(self.hist_block.iter_vals()) {
|
||||
let val = f64_from_fastfield_u64(hist_raw, field_type);
|
||||
let bucket = (get_bucket_pos_f64(val, interval, offset) as i64 - base_pos) as usize;
|
||||
self.next_count_lane = (self.next_count_lane + 1) % LANES;
|
||||
increment_grid_count(counts, term_id, bucket, num_buckets, self.next_count_lane);
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn collect_block_with_bounds<const LANES: usize>(
|
||||
&mut self,
|
||||
term_ids: impl Iterator<Item = u64>,
|
||||
counts: &mut [[u32; LANES]],
|
||||
term_counts: &mut [[u32; LANES]],
|
||||
) {
|
||||
let field_type = self.field_type;
|
||||
let interval = self.interval;
|
||||
let offset = self.offset;
|
||||
let base_pos = self.base_pos;
|
||||
let num_buckets = self.num_buckets;
|
||||
let bounds = self.bounds;
|
||||
let mut next_count_lane = self.next_count_lane;
|
||||
for (term_id, hist_raw) in term_ids.zip(self.hist_block.iter_vals()) {
|
||||
let val = f64_from_fastfield_u64(hist_raw, field_type);
|
||||
if bounds.contains(val) {
|
||||
let bucket = (get_bucket_pos_f64(val, interval, offset) as i64 - base_pos) as usize;
|
||||
increment_grid_count(counts, term_id, bucket, num_buckets, next_count_lane);
|
||||
} else {
|
||||
term_counts[term_id as usize][next_count_lane] += 1;
|
||||
}
|
||||
next_count_lane = (next_count_lane + 1) % LANES;
|
||||
}
|
||||
self.next_count_lane = next_count_lane;
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolver for a small histogram grid. Bucket starts are precomputed in monotonic fast-field
|
||||
/// `u64` space, then scanned linearly. `NUM_BUCKETS` is fixed so the optimizer can unroll the scan.
|
||||
#[derive(Debug)]
|
||||
struct LinearBucketResolver<const NUM_BUCKETS: usize> {
|
||||
hist_block: ColumnBlockAccessor<u64>,
|
||||
next_count_lane: usize,
|
||||
accessor: Column<u64>,
|
||||
is_full: bool,
|
||||
boundaries: [u64; NUM_BUCKETS],
|
||||
num_buckets: usize,
|
||||
}
|
||||
|
||||
impl<const NUM_BUCKETS: usize> LinearBucketResolver<NUM_BUCKETS> {
|
||||
fn new(
|
||||
hist_req_data: &HistogramAggReqData,
|
||||
base_pos: i64,
|
||||
num_time_buckets: usize,
|
||||
) -> Option<Self> {
|
||||
assert!(num_time_buckets > 1 && num_time_buckets <= NUM_BUCKETS);
|
||||
let max_encoded_value = hist_req_data.accessor.max_value();
|
||||
// Padding must compare false for every column value. There is no such `u64` sentinel when
|
||||
// the column contains `u64::MAX`, so that edge case uses the computed resolver instead.
|
||||
let padding = max_encoded_value.checked_add(1)?;
|
||||
let mut boundaries = [padding; NUM_BUCKETS];
|
||||
let mut bucket_start = hist_req_data.accessor.min_value();
|
||||
for bucket in 1..num_time_buckets {
|
||||
bucket_start = first_encoded_value_for_bucket(
|
||||
bucket_start,
|
||||
max_encoded_value,
|
||||
bucket,
|
||||
hist_req_data,
|
||||
base_pos,
|
||||
);
|
||||
boundaries[bucket - 1] = bucket_start;
|
||||
}
|
||||
Some(Self {
|
||||
hist_block: ColumnBlockAccessor::default(),
|
||||
next_count_lane: 0,
|
||||
accessor: hist_req_data.accessor.clone(),
|
||||
is_full: hist_req_data.accessor.get_cardinality().is_full(),
|
||||
boundaries,
|
||||
num_buckets: num_time_buckets,
|
||||
})
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn resolve(boundaries: &[u64; NUM_BUCKETS], hist_raw: u64) -> usize {
|
||||
let mut bucket = 0;
|
||||
for &boundary in boundaries {
|
||||
bucket += (hist_raw >= boundary) as usize;
|
||||
}
|
||||
bucket
|
||||
}
|
||||
}
|
||||
|
||||
impl<const NUM_BUCKETS: usize> BucketResolver for LinearBucketResolver<NUM_BUCKETS> {
|
||||
#[inline]
|
||||
fn prepare_block(&mut self, docs: &[crate::DocId]) {
|
||||
self.hist_block
|
||||
.fetch_block_with_is_full(docs, &self.accessor, self.is_full);
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn num_buckets(&self) -> usize {
|
||||
self.num_buckets
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn collect_block<const LANES: usize>(
|
||||
&mut self,
|
||||
term_ids: impl Iterator<Item = u64>,
|
||||
counts: &mut [[u32; LANES]],
|
||||
) {
|
||||
let num_buckets = self.num_buckets;
|
||||
let boundaries = &self.boundaries;
|
||||
for (term_id, hist_raw) in term_ids.zip(self.hist_block.iter_vals()) {
|
||||
let bucket = Self::resolve(boundaries, hist_raw);
|
||||
self.next_count_lane = (self.next_count_lane + 1) % LANES;
|
||||
increment_grid_count(counts, term_id, bucket, num_buckets, self.next_count_lane);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Finds the first monotonic fast-field value assigned to `target_bucket` or a later bucket. Doing
|
||||
/// this binary search once during construction preserves the existing floating-point edge behavior
|
||||
/// without paying for conversion/division during collection.
|
||||
fn first_encoded_value_for_bucket(
|
||||
mut encoded_lower_bound: u64,
|
||||
mut encoded_upper_bound: u64,
|
||||
target_bucket: usize,
|
||||
hist_req_data: &HistogramAggReqData,
|
||||
base_pos: i64,
|
||||
) -> u64 {
|
||||
while encoded_lower_bound < encoded_upper_bound {
|
||||
let encoded_midpoint =
|
||||
encoded_lower_bound + (encoded_upper_bound - encoded_lower_bound) / 2;
|
||||
let val = f64_from_fastfield_u64(encoded_midpoint, hist_req_data.field_type);
|
||||
let bucket = (get_bucket_pos_f64(val, hist_req_data.req.interval, hist_req_data.offset)
|
||||
as i64
|
||||
- base_pos) as usize;
|
||||
if bucket < target_bucket {
|
||||
encoded_lower_bound = encoded_midpoint + 1;
|
||||
} else {
|
||||
encoded_upper_bound = encoded_midpoint;
|
||||
}
|
||||
}
|
||||
encoded_lower_bound
|
||||
}
|
||||
|
||||
/// Fused collector for `terms` (low cardinality) × a single `histogram`/`date_histogram` leaf with
|
||||
/// nothing nested below it, when the resulting `num_terms × num_time_buckets` grid is small (see
|
||||
/// [`MAX_FUSED_GRID_BUCKETS`]).
|
||||
/// nothing nested below it, when the resulting counter grid is small (see
|
||||
/// [`MAX_FUSED_GRID_COUNTERS`]).
|
||||
///
|
||||
/// It keeps a flat, fully dense 2D counter grid (`counts[term * num_time_buckets + bucket]`) and a
|
||||
/// per-term total. A single pass reads both the term and histogram columns in document order and
|
||||
/// It keeps a flat, fully dense 2D counter grid
|
||||
/// (`counts[term * num_time_buckets + bucket][lane]`) and a per-term total. Cycling writes through
|
||||
/// independent lanes avoids a serial read/modify/write dependency when consecutive documents hit
|
||||
/// the same cell. A single pass reads both the term and histogram columns in document order and
|
||||
/// bumps the counters directly — no doc-id buffering, no per-term scattered re-fetch, no dynamic
|
||||
/// dispatch on flush, no per-bucket key/id storage during collection (keys are derived from the
|
||||
/// index at the end).
|
||||
@@ -40,38 +328,37 @@ const MAX_FUSED_GRID_BUCKETS: usize = 16384;
|
||||
/// handed to the shared intermediate-result builders, so cross-segment merging is identical to the
|
||||
/// general path.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct SegmentTermHistogramCollector {
|
||||
struct SegmentTermHistogramCollector<R: BucketResolver, const LANES: usize> {
|
||||
/// Per-term count of docs *outside* `hard_bounds` (still in `doc_count`, but in no bucket).
|
||||
/// Per-term total = this + the term's `counts` row-sum; left empty when there are no hard
|
||||
/// bounds (every doc is in-bounds, so there's no remainder to track).
|
||||
term_counts: Vec<u32>,
|
||||
/// Flattened `[num_terms * num_time_buckets]` histogram counters (`u32`, see
|
||||
/// `term_counts`).
|
||||
term_counts: Vec<[u32; LANES]>,
|
||||
/// Flattened `[num_terms * num_time_buckets]` histogram counters, each split into independent
|
||||
/// count lanes.
|
||||
///
|
||||
/// Each term id get its own contiguous slice of `num_time_buckets` histogram counter.
|
||||
/// When we count all docs (#nofilter), we can derive the per-term total as the sum over that
|
||||
/// term's slice.
|
||||
counts: Vec<u32>,
|
||||
/// Histogram buckets per term (the dense time-range length).
|
||||
num_time_buckets: usize,
|
||||
/// Each term id gets its own contiguous slice of `num_time_buckets` histogram counters. When
|
||||
/// we count all docs (#nofilter), we can derive the per-term total as the sum over that term's
|
||||
/// slice.
|
||||
counts: Vec<[u32; LANES]>,
|
||||
/// `bucket_pos` mapped to time-bucket index 0.
|
||||
base_pos: i64,
|
||||
terms_req_data: TermsAggReqData,
|
||||
/// The (cloned, normalized) histogram request: its column + interval/offset/bounds.
|
||||
hist_req_data: HistogramAggReqData,
|
||||
/// Private block accessors for both columns. We read them together, so each needs its own
|
||||
/// (the shared `agg_data` scratch accessor only holds one block at a time). Owning them keeps
|
||||
/// `collect` independent of `agg_data`.
|
||||
/// 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<u64>,
|
||||
hist_block: ColumnBlockAccessor<u64>,
|
||||
bucket_resolver: R,
|
||||
/// No hard bounds, so every doc is in-bounds.
|
||||
all_docs_in_bounds: bool,
|
||||
/// Both columns are full (fused-path precondition); cached so `collect` skips the per-block
|
||||
/// cardinality lookup in `fetch_block`.
|
||||
is_full: bool,
|
||||
/// The term column is full (a fused-path precondition); cached so `collect` skips the
|
||||
/// per-block cardinality lookup in `fetch_block`.
|
||||
term_is_full: bool,
|
||||
}
|
||||
|
||||
impl SegmentAggregationCollector for SegmentTermHistogramCollector {
|
||||
impl<R: BucketResolver, const LANES: usize> SegmentAggregationCollector
|
||||
for SegmentTermHistogramCollector<R, LANES>
|
||||
{
|
||||
fn add_intermediate_aggregation_result(
|
||||
&mut self,
|
||||
agg_data: &AggregationsSegmentCtx,
|
||||
@@ -87,14 +374,19 @@ impl SegmentAggregationCollector for SegmentTermHistogramCollector {
|
||||
let mut bucket_id_provider = BucketIdProvider::default();
|
||||
// Per-term total = histogram row-sum (in-bounds) + `term_counts` (out-of-bounds remainder,
|
||||
// empty when there are no hard bounds).
|
||||
let num_time_buckets = self.bucket_resolver.num_buckets();
|
||||
let term_buckets = VecTermBuckets {
|
||||
buckets: self
|
||||
.counts
|
||||
.chunks_exact(self.num_time_buckets)
|
||||
.chunks_exact(num_time_buckets)
|
||||
.enumerate()
|
||||
.map(|(term_id, row)| {
|
||||
let in_bounds: u32 = row.iter().sum();
|
||||
let out_of_bounds = self.term_counts.get(term_id).copied().unwrap_or(0);
|
||||
let in_bounds: u32 = row.iter().flat_map(|lanes| lanes.iter()).sum();
|
||||
let out_of_bounds = self
|
||||
.term_counts
|
||||
.get(term_id)
|
||||
.map(|lanes| lanes.iter().sum())
|
||||
.unwrap_or(0);
|
||||
Bucket {
|
||||
count: in_bounds + out_of_bounds,
|
||||
bucket_id: bucket_id_provider.next_bucket_id(),
|
||||
@@ -105,7 +397,7 @@ impl SegmentAggregationCollector for SegmentTermHistogramCollector {
|
||||
let mut histogram = SegmentHistogramCollector::<()>::from_dense_rows(
|
||||
self.hist_req_data.clone(),
|
||||
self.base_pos,
|
||||
self.num_time_buckets,
|
||||
num_time_buckets,
|
||||
&self.counts,
|
||||
);
|
||||
let name = self.terms_req_data.name.clone();
|
||||
@@ -131,47 +423,27 @@ impl SegmentAggregationCollector for SegmentTermHistogramCollector {
|
||||
"fused term-histogram collector is top-level only"
|
||||
);
|
||||
|
||||
// Fetch both columns into our own accessors (we read them together, so they can't share the
|
||||
// single `agg_data` scratch accessor). The collector owns all its inputs, so `collect`
|
||||
// doesn't touch `agg_data`.
|
||||
self.term_block
|
||||
.fetch_block_with_is_full(docs, &self.terms_req_data.accessor, self.is_full);
|
||||
self.hist_block
|
||||
.fetch_block_with_is_full(docs, &self.hist_req_data.accessor, self.is_full);
|
||||
// The term column is always needed. The resolver fetches the histogram column only when
|
||||
// bucket selection depends on its values; `SingleBucketResolver` makes this a no-op.
|
||||
self.term_block.fetch_block_with_is_full(
|
||||
docs,
|
||||
&self.terms_req_data.accessor,
|
||||
self.term_is_full,
|
||||
);
|
||||
self.bucket_resolver.prepare_block(docs);
|
||||
|
||||
// Hoist the loop-invariant fields into locals: the optimizer can't prove the
|
||||
// `self.counts`/`self.term_counts` writes don't alias these `self` fields, so it can't keep
|
||||
// them in registers and re-reads them from memory every iteration — ~15% slower on
|
||||
// `terms_status_with_date_histogram` when read straight from `self`.
|
||||
// Note: check which are actually relevant.
|
||||
let field_type = self.hist_req_data.field_type;
|
||||
let bounds = self.hist_req_data.bounds;
|
||||
let interval = self.hist_req_data.req.interval;
|
||||
let offset = self.hist_req_data.offset;
|
||||
let base_pos = self.base_pos;
|
||||
let num_time_buckets = self.num_time_buckets;
|
||||
let all_docs_in_bounds = self.all_docs_in_bounds;
|
||||
let term_counts = &mut self.term_counts;
|
||||
let counts = &mut self.counts;
|
||||
|
||||
// Both columns are full (checked at construction), so values align with `docs` positionally
|
||||
// and are read together in one pass.
|
||||
// In-bounds docs bump the `counts` grid, out-of-bounds bump `term_counts`; deriving the
|
||||
// total at flush avoids a per-doc `term_counts` RMW that serializes on
|
||||
// store-to-load forwarding.
|
||||
for (term_id, hist_raw) in self.term_block.iter_vals().zip(self.hist_block.iter_vals()) {
|
||||
let term_id = term_id as usize;
|
||||
let val = f64_from_fastfield_u64(hist_raw, field_type);
|
||||
if all_docs_in_bounds || bounds.contains(val) {
|
||||
let bucket = (get_bucket_pos_f64(val, interval, offset) as i64 - base_pos) as usize;
|
||||
debug_assert!(
|
||||
bucket < num_time_buckets,
|
||||
"histogram bucket outside dense range"
|
||||
);
|
||||
counts[term_id * num_time_buckets + bucket] += 1;
|
||||
} else {
|
||||
term_counts[term_id] += 1;
|
||||
}
|
||||
// Keep separate bounded and unbounded entry points so the common path has no bounds branch,
|
||||
// while the computed resolver can retain a direct hard-bounds loop without materializing an
|
||||
// iterator of `Option<usize>`.
|
||||
if self.all_docs_in_bounds {
|
||||
self.bucket_resolver
|
||||
.collect_block::<LANES>(self.term_block.iter_vals(), &mut self.counts);
|
||||
} else {
|
||||
self.bucket_resolver.collect_block_with_bounds::<LANES>(
|
||||
self.term_block.iter_vals(),
|
||||
&mut self.counts,
|
||||
&mut self.term_counts,
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -206,7 +478,7 @@ impl SegmentAggregationCollector for SegmentTermHistogramCollector {
|
||||
///
|
||||
/// 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 `num_terms × num_time_buckets` grid no larger than [`MAX_FUSED_GRID_BUCKETS`].
|
||||
/// and a physical counter grid no larger than [`MAX_FUSED_GRID_COUNTERS`].
|
||||
pub(super) fn maybe_build_collector(
|
||||
agg_data: &mut AggregationsSegmentCtx,
|
||||
node: &AggRefNode,
|
||||
@@ -218,8 +490,9 @@ pub(super) fn maybe_build_collector(
|
||||
// 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.
|
||||
//
|
||||
// We don't cap the term cardinality here: the flat grid is bounded by the total cell count
|
||||
// (`num_terms * num_time_buckets <= MAX_FUSED_GRID_BUCKETS`) checked below, which subsumes it.
|
||||
// We don't cap the term cardinality here: the flat grid is bounded by the total physical
|
||||
// counter count (`num_terms * num_time_buckets * LANES <= MAX_FUSED_GRID_COUNTERS`) checked
|
||||
// below, which subsumes it.
|
||||
//
|
||||
// We only allow this at the top-level, since we don't know how many buckets are created. We
|
||||
// are less likely to get enough docs for the preallocation to be worth and there's a risk of
|
||||
@@ -248,43 +521,150 @@ pub(super) fn maybe_build_collector(
|
||||
}
|
||||
|
||||
// Clone + normalize the histogram request and get its dense bucket range; only take the fused
|
||||
// path when the flat `num_terms × num_time_buckets` grid is small enough.
|
||||
// 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 num_terms = col_max_val.saturating_add(1) as usize;
|
||||
if num_terms.saturating_mul(range.len) > MAX_FUSED_GRID_BUCKETS {
|
||||
let num_grid_cells = num_terms.saturating_mul(range.len);
|
||||
let use_count_lanes = num_grid_cells <= MAX_NUM_BUCKETS_FOR_COUNT_LANES;
|
||||
let num_count_lanes = if use_count_lanes {
|
||||
NUM_LOW_CARD_COUNT_LANES
|
||||
} else {
|
||||
SINGLE_COUNT_LANE
|
||||
};
|
||||
if num_grid_cells.saturating_mul(num_count_lanes) > MAX_FUSED_GRID_COUNTERS {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// No hard bounds means every doc is in-bounds, letting `collect` short-circuit the bounds
|
||||
// check — and leaving `term_counts` (the out-of-bounds remainder) unused, so we skip allocating
|
||||
// it.
|
||||
let collector = if use_count_lanes {
|
||||
build_collector::<NUM_LOW_CARD_COUNT_LANES>(
|
||||
agg_data,
|
||||
terms_req_data,
|
||||
hist_req_data,
|
||||
num_terms,
|
||||
range.len,
|
||||
range.base_pos,
|
||||
)?
|
||||
} else {
|
||||
build_collector::<SINGLE_COUNT_LANE>(
|
||||
agg_data,
|
||||
terms_req_data,
|
||||
hist_req_data,
|
||||
num_terms,
|
||||
range.len,
|
||||
range.base_pos,
|
||||
)?
|
||||
};
|
||||
Ok(Some(collector))
|
||||
}
|
||||
|
||||
fn build_collector<const LANES: usize>(
|
||||
agg_data: &mut AggregationsSegmentCtx,
|
||||
terms_req_data: &TermsAggReqData,
|
||||
hist_req_data: HistogramAggReqData,
|
||||
num_terms: usize,
|
||||
num_time_buckets: usize,
|
||||
base_pos: i64,
|
||||
) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
|
||||
assert!(LANES > 0, "a fused grid needs at least one count lane");
|
||||
|
||||
let all_docs_in_bounds =
|
||||
hist_req_data.bounds.min == f64::MIN && hist_req_data.bounds.max == f64::MAX;
|
||||
let counts = vec![0u32; num_terms * range.len];
|
||||
if all_docs_in_bounds && num_time_buckets == 1 {
|
||||
return build_collector_with_resolver::<SingleBucketResolver, LANES>(
|
||||
agg_data,
|
||||
terms_req_data,
|
||||
hist_req_data,
|
||||
num_terms,
|
||||
base_pos,
|
||||
SingleBucketResolver::default(),
|
||||
);
|
||||
}
|
||||
if all_docs_in_bounds && num_time_buckets <= NUM_SMALL_LINEAR_BUCKETS {
|
||||
if let Some(resolver) = LinearBucketResolver::<NUM_SMALL_LINEAR_BUCKETS>::new(
|
||||
&hist_req_data,
|
||||
base_pos,
|
||||
num_time_buckets,
|
||||
) {
|
||||
return build_collector_with_resolver::<_, LANES>(
|
||||
agg_data,
|
||||
terms_req_data,
|
||||
hist_req_data,
|
||||
num_terms,
|
||||
base_pos,
|
||||
resolver,
|
||||
);
|
||||
}
|
||||
} else if all_docs_in_bounds && num_time_buckets <= NUM_LARGE_LINEAR_BUCKETS {
|
||||
if let Some(resolver) = LinearBucketResolver::<NUM_LARGE_LINEAR_BUCKETS>::new(
|
||||
&hist_req_data,
|
||||
base_pos,
|
||||
num_time_buckets,
|
||||
) {
|
||||
return build_collector_with_resolver::<_, LANES>(
|
||||
agg_data,
|
||||
terms_req_data,
|
||||
hist_req_data,
|
||||
num_terms,
|
||||
base_pos,
|
||||
resolver,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let resolver = ComputedBucketResolver::new(&hist_req_data, base_pos, num_time_buckets);
|
||||
build_collector_with_resolver::<_, LANES>(
|
||||
agg_data,
|
||||
terms_req_data,
|
||||
hist_req_data,
|
||||
num_terms,
|
||||
base_pos,
|
||||
resolver,
|
||||
)
|
||||
}
|
||||
|
||||
fn build_collector_with_resolver<R: BucketResolver, const LANES: usize>(
|
||||
agg_data: &mut AggregationsSegmentCtx,
|
||||
terms_req_data: &TermsAggReqData,
|
||||
hist_req_data: HistogramAggReqData,
|
||||
num_terms: usize,
|
||||
base_pos: i64,
|
||||
bucket_resolver: R,
|
||||
) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
|
||||
let num_time_buckets = bucket_resolver.num_buckets();
|
||||
// No hard bounds means every doc is in-bounds, letting `collect` use its branch-free loop and
|
||||
// leaving `term_counts` (the out-of-bounds remainder) unused.
|
||||
let all_docs_in_bounds =
|
||||
hist_req_data.bounds.min == f64::MIN && hist_req_data.bounds.max == f64::MAX;
|
||||
let counts = vec![[0u32; LANES]; num_terms * num_time_buckets];
|
||||
let term_counts = if all_docs_in_bounds {
|
||||
Vec::new()
|
||||
} else {
|
||||
vec![0u32; num_terms]
|
||||
vec![[0u32; LANES]; num_terms]
|
||||
};
|
||||
// Charge both grids to the aggregation memory limit.
|
||||
agg_data.context.limits.add_memory_consumed(
|
||||
((counts.len() + term_counts.len()) * std::mem::size_of::<u32>()) as u64,
|
||||
)?;
|
||||
Ok(Some(Box::new(SegmentTermHistogramCollector {
|
||||
// Charge both grids, including replicated lanes, to the aggregation memory limit.
|
||||
let memory_consumption =
|
||||
std::mem::size_of_val(counts.as_slice()) + std::mem::size_of_val(term_counts.as_slice());
|
||||
agg_data
|
||||
.context
|
||||
.limits
|
||||
.add_memory_consumed(memory_consumption as u64)?;
|
||||
|
||||
Ok(Box::new(SegmentTermHistogramCollector::<R, LANES> {
|
||||
term_counts,
|
||||
counts,
|
||||
num_time_buckets: range.len,
|
||||
base_pos: range.base_pos,
|
||||
base_pos,
|
||||
terms_req_data: terms_req_data.clone(),
|
||||
hist_req_data,
|
||||
term_block: ColumnBlockAccessor::default(),
|
||||
hist_block: ColumnBlockAccessor::default(),
|
||||
bucket_resolver,
|
||||
all_docs_in_bounds,
|
||||
is_full: terms_req_data.accessor.get_cardinality().is_full(),
|
||||
})))
|
||||
term_is_full: terms_req_data.accessor.get_cardinality().is_full(),
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -348,6 +728,51 @@ mod tests {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Four and eight histogram buckets take their corresponding fixed-size linear resolvers.
|
||||
/// Negative values also exercise monotonic `f64` fast-field boundaries on both sides of zero.
|
||||
#[test]
|
||||
fn fused_term_histogram_linear_bucket_resolver() -> crate::Result<()> {
|
||||
for num_buckets in [4usize, 8] {
|
||||
// Three terms are coprime with both bucket counts, so every pair occurs 10 times.
|
||||
let docs: Vec<(f64, String)> = (0..3 * num_buckets * 10)
|
||||
.map(|i| {
|
||||
(
|
||||
(i % num_buckets) as f64 - (num_buckets / 2) as f64,
|
||||
["a", "b", "c"][i % 3].to_string(),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
let index = get_test_index_from_values_and_terms(true, &[docs])?;
|
||||
let agg_req: Aggregations = serde_json::from_value(serde_json::json!({
|
||||
"by_term": {
|
||||
"terms": { "field": "string_id", "order": { "_key": "asc" } },
|
||||
"aggs": {
|
||||
"histo": { "histogram": { "field": "score_f64", "interval": 1.0 } }
|
||||
}
|
||||
}
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
let res = exec_request(agg_req, &index)?;
|
||||
for (term_idx, term) in ["a", "b", "c"].iter().enumerate() {
|
||||
assert_eq!(res["by_term"]["buckets"][term_idx]["key"], *term);
|
||||
assert_eq!(
|
||||
res["by_term"]["buckets"][term_idx]["doc_count"],
|
||||
num_buckets * 10
|
||||
);
|
||||
let histo = &res["by_term"]["buckets"][term_idx]["histo"]["buckets"];
|
||||
for bucket in 0..num_buckets {
|
||||
assert_eq!(
|
||||
histo[bucket]["key"],
|
||||
bucket as f64 - (num_buckets / 2) as f64
|
||||
);
|
||||
assert_eq!(histo[bucket]["doc_count"], 10);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// A `missing` config on a *full* term column still takes the fused path (the string sentinel
|
||||
/// is just `col_max + 1`, so the column stays low-cardinality). Since no doc is missing, the
|
||||
/// real term buckets must be exactly as without `missing`.
|
||||
|
||||
Reference in New Issue
Block a user