From 3bd8a6b4d17ce6e2d0eefa40969bf7240657b8a3 Mon Sep 17 00:00:00 2001 From: Pascal Seitz Date: Wed, 29 Jul 2026 17:22:30 +0200 Subject: [PATCH] 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. --- src/aggregation/bucket/histogram/histogram.rs | 15 +- src/aggregation/bucket/term_agg/mod.rs | 47 +- .../bucket/term_agg/term_histogram.rs | 609 +++++++++++++++--- 3 files changed, 552 insertions(+), 119 deletions(-) diff --git a/src/aggregation/bucket/histogram/histogram.rs b/src/aggregation/bucket/histogram/histogram.rs index d755f4fa7..0942a08f7 100644 --- a/src/aggregation/bucket/histogram/histogram.rs +++ b/src/aggregation/bucket/histogram/histogram.rs @@ -632,14 +632,15 @@ impl SegmentHistogramCollector { 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( 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(); diff --git a/src/aggregation/bucket/term_agg/mod.rs b/src/aggregation/bucket/term_agg/mod.rs index 9d3b2f724..1692a3bb5 100644 --- a/src/aggregation/bucket/term_agg/mod.rs +++ b/src/aggregation/bucket/term_agg/mod.rs @@ -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 { - count_lanes: [u32; NUM_LOW_CARD_COUNT_LANES], +struct TermBucketWithLanes { + count_lanes: [u32; LANES], bucket_id: B, } -impl TermBucketWithLanes { +impl TermBucketWithLanes { #[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 TermBucketWithLanes { /// 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 { - buckets: Vec>, +struct VecTermBucketsWithLanes { + buckets: Vec>, next_count_lane: usize, } -impl TermAggregationMap for VecTermBucketsWithLanes { +impl TermAggregationMap for VecTermBucketsWithLanes { type Slot = B; const SORTED_BY_KEY: bool = true; fn get_memory_consumption(&self) -> usize { - self.buckets.capacity() * std::mem::size_of::>() + self.buckets.capacity() * std::mem::size_of::>() } #[inline(always)] @@ -890,12 +895,12 @@ impl TermAggregationMap for VecTermBucketsWithLanes { 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 TermAggregationMap for VecTermBucketsWithLanes { } 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::::new(bucket_id_provider)) + .take(num_terms as usize) + .collect(); Self { buckets, next_count_lane: 0, diff --git a/src/aggregation/bucket/term_agg/term_histogram.rs b/src/aggregation/bucket/term_agg/term_histogram.rs index 7f60970a7..31cd712b7 100644 --- a/src/aggregation/bucket/term_agg/term_histogram.rs +++ b/src/aggregation/bucket/term_agg/term_histogram.rs @@ -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( + &mut self, + term_ids: impl Iterator, + 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( + &mut self, + term_ids: impl Iterator, + counts: &mut [[u32; LANES]], + _term_counts: &mut [[u32; LANES]], + ) { + self.collect_block(term_ids, counts); + } +} + +#[inline] +fn increment_grid_count( + 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( + &mut self, + term_ids: impl Iterator, + 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, + next_count_lane: usize, + accessor: Column, + 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( + &mut self, + term_ids: impl Iterator, + 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( + &mut self, + term_ids: impl Iterator, + 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 { + hist_block: ColumnBlockAccessor, + next_count_lane: usize, + accessor: Column, + is_full: bool, + boundaries: [u64; NUM_BUCKETS], + num_buckets: usize, +} + +impl LinearBucketResolver { + fn new( + hist_req_data: &HistogramAggReqData, + base_pos: i64, + num_time_buckets: usize, + ) -> Option { + 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 BucketResolver for LinearBucketResolver { + #[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( + &mut self, + term_ids: impl Iterator, + 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 { /// 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, - /// 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, - /// 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, - hist_block: ColumnBlockAccessor, + 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 SegmentAggregationCollector + for SegmentTermHistogramCollector +{ 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`. + if self.all_docs_in_bounds { + self.bucket_resolver + .collect_block::(self.term_block.iter_vals(), &mut self.counts); + } else { + self.bucket_resolver.collect_block_with_bounds::( + 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::( + agg_data, + terms_req_data, + hist_req_data, + num_terms, + range.len, + range.base_pos, + )? + } else { + build_collector::( + agg_data, + terms_req_data, + hist_req_data, + num_terms, + range.len, + range.base_pos, + )? + }; + Ok(Some(collector)) +} + +fn build_collector( + agg_data: &mut AggregationsSegmentCtx, + terms_req_data: &TermsAggReqData, + hist_req_data: HistogramAggReqData, + num_terms: usize, + num_time_buckets: usize, + base_pos: i64, +) -> crate::Result> { + 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::( + 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::::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::::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( + agg_data: &mut AggregationsSegmentCtx, + terms_req_data: &TermsAggReqData, + hist_req_data: HistogramAggReqData, + num_terms: usize, + base_pos: i64, + bucket_resolver: R, +) -> crate::Result> { + 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::()) 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:: { 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`.