From 507e1760aada7aa9cee52d5920ea1588a2acdec4 Mon Sep 17 00:00:00 2001 From: Paul Masurel Date: Tue, 22 Sep 2026 10:47:49 +0200 Subject: [PATCH] Minor refactor of the cardinality aggregation. (#3120) The Str path and the non string path are very different. This PR isolates them as different segment aggregation collector. Co-authored-by: Paul Masurel --- src/aggregation/agg_data.rs | 39 +- src/aggregation/metric/cardinality.rs | 1407 ----------------- src/aggregation/metric/cardinality/mod.rs | 609 +++++++ .../metric/cardinality/numeric_collector.rs | 255 +++ .../metric/cardinality/str_collector.rs | 397 +++++ .../cardinality/term_ord_accumulator.rs | 343 ++++ 6 files changed, 1607 insertions(+), 1443 deletions(-) delete mode 100644 src/aggregation/metric/cardinality.rs create mode 100644 src/aggregation/metric/cardinality/mod.rs create mode 100644 src/aggregation/metric/cardinality/numeric_collector.rs create mode 100644 src/aggregation/metric/cardinality/str_collector.rs create mode 100644 src/aggregation/metric/cardinality/term_ord_accumulator.rs diff --git a/src/aggregation/agg_data.rs b/src/aggregation/agg_data.rs index bcb0acedd..16a08c0e4 100644 --- a/src/aggregation/agg_data.rs +++ b/src/aggregation/agg_data.rs @@ -22,9 +22,8 @@ use crate::aggregation::bucket::{ use crate::aggregation::metric::{ build_segment_stats_collector, AverageAggregation, CardinalityAggReqData, CardinalityAggregationReq, CountAggregation, ExtendedStatsAggregation, MaxAggregation, - MetricAggReqData, MinAggregation, SegmentCardinalityCollector, SegmentExtendedStatsCollector, - SegmentPercentilesCollector, StatsAggregation, StatsType, SumAggregation, TermOrdSet, - TopHitsAggReqData, TopHitsSegmentCollector, BITSET_MAX_TERM_ORD, + MetricAggReqData, MinAggregation, SegmentExtendedStatsCollector, SegmentPercentilesCollector, + StatsAggregation, StatsType, SumAggregation, TopHitsAggReqData, TopHitsSegmentCollector, }; use crate::aggregation::segment_agg_result::{ GenericSegmentAggregationResultsCollector, SegmentAggregationCollector, @@ -286,39 +285,7 @@ pub(crate) fn build_segment_agg_collector( Ok(Box::new(TermMissingAgg::new(req, node)?)) } AggKind::Cardinality => { - let req_data = req.get_cardinality_req_data(node.idx_in_req_data); - // For str columns, choose the per-bucket entries representation - // based on the segment's column.max_value(): - // * small (< BITSET_MAX_TERM_ORD): `BitSet`, pre-allocated, no promotion machinery. - // * large: `TermOrdSet` (sparse FxHashSet that promotes to a paged bitset). - // For non-str columns the `entries` field is unused (values go - // straight into the HLL sketch); we still pick `TermOrdSet` - // because its empty Sparse(FxHashSet) costs nothing. - let is_str = req_data.column_type == ColumnType::Str; - let max_term_ord_inclusive = if is_str { - req_data.accessor.max_value() - } else { - 0 - }; - let collector: Box = - if is_str && max_term_ord_inclusive < BITSET_MAX_TERM_ORD { - Box::new(SegmentCardinalityCollector::::from_req( - req_data.column_type, - node.idx_in_req_data, - req_data.accessor.clone(), - req_data.missing_value_for_accessor, - max_term_ord_inclusive, - )) - } else { - Box::new(SegmentCardinalityCollector::::from_req( - req_data.column_type, - node.idx_in_req_data, - req_data.accessor.clone(), - req_data.missing_value_for_accessor, - max_term_ord_inclusive, - )) - }; - Ok(collector) + crate::aggregation::metric::build_segment_cardinality_collector(req, node) } AggKind::StatsKind(stats_type) => { let req_data = &mut req.per_request.stats_metric_req_data[node.idx_in_req_data]; diff --git a/src/aggregation/metric/cardinality.rs b/src/aggregation/metric/cardinality.rs deleted file mode 100644 index d605125e2..000000000 --- a/src/aggregation/metric/cardinality.rs +++ /dev/null @@ -1,1407 +0,0 @@ -use std::fmt::Debug; -use std::hash::Hash; -use std::io; - -use columnar::column_values::CompactSpaceU64Accessor; -use columnar::{Column, ColumnType, Dictionary, StrColumn}; -use common::{BitSet, TinySet}; -use datasketches::hll::{Coupon, HllSketch, HllType, HllUnion}; -use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet}; -use serde::{Deserialize, Deserializer, Serialize, Serializer}; - -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::*; -use crate::TantivyError; - -/// Log2 of the number of registers for the HLL sketch. -/// 2^11 = 2048 registers, giving ~2.3% relative error and ~1KB per sketch (Hll4). -const LG_K: u8 = 11; - -/// Promote FxHashSet -> PagedBitset at ~3% density (`len * 32 > -/// dict_num_terms`). Past this point the bitset (~`dict_num_terms / 7.5` -/// bytes) is smaller than the hashset (~10 B/entry minimum) and avoids -/// the per-insert hash. -const PROMOTION_RATIO: u64 = 32; - -/// # Cardinality -/// -/// The cardinality aggregation allows for computing an estimate -/// of the number of different values in a data set based on the -/// Apache DataSketches HyperLogLog algorithm. This is particularly useful for -/// understanding the uniqueness of values in a large dataset where counting -/// each unique value individually would be computationally expensive. -/// -/// For example, you might use a cardinality aggregation to estimate the number -/// of unique visitors to a website by aggregating on a field that contains -/// user IDs or session IDs. -/// -/// To use the cardinality aggregation, you'll need to provide a field to -/// aggregate on. The following example demonstrates a request for the cardinality -/// of the "user_id" field: -/// -/// ```JSON -/// { -/// "cardinality": { -/// "field": "user_id" -/// } -/// } -/// ``` -/// -/// This request will return an estimate of the number of unique values in the -/// "user_id" field. -/// -/// ## Missing Values -/// -/// The `missing` parameter defines how documents that are missing a value should be treated. -/// By default, documents without a value for the specified field are ignored. However, you can -/// specify a default value for these documents using the `missing` parameter. This can be useful -/// when you want to include documents with missing values in the aggregation. -/// -/// For example, the following request treats documents with missing values in the "user_id" -/// field as if they had a value of "unknown": -/// -/// ```JSON -/// { -/// "cardinality": { -/// "field": "user_id", -/// "missing": "unknown" -/// } -/// } -/// ``` -/// -/// # Estimation Accuracy -/// -/// The cardinality aggregation provides an approximate count, which is usually -/// accurate within a small error range. This trade-off allows for efficient -/// computation even on very large datasets. -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct CardinalityAggregationReq { - /// The field name to compute the percentiles on. - pub field: String, - /// The missing parameter defines how documents that are missing a value should be treated. - /// By default they will be ignored but it is also possible to treat them as if they had a - /// value. Examples in JSON format: - /// { "field": "my_numbers", "missing": "10.0" } - #[serde(skip_serializing_if = "Option::is_none", default)] - pub missing: Option, -} - -/// Contains all information required by the SegmentCardinalityCollector to perform the -/// cardinality aggregation on a segment. -pub(crate) struct CardinalityAggReqData { - /// The column accessor to access the fast field values. - pub(crate) accessor: Column, - /// The column_type of the field. - pub(crate) column_type: ColumnType, - /// The string dictionary column if the field is of type string. - pub(crate) str_dict_column: Option, - /// The missing value normalized to the internal u64 representation of the field type. - pub(crate) missing_value_for_accessor: Option, - /// The name of the aggregation. - pub(crate) name: String, - /// The aggregation request. - pub(crate) req: CardinalityAggregationReq, -} - -impl CardinalityAggReqData { - /// Estimate the memory consumption of this struct in bytes. - pub fn get_memory_consumption(&self) -> usize { - std::mem::size_of::() - } -} - -impl CardinalityAggregationReq { - /// Creates a new [`CardinalityAggregationReq`] instance from a field name. - pub fn from_field_name(field_name: String) -> Self { - Self { - field: field_name, - missing: None, - } - } - /// Returns the field name the aggregation is computed on. - pub fn field_name(&self) -> &str { - &self.field - } -} - -/// A CouponCache is here to cache the mapping term ordinal -> coupon (see above). -/// The idea is that we do not want to fetch terms associated to several term ordinals, -/// several times due to the fact that we have several buckets. -enum CouponCache { - Dense { - coupon_map: Vec, - missing_coupon_opt: Option, - }, - Sparse { - coupon_map: FxHashMap, - missing_coupon_opt: Option, - }, -} - -impl CouponCache { - fn new( - term_ords: Vec, - coupons: Vec, - missing_coupon_opt: Option, - ) -> CouponCache { - let num_terms = term_ords.len(); - assert_eq!(num_terms, coupons.len()); - if term_ords.is_empty() { - return CouponCache::Dense { - coupon_map: Vec::new(), - missing_coupon_opt, - }; - } - let highest_term_ord = term_ords.last().copied().unwrap_or(0u64); - // We prefer the dense implementation, if it is not too wasteful. - // There are two cases for which we can use it. - // 1- if the data is small. - // 2- if the data is not necessarily small, but due to a high occupancy ratio, the RAM usage - // is not that much bigger than if we had used a HashSet. (occupancy ratio + extra - // metadata ~ x2.25) - let should_use_dense = - highest_term_ord < 1_000_000u64 || highest_term_ord < num_terms as u64 * 3u64; - if should_use_dense { - // We don't really care about the value here. We will populate all the values we will - // read anyway. - let uninitialized_coupon = Coupon::from_hash(0); - let mut coupon_map: Vec = - vec![uninitialized_coupon; highest_term_ord as usize + 1]; - - for (term_ord, coupon) in term_ords.into_iter().zip(coupons) { - coupon_map[term_ord as usize] = coupon; - } - CouponCache::Dense { - coupon_map, - missing_coupon_opt, - } - } else { - let coupon_map: FxHashMap = term_ords.into_iter().zip(coupons).collect(); - CouponCache::Sparse { - coupon_map, - missing_coupon_opt, - } - } - } -} - -// ================================================================= -// PagedBitset: a sparse bitset indexed by term_ord. -// -// Used as the dense alternative to FxHashSet once a string -// cardinality bucket has accumulated enough unique term ordinals. -// Memory is bounded to (touched pages) * (page bytes), not -// (max_term_ord / 8). -// -// Page geometry mirrors `PagedTermMap` in `term_agg.rs`: 1024 ords -// per page, lazy `Vec>>` directory. -// ================================================================= -const BITSET_PAGE_SHIFT: u32 = 10; -const BITSET_PAGE_BITS: u64 = 1u64 << BITSET_PAGE_SHIFT; // 1024 -const BITSET_PAGE_MASK: u64 = BITSET_PAGE_BITS - 1; -const BITSET_WORDS_PER_PAGE: usize = (BITSET_PAGE_BITS / 64) as usize; // 16 - -#[derive(Clone)] -struct PagedBitsetPage { - words: [TinySet; BITSET_WORDS_PER_PAGE], -} - -impl PagedBitsetPage { - fn new() -> Self { - Self { - words: [TinySet::empty(); BITSET_WORDS_PER_PAGE], - } - } -} - -pub(crate) struct PagedBitset { - pages: Vec>>, - /// Cached number of set bits, maintained on insert. - count: u64, -} - -impl PagedBitset { - /// Allocates a directory big enough to hold ords up to and including - /// `max_term_ord`. Pages are allocated lazily on first set. - fn with_max_term_ord(max_term_ord: u64) -> Self { - let max_page_idx = (max_term_ord >> BITSET_PAGE_SHIFT) as usize; - let num_pages = max_page_idx + 1; - Self { - pages: vec![None; num_pages], - count: 0, - } - } - - #[inline] - fn insert(&mut self, term_ord: u64) { - let page_idx = (term_ord >> BITSET_PAGE_SHIFT) as usize; - let intra = term_ord & BITSET_PAGE_MASK; - let word_idx = (intra >> 6) as usize; - let bit_idx = (intra & 63) as u32; - - let page = match &mut self.pages[page_idx] { - Some(p) => p, - None => { - self.pages[page_idx] = Some(Box::new(PagedBitsetPage::new())); - self.pages[page_idx].as_mut().unwrap() - } - }; - if page.words[word_idx].insert_mut(bit_idx) { - self.count += 1; - } - } - - /// Number of set bits. O(1). - #[inline] - fn len(&self) -> u64 { - self.count - } - - /// Iterate set ords in ascending order. - fn iter_sorted(&self) -> impl Iterator + '_ { - self.pages - .iter() - .enumerate() - .filter_map(|(page_idx, page_opt)| page_opt.as_ref().map(|p| (page_idx, p))) - .flat_map(|(page_idx, page)| { - let page_base_ord = (page_idx as u64) << BITSET_PAGE_SHIFT; - page.words - .iter() - .enumerate() - .flat_map(move |(word_idx, &word)| { - let word_base_ord = page_base_ord + (word_idx as u64) * 64; - word.into_iter() - .map(move |bit| word_base_ord + u64::from(bit)) - }) - }) - } -} - -/// Threshold below which we use `BitSet` instead of `TermOrdSet`. -/// -/// Both `BitSet` and `FxHashSet` have the same 32-byte struct, so the comparison is heap only: -/// * `BitSet` at T=256: 5 `TinySet` words covering 258 bits (with the missing-value sentinel) = -/// 40 bytes. -/// * `FxHashSet` after one insert: 4-bucket hashbrown table ≈ 56 bytes -pub(crate) const BITSET_MAX_TERM_ORD: u64 = 256; - -// ================================================================= -// TermOrdAccumulator: per-bucket abstraction over the entries set. -// -// Implementations: -// - `BitSet` (from `common`): used when `column.max_value()` is small (< BITSET_MAX_TERM_ORD). -// Pre-allocated, no promotion. -// - `TermOrdSet`: adaptive, starts as FxHashSet and promotes to a paged bitset when occupancy -// crosses the density threshold (only if promotion is enabled — typically gated on top-level -// aggregation). -// -// The trait lets `SegmentCardinalityCollector` be generic over the choice -// so the hot collect() loop monomorphizes to a direct call (no enum -// dispatch per insert). -// ================================================================= -pub(crate) trait TermOrdAccumulator: Sized { - /// Construct an empty accumulator. - /// `max_term_ord_inclusive` is the largest term_ord that may be - /// inserted (used to size pre-allocated bitsets and the dense bitset - /// on promotion). - fn new(max_term_ord_inclusive: u64) -> Self; - fn insert(&mut self, term_ord: u64); - /// Bulk insert. Implementations may override to hoist any inner - /// dispatch outside the loop. Default loops `insert`. - #[inline] - fn extend_from_iter>(&mut self, ords: I) { - for ord in ords { - self.insert(ord); - } - } - /// Hook called once per ingested block. Adaptive impls use this to - /// decide on sparse->dense promotion. - fn maybe_compact(&mut self) {} - fn len(&self) -> usize; - fn iter_ords(&self) -> impl Iterator + '_; -} - -impl TermOrdAccumulator for BitSet { - #[inline] - fn new(max_term_ord_inclusive: u64) -> Self { - // `BitSet::with_max_value(M)` accepts ords in [0, M). - // We need ords up to and including `max_term_ord_inclusive`, plus - // the missing-value sentinel `column.max_value() + 1`. - BitSet::with_max_value((max_term_ord_inclusive + 2) as u32) - } - #[inline] - fn insert(&mut self, term_ord: u64) { - BitSet::insert(self, term_ord as u32); - } - #[inline] - fn len(&self) -> usize { - BitSet::len(self) - } - fn iter_ords(&self) -> impl Iterator + '_ { - // `BitSet` itself doesn't expose iteration, but - // `BitSet::tinyset(bucket)` does. Walk per-bucket and yield each - // set bit. The capacity is `max_value()`; iterating to - // `div_ceil(64)` covers every possible ord exactly once. - let num_buckets = self.max_value().div_ceil(64); - (0..num_buckets).flat_map(move |bucket| { - let chunk_base = u64::from(bucket) * 64; - self.tinyset(bucket) - .into_iter() - .map(move |bit| chunk_base + u64::from(bit)) - }) - } -} - -// ================================================================= -// TermOrdSet: adaptive sparse->dense accumulator. -// -// Starts as an FxHashSet (cheap when few ords are seen). When occupancy -// crosses `len * PROMOTION_RATIO > max_term_ord_inclusive`, drains into -// a `PagedBitset` and continues dense. Promotion is one-way. -// ================================================================= -pub(crate) struct TermOrdSet { - inner: TermOrdSetInner, - /// Largest term_ord that may be inserted. Used for both sizing the - /// dense bitset on promotion and as the promotion-threshold reference. - max_term_ord_inclusive: u64, -} - -enum TermOrdSetInner { - Sparse(FxHashSet), - Dense(PagedBitset), -} - -impl TermOrdAccumulator for TermOrdSet { - fn new(max_term_ord_inclusive: u64) -> Self { - Self { - inner: TermOrdSetInner::Sparse(FxHashSet::default()), - max_term_ord_inclusive, - } - } - - #[inline] - fn insert(&mut self, term_ord: u64) { - match &mut self.inner { - TermOrdSetInner::Sparse(set) => { - set.insert(term_ord); - } - TermOrdSetInner::Dense(bitset) => bitset.insert(term_ord), - } - } - - /// Hoist the Sparse/Dense match outside the per-ord loop so that a - /// block of inserts dispatches once. - fn extend_from_iter>(&mut self, ords: I) { - match &mut self.inner { - TermOrdSetInner::Sparse(set) => { - for ord in ords { - set.insert(ord); - } - } - TermOrdSetInner::Dense(bitset) => { - for ord in ords { - bitset.insert(ord); - } - } - } - } - - fn maybe_compact(&mut self) { - let TermOrdSetInner::Sparse(set) = &mut self.inner else { - return; - }; - if set.len() as u64 * PROMOTION_RATIO <= self.max_term_ord_inclusive { - return; - } - // Size for ord <= max_term_ord_inclusive plus the missing sentinel - // (column.max_value() + 1, which may equal max_term_ord_inclusive - // when the column references every dictionary term). - let mut bitset = PagedBitset::with_max_term_ord(self.max_term_ord_inclusive + 1); - let set = std::mem::take(set); - for ord in set { - bitset.insert(ord); - } - self.inner = TermOrdSetInner::Dense(bitset); - } - - fn len(&self) -> usize { - match &self.inner { - TermOrdSetInner::Sparse(set) => set.len(), - TermOrdSetInner::Dense(bitset) => bitset.len() as usize, - } - } - - fn iter_ords(&self) -> impl Iterator + '_ { - match &self.inner { - TermOrdSetInner::Sparse(set) => itertools::Either::Left(set.iter().copied()), - TermOrdSetInner::Dense(bitset) => itertools::Either::Right(bitset.iter_sorted()), - } - } -} - -pub(crate) struct SegmentCardinalityCollector { - /// Buckets are Some(_) until they get consumed by into_intermediate_results(). - buckets: Vec>>, - accessor_idx: usize, - /// The column accessor to access the fast field values. - accessor: Column, - /// The column_type of the field. - column_type: ColumnType, - /// The missing value normalized to the internal u64 representation of the field type. - missing_value_for_accessor: Option, - coupon_cache: Option, - /// Largest term_ord that may be inserted into a bucket. For str columns - /// this is `accessor.max_value()`; for non-str columns this is unused - /// (no inserts go into `entries`) and set to 0. - max_term_ord_inclusive: u64, -} - -impl Debug for SegmentCardinalityCollector { - fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - f.debug_struct("SegmentCardinalityCollector") - .field("column_type", &self.column_type) - .field( - "missing_value_for_accessor", - &self.missing_value_for_accessor, - ) - .finish() - } -} - -/// Per-bucket state. Shape depends on column kind: str columns dedup -/// term ords and only build the HLL sketch at finalization (saves the -/// ~96 B `CardinalityCollector` per bucket during collect); numeric/IpAddr -/// columns feed the sketch directly during collect. -pub(crate) enum SegmentCardinalityCollectorBucket { - Str(S), - Numeric(CardinalityCollector), -} -impl SegmentCardinalityCollectorBucket { - #[inline(always)] - pub fn new(column_type: ColumnType, max_term_ord_inclusive: u64) -> Self { - if column_type == ColumnType::Str { - Self::Str(S::new(max_term_ord_inclusive)) - } else { - Self::Numeric(CardinalityCollector::new(column_type as u8)) - } - } - - // Returns a intermediate metric result. - // - // If the column is not str, the values have been added to the - // sketch during collection. - // - // If the column is str, then the values are dictionary encoded - // and have not been added to the sketch yet. - // We need to resolves the term ords accumulated in the str entries - // with the coupon cache, and append the results to a fresh sketch. - fn into_intermediate_metric_result( - self, - coupon_cache_opt: Option<&CouponCache>, - ) -> crate::Result { - let cardinality = match self { - Self::Str(entries) => { - let mut cardinality = CardinalityCollector::new(ColumnType::Str as u8); - if let Some(coupon_cache) = coupon_cache_opt { - // Sketch must be empty for str columns: coupons are appended here - // from the term_ord set (and not directly during collection). - assert!(cardinality.sketch.is_empty()); - append_to_sketch(&entries, coupon_cache, &mut cardinality); - } - cardinality - } - Self::Numeric(cardinality) => cardinality, - }; - Ok(IntermediateMetricResult::Cardinality(cardinality)) - } -} - -/// Builds a coupon cache from the given buckets, dictionary, and optional missing value. -/// Returns a mapping from term_ord to the hash (coupon) of the associated term. -fn build_coupon_cache( - buckets: &[Option>], - dictionary: &Dictionary, - missing_value_opt: Option<&Key>, -) -> io::Result { - // Caller restricts this to str cardinality collectors, so every - // present bucket must be the `Str` variant. Pass 1 validates and - // computes the capacity hint; pass 2 inserts. - let mut max_bucket_len = 0usize; - for bucket in buckets.iter().flatten() { - match bucket { - SegmentCardinalityCollectorBucket::Str(entries) => { - max_bucket_len = max_bucket_len.max(entries.len()); - } - SegmentCardinalityCollectorBucket::Numeric(_) => { - return Err(io::Error::other( - "build_coupon_cache invoked with a non-str bucket", - )); - } - } - } - let mut term_ords_set = FxHashSet::with_capacity_and_hasher(max_bucket_len * 2, FxBuildHasher); - for bucket in buckets.iter().flatten() { - if let SegmentCardinalityCollectorBucket::Str(entries) = bucket { - term_ords_set.extend(entries.iter_ords()); - } - } - let mut term_ords: Vec = term_ords_set.into_iter().collect(); - term_ords.sort_unstable(); - - term_ords.pop_if(|highest_term_ord| *highest_term_ord >= dictionary.num_terms() as u64); - - let mut coupons: Vec = Vec::with_capacity(term_ords.len()); - let all_term_ords_found: bool = - dictionary.sorted_ords_to_term_cb(&term_ords, |term_bytes| { - let coupon: Coupon = Coupon::from_hash(term_bytes); - coupons.push(coupon); - })?; - assert!(all_term_ords_found); - - // Regardless of whether or not there is effectively a missing value in one of the buckets, - // we populate the cache with the missing key too (if any). - let missing_coupon_opt: Option = missing_value_opt.map(|missing_key| { - if let Key::Str(missing_value_str) = missing_key { - Coupon::from_hash(missing_value_str.as_bytes()) - } else { - // See https://github.com/quickwit-oss/tantivy/issues/2891 - // A missing key with a type different from Str will not work as intended - // for the moment. - // - // Right now this is just a partial workaround. - Coupon::from_hash("__tantivy_missing_non_str__".as_bytes()) - } - }); - Ok(CouponCache::new(term_ords, coupons, missing_coupon_opt)) -} - -fn append_to_sketch( - term_ords: &S, - coupon_cache: &CouponCache, - sketch: &mut CardinalityCollector, -) { - match coupon_cache { - CouponCache::Dense { - coupon_map, - missing_coupon_opt, - } => { - for term_ord in term_ords.iter_ords() { - if let Some(coupon) = coupon_map - .get(term_ord as usize) - .copied() - .or(*missing_coupon_opt) - { - sketch.insert_coupon(coupon); - } - } - } - CouponCache::Sparse { - coupon_map, - missing_coupon_opt, - } => { - for term_ord in term_ords.iter_ords() { - if let Some(coupon) = coupon_map.get(&term_ord).copied().or(*missing_coupon_opt) { - sketch.insert_coupon(coupon); - } - } - } - } -} - -impl SegmentCardinalityCollector { - pub fn from_req( - column_type: ColumnType, - accessor_idx: usize, - accessor: Column, - missing_value_for_accessor: Option, - max_term_ord_inclusive: u64, - ) -> Self { - Self { - buckets: Vec::new(), - column_type, - accessor_idx, - accessor, - missing_value_for_accessor, - coupon_cache: None, - max_term_ord_inclusive, - } - } - - fn fetch_block_with_field( - &mut self, - docs: &[crate::DocId], - agg_data: &mut AggregationsSegmentCtx, - ) { - agg_data.column_block_accessor.fetch_block_with_missing( - docs, - &self.accessor, - self.missing_value_for_accessor, - ); - } -} - -impl SegmentAggregationCollector - for SegmentCardinalityCollector -{ - fn add_intermediate_aggregation_result( - &mut self, - agg_data: &AggregationsSegmentCtx, - results: &mut IntermediateAggregationResults, - 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); - // Strings are dictionary encoded. Fetching the terms associated to strings - // is expensive. For this reason, we do that once for all buckets and cache the results - // here. - if let Some(str_dict_column) = &req_data.str_dict_column { - // Ensure the coupon cache is populated. - // A mapping from term_ord to the hash of the associated term. - // The missing value sentinel will be associated to the hash of the missing value if - // any. - if self.coupon_cache.is_none() { - self.coupon_cache = Some(build_coupon_cache( - &self.buckets, - str_dict_column.dictionary(), - req_data.req.missing.as_ref(), - )?); - } - } - let name = req_data.name.to_string(); - // take the bucket in buckets and replace it with a new empty one - let Some(bucket) = self.buckets[bucket_id as usize].take() else { - return Err(crate::TantivyError::InternalError( - "the same bucket should not be finalized twice.".to_string(), - )); - }; - let intermediate_result = - bucket.into_intermediate_metric_result(self.coupon_cache.as_ref())?; - results.push( - name, - IntermediateAggregationResult::Metric(intermediate_result), - )?; - - Ok(()) - } - - fn collect( - &mut self, - parent_bucket_id: BucketId, - docs: &[crate::DocId], - agg_data: &mut AggregationsSegmentCtx, - ) -> crate::Result<()> { - self.fetch_block_with_field(docs, agg_data); - let Some(bucket) = &mut self.buckets[parent_bucket_id as usize].as_mut() else { - return Err(crate::TantivyError::InternalError( - "collection should not happen after finalization".to_string(), - )); - }; - let col_block_accessor = &agg_data.column_block_accessor; - match bucket { - SegmentCardinalityCollectorBucket::Str(entries) => { - // Promotion check runs on the pre-block state: the first call - // sees an empty set (no-op), and the last block of inserts - // doesn't trigger a promotion of a set we won't grow further. - // The trait dispatches once per block (via `extend_from_iter`) - // for adaptive variants and inlines to a tight loop for the - // BitSet path. - entries.maybe_compact(); - entries.extend_from_iter(col_block_accessor.iter_vals()); - } - SegmentCardinalityCollectorBucket::Numeric(cardinality) => { - if self.column_type == ColumnType::IpAddr { - let compact_space_accessor = self - .accessor - .values - .clone() - .downcast_arc::() - .map_err(|_| { - TantivyError::AggregationError( - crate::aggregation::AggregationError::InternalError( - "Type mismatch: Could not downcast to CompactSpaceU64Accessor" - .to_string(), - ), - ) - })?; - for val in col_block_accessor.iter_vals() { - let val: u128 = compact_space_accessor.compact_to_u128(val as u32); - cardinality.insert(val); - } - } else { - for val in col_block_accessor.iter_vals() { - cardinality.insert(val); - } - } - } - } - - Ok(()) - } - - fn prepare_max_bucket( - &mut self, - max_bucket: BucketId, - _agg_data: &AggregationsSegmentCtx, - ) -> crate::Result<()> { - if max_bucket as usize >= self.buckets.len() { - let column_type = self.column_type; - let max_term_ord_inclusive = self.max_term_ord_inclusive; - self.buckets.resize_with(max_bucket as usize + 1, || { - Some(SegmentCardinalityCollectorBucket::::new( - column_type, - max_term_ord_inclusive, - )) - }); - } - Ok(()) - } - - fn compute_metric_value( - &self, - bucket_id: BucketId, - sub_agg_name: &str, - sub_agg_property: &str, - agg_data: &AggregationsSegmentCtx, - ) -> Option { - let req_data = &agg_data.get_cardinality_req_data(self.accessor_idx); - if req_data.name != sub_agg_name || !sub_agg_property.is_empty() { - return None; - } - let bucket = self.buckets.get(bucket_id as usize)?.as_ref()?; - // For string columns the sketch isn't built until finalization; the - // term_ord set's len is the exact distinct count. For numeric columns - // the sketch is populated during collect. - match bucket { - SegmentCardinalityCollectorBucket::Str(entries) => Some(entries.len() as f64), - SegmentCardinalityCollectorBucket::Numeric(cardinality) => { - Some(cardinality.sketch.estimate().trunc()) - } - } - } -} - -#[derive(Clone, Debug)] -/// The cardinality collector used during segment collection and for merging results. -/// Uses Apache DataSketches HLL (lg_k=11, Hll4) for compact binary serialization -/// and cross-language compatibility (e.g. Java `datasketches` library). -pub struct CardinalityCollector { - sketch: HllSketch, - /// Salt derived from `ColumnType`, used to differentiate values of different column types - /// that map to the same u64 (e.g. bool `false` = 0 vs i64 `0`). - /// Not serialized — only needed during insertion, not after sketch registers are populated. - salt: u8, -} - -impl Default for CardinalityCollector { - fn default() -> Self { - Self::new(0) - } -} - -impl PartialEq for CardinalityCollector { - fn eq(&self, _other: &Self) -> bool { - false - } -} - -impl Serialize for CardinalityCollector { - fn serialize(&self, serializer: S) -> Result { - let bytes = self.sketch.serialize(); - serializer.serialize_bytes(&bytes) - } -} - -impl<'de> Deserialize<'de> for CardinalityCollector { - fn deserialize>(deserializer: D) -> Result { - let bytes: Vec = Deserialize::deserialize(deserializer)?; - let sketch = HllSketch::deserialize(&bytes).map_err(serde::de::Error::custom)?; - Ok(Self { sketch, salt: 0 }) - } -} - -impl CardinalityCollector { - fn new(salt: u8) -> Self { - Self { - sketch: HllSketch::new(LG_K, HllType::Hll8), - salt, - } - } - - /// Insert a value into the HLL sketch, salted by the column type. - /// The salt ensures that identical u64 values from different column types - /// (e.g. bool `false` vs i64 `0`) are counted as distinct. - fn insert(&mut self, value: T) { - self.sketch.update((self.salt, value)); - } - - fn insert_coupon(&mut self, coupon: Coupon) { - self.sketch.update_with_coupon(coupon); - } - - /// Compute the final cardinality estimate. - pub fn finalize(self) -> Option { - Some(self.sketch.estimate().trunc()) - } - - /// Serialize the HLL sketch to its compact binary representation. - /// The format is cross-language compatible with Apache DataSketches (Java, C++, Python). - pub fn to_sketch_bytes(&self) -> Vec { - self.sketch.serialize() - } - - pub(crate) fn merge_fruits(&mut self, right: CardinalityCollector) -> crate::Result<()> { - let mut union = HllUnion::new(LG_K); - union.update(&self.sketch); - union.update(&right.sketch); - self.sketch = union.to_sketch(HllType::Hll8); - Ok(()) - } -} - -#[cfg(test)] -mod tests { - - use std::net::IpAddr; - use std::str::FromStr; - - use columnar::MonotonicallyMappableToU64; - - use crate::aggregation::agg_req::Aggregations; - use crate::aggregation::tests::{exec_request, get_test_index_from_terms}; - use crate::schema::{IntoIpv6Addr, Schema, FAST, STRING}; - use crate::Index; - - #[test] - fn cardinality_aggregation_test_empty_index() -> crate::Result<()> { - let values = vec![]; - let index = get_test_index_from_terms(false, &values)?; - let agg_req: Aggregations = serde_json::from_value(json!({ - "cardinality": { - "cardinality": { - "field": "string_id", - } - }, - })) - .unwrap(); - - let res = exec_request(agg_req, &index)?; - assert_eq!(res["cardinality"]["value"], 0.0); - - Ok(()) - } - - #[test] - fn cardinality_aggregation_test_single_segment() -> crate::Result<()> { - cardinality_aggregation_test_merge_segment(true) - } - #[test] - fn cardinality_aggregation_test() -> crate::Result<()> { - cardinality_aggregation_test_merge_segment(false) - } - fn cardinality_aggregation_test_merge_segment(merge_segments: bool) -> crate::Result<()> { - let segment_and_terms = vec![ - vec!["terma"], - vec!["termb"], - vec!["termc"], - vec!["terma"], - vec!["terma"], - vec!["terma"], - vec!["termb"], - vec!["terma"], - ]; - let index = get_test_index_from_terms(merge_segments, &segment_and_terms)?; - let agg_req: Aggregations = serde_json::from_value(json!({ - "cardinality": { - "cardinality": { - "field": "string_id", - } - }, - })) - .unwrap(); - - let res = exec_request(agg_req, &index)?; - assert_eq!(res["cardinality"]["value"], 3.0); - - Ok(()) - } - - /// Build a single-segment string-cardinality index with 32 unique terms. - /// `column.max_value() = 31` is well below `BITSET_MAX_TERM_ORD`, - /// so the bucket exercises the `BitSet` path end to end. - #[test] - fn cardinality_aggregation_test_str_bitset() -> crate::Result<()> { - let terms: Vec = (0..32).map(|i| format!("term_{i}")).collect(); - let term_refs: Vec> = terms.iter().map(|t| vec![t.as_str()]).collect::>(); - // single segment so we have a single dictionary of 32 terms. - let index = get_test_index_from_terms(true, &term_refs)?; - - let agg_req: Aggregations = serde_json::from_value(json!({ - "cardinality": { - "cardinality": { "field": "string_id" } - }, - })) - .unwrap(); - - let res = exec_request(agg_req, &index)?; - assert_eq!(res["cardinality"]["value"], 32.0); - Ok(()) - } - - /// `BitSet` path with a `missing` parameter: the column-level missing - /// sentinel (`column.max_value() + 1`) flows into the bitset, the - /// dict lookup filter at finalization drops it, and the missing - /// coupon is applied separately. - #[test] - fn cardinality_aggregation_test_str_bitset_with_missing() { - let mut schema_builder = Schema::builder(); - let name_field = schema_builder.add_text_field("name", STRING | FAST); - let index = Index::create_in_ram(schema_builder.build()); - let mut writer = index.writer_for_tests().unwrap(); - for i in 0..16 { - let term = format!("t{i:02}"); - writer.add_document(doc!(name_field => term)).unwrap(); - } - // One empty doc, exercising the missing sentinel. - writer.add_document(doc!()).unwrap(); - writer.commit().unwrap(); - - let agg_req: Aggregations = serde_json::from_value(json!({ - "cardinality": { - "cardinality": { - "field": "name", - "missing": "MISSING_SENTINEL_KEY", - } - }, - })) - .unwrap(); - - let res = exec_request(agg_req, &index).unwrap(); - // 16 distinct real terms + 1 distinct "missing" value = 17. - assert_eq!(res["cardinality"]["value"], 17.0); - } - - /// Unit-test the PagedBitset itself: cross-page inserts produce sorted - /// iteration, len() matches the inserted set, and duplicates are - /// idempotent. - #[test] - fn paged_bitset_basic() { - use super::PagedBitset; - // Span several pages: BITSET_PAGE_BITS = 1024, so ords > 1024 land - // on the second page, > 2048 on the third, etc. - let ords = [0u64, 1, 63, 64, 1023, 1024, 1025, 4096, 4097, 9999, 10_000]; - let max_ord = *ords.iter().max().unwrap(); - let mut bitset = PagedBitset::with_max_term_ord(max_ord); - for &ord in &ords { - bitset.insert(ord); - // Idempotent: inserting again must not increase count. - bitset.insert(ord); - } - assert_eq!(bitset.len(), ords.len() as u64); - let collected: Vec = bitset.iter_sorted().collect(); - let mut expected: Vec = ords.to_vec(); - expected.sort_unstable(); - assert_eq!(collected, expected); - } - - /// Unit-test `TermOrdSet`: starts Sparse, promotes to Dense on - /// `maybe_compact` once the density threshold is crossed, and - /// `iter_ords()` yields the same set in either state. Ords spanning - /// multiple paged-bitset pages exercise the Dense iter ordering. - #[test] - fn term_ord_set_promotes_on_maybe_compact() { - use super::{TermOrdAccumulator, TermOrdSet, PROMOTION_RATIO}; - // Pick max so promotion needs few inserts: len * RATIO > max with - // RATIO=32 and max=64 trips at len=3 (3*32=96 > 64). - let max_term_ord = 64u64; - let mut set = ::new(max_term_ord); - // Two inserts: should stay Sparse after maybe_compact (2 * RATIO = 64, not > 64). - set.insert(0); - set.insert(7); - set.maybe_compact(); - assert_eq!(set.len(), 2); - - // Third insert promotes on next maybe_compact. - set.insert(20); - assert_eq!(set.len(), 3); - // Sanity check: at len=3, 3 * PROMOTION_RATIO = 96 > 64. - assert!(3u64 * PROMOTION_RATIO > max_term_ord); - set.maybe_compact(); - - // Post-promotion: extending continues to work. - set.insert(15); - set.insert(15); // dup - assert_eq!(set.len(), 4); - - let mut collected: Vec = set.iter_ords().collect(); - collected.sort_unstable(); - assert_eq!(collected, vec![0, 7, 15, 20]); - } - - /// Unit-test the `BitSet` impl of `TermOrdAccumulator`: insert, - /// dedup, and iter_ords order. - #[test] - fn bitset_accumulator_basic() { - use common::BitSet; - - use super::TermOrdAccumulator; - let mut set = ::new(255); - for ord in [0u64, 1, 63, 64, 65, 128, 200, 200, 0] { - ::insert(&mut set, ord); - } - assert_eq!(::len(&set), 7); - let collected: Vec = set.iter_ords().collect(); - assert_eq!(collected, vec![0, 1, 63, 64, 65, 128, 200]); - } - - #[test] - fn cardinality_aggregation_u64() -> crate::Result<()> { - let mut schema_builder = Schema::builder(); - let id_field = schema_builder.add_u64_field("id", FAST); - let index = Index::create_in_ram(schema_builder.build()); - { - let mut writer = index.writer_for_tests()?; - writer.add_document(doc!(id_field => 1u64))?; - writer.add_document(doc!(id_field => 2u64))?; - writer.add_document(doc!(id_field => 3u64))?; - writer.add_document(doc!())?; - writer.commit()?; - } - - let agg_req: Aggregations = serde_json::from_value(json!({ - "cardinality": { - "cardinality": { - "field": "id", - "missing": 0u64 - }, - } - })) - .unwrap(); - - let res = exec_request(agg_req, &index)?; - assert_eq!(res["cardinality"]["value"], 4.0); - - Ok(()) - } - - #[test] - fn cardinality_aggregation_ip_addr() -> crate::Result<()> { - let mut schema_builder = Schema::builder(); - let field = schema_builder.add_ip_addr_field("ip_field", FAST); - let index = Index::create_in_ram(schema_builder.build()); - { - let mut writer = index.writer_for_tests()?; - // IpV6 loopback - writer.add_document(doc!(field=>IpAddr::from_str("::1").unwrap().into_ipv6_addr()))?; - writer.add_document(doc!(field=>IpAddr::from_str("::1").unwrap().into_ipv6_addr()))?; - // IpV4 - writer.add_document( - doc!(field=>IpAddr::from_str("127.0.0.1").unwrap().into_ipv6_addr()), - )?; - writer.commit()?; - } - - let agg_req: Aggregations = serde_json::from_value(json!({ - "cardinality": { - "cardinality": { - "field": "ip_field" - }, - } - })) - .unwrap(); - - let res = exec_request(agg_req, &index)?; - assert_eq!(res["cardinality"]["value"], 2.0); - - Ok(()) - } - - #[test] - fn cardinality_aggregation_bytes_excluded_from_accessors() -> crate::Result<()> { - // `Bytes` columns are opened as raw per-segment dictionary ordinals (like `Str`), but - // unlike `Str`, cardinality has no dictionary-resolution path for them: it would hash - // the raw ordinal directly, which are segment dependant. Ignore bytes values instead of - // counting them wrong. - let mut schema_builder = Schema::builder(); - let field = schema_builder.add_bytes_field("raw", FAST); - let index = Index::create_in_ram(schema_builder.build()); - { - let mut writer = index.writer_for_tests()?; - writer.add_document(doc!(field => vec![1u8]))?; - writer.add_document(doc!(field => vec![2u8]))?; - writer.commit()?; - writer.add_document(doc!(field => vec![3u8]))?; - writer.add_document(doc!(field => vec![4u8]))?; - writer.commit()?; - } - - let agg_req: Aggregations = serde_json::from_value(json!({ - "cardinality": { - "cardinality": { - "field": "raw" - }, - } - })) - .unwrap(); - - let res = exec_request(agg_req, &index)?; - assert_eq!(res["cardinality"]["value"], 0.0); - - Ok(()) - } - - #[test] - fn cardinality_aggregation_json() -> crate::Result<()> { - let mut schema_builder = Schema::builder(); - let field = schema_builder.add_json_field("json", FAST); - let index = Index::create_in_ram(schema_builder.build()); - { - let mut writer = index.writer_for_tests()?; - writer.add_document(doc!(field => json!({"value": false})))?; - writer.add_document(doc!(field => json!({"value": true})))?; - writer.add_document(doc!(field => json!({"value": i64::from_u64(0u64)})))?; - writer.add_document(doc!(field => json!({"value": i64::from_u64(1u64)})))?; - writer.commit()?; - } - - let agg_req: Aggregations = serde_json::from_value(json!({ - "cardinality": { - "cardinality": { - "field": "json.value" - }, - } - })) - .unwrap(); - - let res = exec_request(agg_req, &index)?; - assert_eq!(res["cardinality"]["value"], 4.0); - - Ok(()) - } - - /// A JSON path that resolves to both a Str column and a numeric column - /// produces two collector instances per segment — one with `Str` buckets - /// and one with `Numeric` buckets. Their `IntermediateMetricResult`s must - /// merge into the union cardinality. - #[test] - fn cardinality_aggregation_json_str_and_numeric() -> crate::Result<()> { - let mut schema_builder = Schema::builder(); - let field = schema_builder.add_json_field("json", FAST); - let index = Index::create_in_ram(schema_builder.build()); - { - let mut writer = index.writer_for_tests()?; - writer.add_document(doc!(field => json!({"value": "hello"})))?; - writer.add_document(doc!(field => json!({"value": "world"})))?; - writer.add_document(doc!(field => json!({"value": "hello"})))?; // dup str - writer.add_document(doc!(field => json!({"value": i64::from_u64(7u64)})))?; - writer.add_document(doc!(field => json!({"value": i64::from_u64(42u64)})))?; - writer.add_document(doc!(field => json!({"value": i64::from_u64(7u64)})))?; // dup num - writer.commit()?; - } - - let agg_req: Aggregations = serde_json::from_value(json!({ - "cardinality": { - "cardinality": { - "field": "json.value" - }, - } - })) - .unwrap(); - - let res = exec_request(agg_req, &index)?; - // 4 distinct values: "hello", "world", 7, 42. - assert_eq!(res["cardinality"]["value"], 4.0); - - Ok(()) - } - - #[test] - fn cardinality_collector_serde_roundtrip() { - use super::CardinalityCollector; - - let mut collector = CardinalityCollector::default(); - collector.insert("hello"); - collector.insert("world"); - collector.insert("hello"); // duplicate - - let serialized = serde_json::to_vec(&collector).unwrap(); - let deserialized: CardinalityCollector = serde_json::from_slice(&serialized).unwrap(); - - let original_estimate = collector.finalize().unwrap(); - let roundtrip_estimate = deserialized.finalize().unwrap(); - assert_eq!(original_estimate, roundtrip_estimate); - assert_eq!(original_estimate, 2.0); - } - - #[test] - fn cardinality_collector_merge() { - use super::CardinalityCollector; - - let mut left = CardinalityCollector::default(); - left.insert("a"); - left.insert("b"); - - let mut right = CardinalityCollector::default(); - right.insert("b"); - right.insert("c"); - - left.merge_fruits(right).unwrap(); - let estimate = left.finalize().unwrap(); - assert_eq!(estimate, 3.0); - } - - /// Verifies that merging two small sketches (both in List/Set coupon mode) - /// produces an exact result — i.e. the HllUnion does not unnecessarily - /// promote to the full HLL array when the combined cardinality is small. - #[test] - fn cardinality_collector_merge_stays_exact_for_small_sets() { - use super::CardinalityCollector; - - let mut left = CardinalityCollector::default(); - for i in 0u64..50 { - left.insert(i); - } - - let mut right = CardinalityCollector::default(); - for i in 30u64..100 { - right.insert(i); - } - - left.merge_fruits(right).unwrap(); - let estimate = left.finalize().unwrap(); - // 100 distinct values (0..100). Both sketches are in Set mode (< 192 coupons), - // so the union should stay in coupon mode and give an exact count. - assert_eq!(estimate, 100.0); - } - - #[test] - fn cardinality_collector_serialize_deserialize_binary() { - use datasketches::hll::HllSketch; - - use super::CardinalityCollector; - - let mut collector = CardinalityCollector::default(); - collector.insert("apple"); - collector.insert("banana"); - collector.insert("cherry"); - - let bytes = collector.to_sketch_bytes(); - let deserialized = HllSketch::deserialize(&bytes).unwrap(); - assert!((deserialized.estimate() - 3.0).abs() < 0.01); - } - - /// Tests that the `missing` parameter correctly counts a single empty document - /// for both u64 and str columns. - #[test] - fn cardinality_aggregation_missing_value_single_empty_doc() { - let mut schema_builder = Schema::builder(); - let id_field = schema_builder.add_u64_field("id", FAST); - let name_field = schema_builder.add_text_field("name", STRING | FAST); - let index = Index::create_in_ram(schema_builder.build()); - let mut writer = index.writer_for_tests().unwrap(); - writer - .add_document(doc!(id_field=>1u64,name_field=>"some_name")) - .unwrap(); - writer.add_document(doc!()).unwrap(); - writer.commit().unwrap(); - - { - // int colum with missing value non redundant - let agg_req: Aggregations = serde_json::from_value(json!({ - "cardinality": { - "cardinality": { - "field": "id", - "missing": 42u64 - }, - } - })) - .unwrap(); - let res = exec_request(agg_req, &index).unwrap(); - assert_eq!(res["cardinality"]["value"], 2.0); - } - - { - // int colum with missing value redundant - let agg_req: Aggregations = serde_json::from_value(json!({ - "cardinality": { - "cardinality": { - "field": "id", - "missing": 1u64 - }, - } - })) - .unwrap(); - let res = exec_request(agg_req, &index).unwrap(); - assert_eq!(res["cardinality"]["value"], 1.0); - } - - { - // str colum with missing value non redundant - // With more than one segment, this is not well handled. - let agg_req: Aggregations = serde_json::from_value(json!({ - "cardinality": { - "cardinality": { - "field": "name", - "missing": "other_name" - }, - } - })) - .unwrap(); - let res = exec_request(agg_req, &index).unwrap(); - assert_eq!(res["cardinality"]["value"], 2.0); - } - - { - // str colum with missing value redundant - let agg_req: Aggregations = serde_json::from_value(json!({ - "cardinality": { - "cardinality": { - "field": "name", - "missing": "some_name" - }, - } - })) - .unwrap(); - let res = exec_request(agg_req, &index).unwrap(); - assert_eq!(res["cardinality"]["value"], 1.0); - } - - { - // str column with missing value with a number type. - let agg_req: Aggregations = serde_json::from_value(json!({ - "cardinality": { - "cardinality": { - "field": "name", - "missing": 3, - }, - } - })) - .unwrap(); - let res = exec_request(agg_req, &index).unwrap(); - assert_eq!(res["cardinality"]["value"], 2.0); - } - } - - #[test] - fn cardinality_collector_salt_differentiates_types() { - use super::CardinalityCollector; - - // Without salt, same u64 value from different column types would collide - let mut collector_bool = CardinalityCollector::new(5); // e.g. ColumnType::Bool - collector_bool.insert(0u64); // false - collector_bool.insert(1u64); // true - - let mut collector_i64 = CardinalityCollector::new(2); // e.g. ColumnType::I64 - collector_i64.insert(0u64); - collector_i64.insert(1u64); - - // Merge them - collector_bool.merge_fruits(collector_i64).unwrap(); - let estimate = collector_bool.finalize().unwrap(); - // Should be 4 because salt makes (5, 0) != (2, 0) and (5, 1) != (2, 1) - assert_eq!(estimate, 4.0); - } -} diff --git a/src/aggregation/metric/cardinality/mod.rs b/src/aggregation/metric/cardinality/mod.rs new file mode 100644 index 000000000..72d43de18 --- /dev/null +++ b/src/aggregation/metric/cardinality/mod.rs @@ -0,0 +1,609 @@ +//! Cardinality aggregation. +//! +//! * [`str_collector`] holds the `ColumnType::Str` segment collector, which accumulates term +//! ordinals and resolves them to HLL coupons at finalization. +//! * [`numeric_collector`] holds the segment collector for every other column type, which feeds +//! the HLL sketch directly during collection. +//! +//! Both segment collectors converge on the same +//! [`IntermediateMetricResult::Cardinality`] payload, so results coming from a +//! str column and from a numeric column (e.g. a JSON path that resolves to +//! both) merge through the single [`CardinalityCollector::merge_fruits`] +//! implementation here. + +mod numeric_collector; +mod str_collector; +mod term_ord_accumulator; + +use std::hash::Hash; + +use columnar::{Column, ColumnType, StrColumn}; +use common::BitSet; +use datasketches::hll::{Coupon, HllSketch, HllType, HllUnion}; +pub(crate) use numeric_collector::SegmentNumericCardinalityCollector; +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::*; + +/// Log2 of the number of registers for the HLL sketch. +/// 2^11 = 2048 registers, giving ~2.3% relative error and ~1KB per sketch (Hll4). +const LG_K: u8 = 11; + +/// # Cardinality +/// +/// The cardinality aggregation allows for computing an estimate +/// of the number of different values in a data set based on the +/// Apache DataSketches HyperLogLog algorithm. This is particularly useful for +/// understanding the uniqueness of values in a large dataset where counting +/// each unique value individually would be computationally expensive. +/// +/// For example, you might use a cardinality aggregation to estimate the number +/// of unique visitors to a website by aggregating on a field that contains +/// user IDs or session IDs. +/// +/// To use the cardinality aggregation, you'll need to provide a field to +/// aggregate on. The following example demonstrates a request for the cardinality +/// of the "user_id" field: +/// +/// ```JSON +/// { +/// "cardinality": { +/// "field": "user_id" +/// } +/// } +/// ``` +/// +/// This request will return an estimate of the number of unique values in the +/// "user_id" field. +/// +/// ## Missing Values +/// +/// The `missing` parameter defines how documents that are missing a value should be treated. +/// By default, documents without a value for the specified field are ignored. However, you can +/// specify a default value for these documents using the `missing` parameter. This can be useful +/// when you want to include documents with missing values in the aggregation. +/// +/// For example, the following request treats documents with missing values in the "user_id" +/// field as if they had a value of "unknown": +/// +/// ```JSON +/// { +/// "cardinality": { +/// "field": "user_id", +/// "missing": "unknown" +/// } +/// } +/// ``` +/// +/// # Estimation Accuracy +/// +/// The cardinality aggregation provides an approximate count, which is usually +/// accurate within a small error range. This trade-off allows for efficient +/// computation even on very large datasets. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct CardinalityAggregationReq { + /// The field name to compute the percentiles on. + pub field: String, + /// The missing parameter defines how documents that are missing a value should be treated. + /// By default they will be ignored but it is also possible to treat them as if they had a + /// value. Examples in JSON format: + /// { "field": "my_numbers", "missing": "10.0" } + #[serde(skip_serializing_if = "Option::is_none", default)] + pub missing: Option, +} + +/// Contains all information required by the segment cardinality collectors to perform the +/// cardinality aggregation on a segment. +pub(crate) struct CardinalityAggReqData { + /// The column accessor to access the fast field values. + pub(crate) accessor: Column, + /// The column_type of the field. + pub(crate) column_type: ColumnType, + /// The string dictionary column if the field is of type string. + pub(crate) str_dict_column: Option, + /// The missing value normalized to the internal u64 representation of the field type. + pub(crate) missing_value_for_accessor: Option, + /// The name of the aggregation. + pub(crate) name: String, + /// The aggregation request. + pub(crate) req: CardinalityAggregationReq, +} + +impl CardinalityAggReqData { + /// Estimate the memory consumption of this struct in bytes. + pub fn get_memory_consumption(&self) -> usize { + std::mem::size_of::() + } +} + +impl CardinalityAggregationReq { + /// Creates a new [`CardinalityAggregationReq`] instance from a field name. + pub fn from_field_name(field_name: String) -> Self { + Self { + field: field_name, + missing: None, + } + } + /// Returns the field name the aggregation is computed on. + pub fn field_name(&self) -> &str { + &self.field + } +} + +#[derive(Clone, Debug)] +/// The cardinality collector used during segment collection and for merging results. +/// Uses Apache DataSketches HLL (lg_k=11, Hll4) for compact binary serialization +/// and cross-language compatibility (e.g. Java `datasketches` library). +pub struct CardinalityCollector { + sketch: HllSketch, + /// Salt derived from `ColumnType`, used to differentiate values of different column types + /// that map to the same u64 (e.g. bool `false` = 0 vs i64 `0`). + /// Not serialized — only needed during insertion, not after sketch registers are populated. + salt: u8, +} + +impl Default for CardinalityCollector { + fn default() -> Self { + Self::new(0) + } +} + +impl PartialEq for CardinalityCollector { + fn eq(&self, _other: &Self) -> bool { + false + } +} + +impl Serialize for CardinalityCollector { + fn serialize(&self, serializer: S) -> Result { + let bytes = self.sketch.serialize(); + serializer.serialize_bytes(&bytes) + } +} + +impl<'de> Deserialize<'de> for CardinalityCollector { + fn deserialize>(deserializer: D) -> Result { + let bytes: Vec = Deserialize::deserialize(deserializer)?; + let sketch = HllSketch::deserialize(&bytes).map_err(serde::de::Error::custom)?; + Ok(Self { sketch, salt: 0 }) + } +} + +impl CardinalityCollector { + fn new(salt: u8) -> Self { + Self { + sketch: HllSketch::new(LG_K, HllType::Hll8), + salt, + } + } + + /// Insert a value into the HLL sketch, salted by the column type. + /// The salt ensures that identical u64 values from different column types + /// (e.g. bool `false` vs i64 `0`) are counted as distinct. + fn insert(&mut self, value: impl Hash) { + self.sketch.update((self.salt, value)); + } + + fn insert_coupon(&mut self, coupon: Coupon) { + self.sketch.update_with_coupon(coupon); + } + + /// Compute the final cardinality estimate. + pub fn finalize(self) -> Option { + Some(self.sketch.estimate().trunc()) + } + + /// Serialize the HLL sketch to its compact binary representation. + /// The format is cross-language compatible with Apache DataSketches (Java, C++, Python). + pub fn to_sketch_bytes(&self) -> Vec { + self.sketch.serialize() + } + + pub(crate) fn merge_fruits(&mut self, right: CardinalityCollector) -> crate::Result<()> { + let mut union = HllUnion::new(LG_K); + union.update(&self.sketch); + union.update(&right.sketch); + self.sketch = union.to_sketch(HllType::Hll8); + Ok(()) + } +} + +/// Builds the segment collector for a cardinality aggregation. +/// +/// str and non-str columns use two entirely different collectors: str +/// accumulates term ordinals and resolves them into HLL coupons at +/// 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, +) -> crate::Result> { + let req_data = req.get_cardinality_req_data(node.idx_in_req_data); + if req_data.column_type != ColumnType::Str { + return Ok(Box::new(SegmentNumericCardinalityCollector::from_req( + req_data.column_type, + node.idx_in_req_data, + req_data.accessor.clone(), + req_data.missing_value_for_accessor, + )?)); + } + // For str columns, we need to collect the set of term ordinals encounterred. + // We choose a different representation depending on the number of maximum + // number of terms. + // * small (< BITSET_MAX_TERM_ORD): `BitSet`, pre-allocated. + // * large: `TermOrdSet` (sparse HashSet that promotes to a paged bitset). + let max_term_ord_inclusive = req_data.accessor.max_value(); + if max_term_ord_inclusive < BITSET_MAX_TERM_ORD { + Ok(Box::new( + SegmentStrCardinalityCollector::::from_req( + node.idx_in_req_data, + req_data.accessor.clone(), + req_data.missing_value_for_accessor, + max_term_ord_inclusive, + ), + )) + } else { + Ok(Box::new( + SegmentStrCardinalityCollector::::from_req( + node.idx_in_req_data, + req_data.accessor.clone(), + req_data.missing_value_for_accessor, + max_term_ord_inclusive, + ), + )) + } +} + +#[cfg(test)] +mod tests { + use columnar::MonotonicallyMappableToU64; + + use crate::aggregation::agg_req::Aggregations; + use crate::aggregation::tests::{exec_request, get_test_index_from_terms}; + use crate::schema::{Schema, FAST, STRING}; + use crate::Index; + + #[test] + fn cardinality_aggregation_test_empty_index() -> crate::Result<()> { + let values = vec![]; + let index = get_test_index_from_terms(false, &values)?; + let agg_req: Aggregations = serde_json::from_value(json!({ + "cardinality": { + "cardinality": { + "field": "string_id", + } + }, + })) + .unwrap(); + + let res = exec_request(agg_req, &index)?; + assert_eq!(res["cardinality"]["value"], 0.0); + + Ok(()) + } + + #[test] + fn cardinality_aggregation_test_single_segment() -> crate::Result<()> { + cardinality_aggregation_test_merge_segment(true) + } + #[test] + fn cardinality_aggregation_test() -> crate::Result<()> { + cardinality_aggregation_test_merge_segment(false) + } + fn cardinality_aggregation_test_merge_segment(merge_segments: bool) -> crate::Result<()> { + let segment_and_terms = vec![ + vec!["terma"], + vec!["termb"], + vec!["termc"], + vec!["terma"], + vec!["terma"], + vec!["terma"], + vec!["termb"], + vec!["terma"], + ]; + let index = get_test_index_from_terms(merge_segments, &segment_and_terms)?; + let agg_req: Aggregations = serde_json::from_value(json!({ + "cardinality": { + "cardinality": { + "field": "string_id", + } + }, + })) + .unwrap(); + + let res = exec_request(agg_req, &index)?; + assert_eq!(res["cardinality"]["value"], 3.0); + + Ok(()) + } + + #[test] + fn cardinality_aggregation_bytes_excluded_from_accessors() -> crate::Result<()> { + // `Bytes` columns are opened as raw per-segment dictionary ordinals (like `Str`), but + // unlike `Str`, cardinality has no dictionary-resolution path for them: it would hash + // the raw ordinal directly, which are segment dependant. Ignore bytes values instead of + // counting them wrong. + let mut schema_builder = Schema::builder(); + let field = schema_builder.add_bytes_field("raw", FAST); + let index = Index::create_in_ram(schema_builder.build()); + { + let mut writer = index.writer_for_tests()?; + writer.add_document(doc!(field => vec![1u8]))?; + writer.add_document(doc!(field => vec![2u8]))?; + writer.commit()?; + writer.add_document(doc!(field => vec![3u8]))?; + writer.add_document(doc!(field => vec![4u8]))?; + writer.commit()?; + } + + let agg_req: Aggregations = serde_json::from_value(json!({ + "cardinality": { + "cardinality": { + "field": "raw" + }, + } + })) + .unwrap(); + + let res = exec_request(agg_req, &index)?; + assert_eq!(res["cardinality"]["value"], 0.0); + + Ok(()) + } + + #[test] + fn cardinality_aggregation_json() -> crate::Result<()> { + let mut schema_builder = Schema::builder(); + let field = schema_builder.add_json_field("json", FAST); + let index = Index::create_in_ram(schema_builder.build()); + { + let mut writer = index.writer_for_tests()?; + writer.add_document(doc!(field => json!({"value": false})))?; + writer.add_document(doc!(field => json!({"value": true})))?; + writer.add_document(doc!(field => json!({"value": i64::from_u64(0u64)})))?; + writer.add_document(doc!(field => json!({"value": i64::from_u64(1u64)})))?; + writer.commit()?; + } + + let agg_req: Aggregations = serde_json::from_value(json!({ + "cardinality": { + "cardinality": { + "field": "json.value" + }, + } + })) + .unwrap(); + + let res = exec_request(agg_req, &index)?; + assert_eq!(res["cardinality"]["value"], 4.0); + + Ok(()) + } + + /// A JSON path that resolves to both a Str column and a numeric column + /// produces two collector instances per segment — one with `Str` buckets + /// and one with `Numeric` buckets. Their `IntermediateMetricResult`s must + /// merge into the union cardinality. + #[test] + fn cardinality_aggregation_json_str_and_numeric() -> crate::Result<()> { + let mut schema_builder = Schema::builder(); + let field = schema_builder.add_json_field("json", FAST); + let index = Index::create_in_ram(schema_builder.build()); + { + let mut writer = index.writer_for_tests()?; + writer.add_document(doc!(field => json!({"value": "hello"})))?; + writer.add_document(doc!(field => json!({"value": "world"})))?; + writer.add_document(doc!(field => json!({"value": "hello"})))?; // dup str + writer.add_document(doc!(field => json!({"value": i64::from_u64(7u64)})))?; + writer.add_document(doc!(field => json!({"value": i64::from_u64(42u64)})))?; + writer.add_document(doc!(field => json!({"value": i64::from_u64(7u64)})))?; // dup num + writer.commit()?; + } + + let agg_req: Aggregations = serde_json::from_value(json!({ + "cardinality": { + "cardinality": { + "field": "json.value" + }, + } + })) + .unwrap(); + + let res = exec_request(agg_req, &index)?; + // 4 distinct values: "hello", "world", 7, 42. + assert_eq!(res["cardinality"]["value"], 4.0); + + Ok(()) + } + + #[test] + fn cardinality_collector_serde_roundtrip() { + use super::CardinalityCollector; + + let mut collector = CardinalityCollector::default(); + collector.insert("hello"); + collector.insert("world"); + collector.insert("hello"); // duplicate + + let serialized = serde_json::to_vec(&collector).unwrap(); + let deserialized: CardinalityCollector = serde_json::from_slice(&serialized).unwrap(); + + let original_estimate = collector.finalize().unwrap(); + let roundtrip_estimate = deserialized.finalize().unwrap(); + assert_eq!(original_estimate, roundtrip_estimate); + assert_eq!(original_estimate, 2.0); + } + + #[test] + fn cardinality_collector_merge() { + use super::CardinalityCollector; + + let mut left = CardinalityCollector::default(); + left.insert("a"); + left.insert("b"); + + let mut right = CardinalityCollector::default(); + right.insert("b"); + right.insert("c"); + + left.merge_fruits(right).unwrap(); + let estimate = left.finalize().unwrap(); + assert_eq!(estimate, 3.0); + } + + /// Verifies that merging two small sketches (both in List/Set coupon mode) + /// produces an exact result — i.e. the HllUnion does not unnecessarily + /// promote to the full HLL array when the combined cardinality is small. + #[test] + fn cardinality_collector_merge_stays_exact_for_small_sets() { + use super::CardinalityCollector; + + let mut left = CardinalityCollector::default(); + for i in 0u64..50 { + left.insert(i); + } + + let mut right = CardinalityCollector::default(); + for i in 30u64..100 { + right.insert(i); + } + + left.merge_fruits(right).unwrap(); + let estimate = left.finalize().unwrap(); + // 100 distinct values (0..100). Both sketches are in Set mode (< 192 coupons), + // so the union should stay in coupon mode and give an exact count. + assert_eq!(estimate, 100.0); + } + + #[test] + fn cardinality_collector_serialize_deserialize_binary() { + use datasketches::hll::HllSketch; + + use super::CardinalityCollector; + + let mut collector = CardinalityCollector::default(); + collector.insert("apple"); + collector.insert("banana"); + collector.insert("cherry"); + + let bytes = collector.to_sketch_bytes(); + let deserialized = HllSketch::deserialize(&bytes).unwrap(); + assert!((deserialized.estimate() - 3.0).abs() < 0.01); + } + + /// Tests that the `missing` parameter correctly counts a single empty document + /// for both u64 and str columns. + #[test] + fn cardinality_aggregation_missing_value_single_empty_doc() { + let mut schema_builder = Schema::builder(); + let id_field = schema_builder.add_u64_field("id", FAST); + let name_field = schema_builder.add_text_field("name", STRING | FAST); + let index = Index::create_in_ram(schema_builder.build()); + let mut writer = index.writer_for_tests().unwrap(); + writer + .add_document(doc!(id_field=>1u64,name_field=>"some_name")) + .unwrap(); + writer.add_document(doc!()).unwrap(); + writer.commit().unwrap(); + + { + // int colum with missing value non redundant + let agg_req: Aggregations = serde_json::from_value(json!({ + "cardinality": { + "cardinality": { + "field": "id", + "missing": 42u64 + }, + } + })) + .unwrap(); + let res = exec_request(agg_req, &index).unwrap(); + assert_eq!(res["cardinality"]["value"], 2.0); + } + + { + // int colum with missing value redundant + let agg_req: Aggregations = serde_json::from_value(json!({ + "cardinality": { + "cardinality": { + "field": "id", + "missing": 1u64 + }, + } + })) + .unwrap(); + let res = exec_request(agg_req, &index).unwrap(); + assert_eq!(res["cardinality"]["value"], 1.0); + } + + { + // str colum with missing value non redundant + // With more than one segment, this is not well handled. + let agg_req: Aggregations = serde_json::from_value(json!({ + "cardinality": { + "cardinality": { + "field": "name", + "missing": "other_name" + }, + } + })) + .unwrap(); + let res = exec_request(agg_req, &index).unwrap(); + assert_eq!(res["cardinality"]["value"], 2.0); + } + + { + // str colum with missing value redundant + let agg_req: Aggregations = serde_json::from_value(json!({ + "cardinality": { + "cardinality": { + "field": "name", + "missing": "some_name" + }, + } + })) + .unwrap(); + let res = exec_request(agg_req, &index).unwrap(); + assert_eq!(res["cardinality"]["value"], 1.0); + } + + { + // str column with missing value with a number type. + let agg_req: Aggregations = serde_json::from_value(json!({ + "cardinality": { + "cardinality": { + "field": "name", + "missing": 3, + }, + } + })) + .unwrap(); + let res = exec_request(agg_req, &index).unwrap(); + assert_eq!(res["cardinality"]["value"], 2.0); + } + } + + #[test] + fn cardinality_collector_salt_differentiates_types() { + use super::CardinalityCollector; + + // Without salt, same u64 value from different column types would collide + let mut collector_bool = CardinalityCollector::new(5); // e.g. ColumnType::Bool + collector_bool.insert(0u64); // false + collector_bool.insert(1u64); // true + + let mut collector_i64 = CardinalityCollector::new(2); // e.g. ColumnType::I64 + collector_i64.insert(0u64); + collector_i64.insert(1u64); + + // Merge them + collector_bool.merge_fruits(collector_i64).unwrap(); + let estimate = collector_bool.finalize().unwrap(); + // Should be 4 because salt makes (5, 0) != (2, 0) and (5, 1) != (2, 1) + assert_eq!(estimate, 4.0); + } +} diff --git a/src/aggregation/metric/cardinality/numeric_collector.rs b/src/aggregation/metric/cardinality/numeric_collector.rs new file mode 100644 index 000000000..072f922a6 --- /dev/null +++ b/src/aggregation/metric/cardinality/numeric_collector.rs @@ -0,0 +1,255 @@ +//! Segment collector for `cardinality` over any non-str column +//! (numeric, bool, date, IpAddr). +//! +//! Unlike the str case there is no dictionary to resolve, so values go +//! straight into a per-bucket HLL sketch during collection. The produced +//! [`CardinalityCollector`] is the same type the str collector produces, so +//! intermediate merging stays in the parent module. + +use std::fmt::Debug; +use std::sync::Arc; + +use columnar::column_values::CompactSpaceU64Accessor; +use columnar::{Column, ColumnType}; + +use super::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::*; +use crate::TantivyError; + +/// Segment collector for `cardinality` over any non-str column +/// (numeric, bool, date, IpAddr). +/// +/// Hidden contract: `column_type` must not be `ColumnType::Str`. Values are +/// inserted into the HLL sketch during collection, so the sketch of a bucket +/// is already complete when the bucket is finalized. +pub(crate) struct SegmentNumericCardinalityCollector { + /// Buckets are Some(_) until they get consumed by + /// `add_intermediate_aggregation_result`. + buckets: Vec>, + accessor_idx: usize, + /// The column accessor to access the fast field values. + accessor: Column, + /// The column_type of the field. + column_type: ColumnType, + /// Set iff `column_type == ColumnType::IpAddr`. Resolved once at + /// construction: the raw column values are compact-space codes that must + /// be expanded to their u128 ip representation before hashing. + compact_space_accessor: Option>, + /// The missing value normalized to the internal u64 representation of the field type. + missing_value_for_accessor: Option, +} + +impl Debug for SegmentNumericCardinalityCollector { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + f.debug_struct("SegmentNumericCardinalityCollector") + .field("column_type", &self.column_type) + .field( + "missing_value_for_accessor", + &self.missing_value_for_accessor, + ) + .finish() + } +} + +impl SegmentNumericCardinalityCollector { + pub fn from_req( + column_type: ColumnType, + accessor_idx: usize, + accessor: Column, + missing_value_for_accessor: Option, + ) -> crate::Result { + assert_ne!(column_type, ColumnType::Str); + let compact_space_accessor = if column_type == ColumnType::IpAddr { + let compact_space_accessor = accessor + .values + .clone() + .downcast_arc::() + .map_err(|_| { + TantivyError::AggregationError( + crate::aggregation::AggregationError::InternalError( + "Type mismatch: Could not downcast to CompactSpaceU64Accessor" + .to_string(), + ), + ) + })?; + Some(compact_space_accessor) + } else { + None + }; + Ok(Self { + buckets: Vec::new(), + accessor_idx, + accessor, + column_type, + compact_space_accessor, + missing_value_for_accessor, + }) + } +} + +impl SegmentAggregationCollector for SegmentNumericCardinalityCollector { + fn add_intermediate_aggregation_result( + &mut self, + agg_data: &AggregationsSegmentCtx, + results: &mut IntermediateAggregationResults, + 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(); + // 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( + "the same bucket should not be finalized twice.".to_string(), + )); + }; + results.push( + name, + IntermediateAggregationResult::Metric(IntermediateMetricResult::Cardinality( + cardinality, + )), + )?; + Ok(()) + } + + fn collect( + &mut self, + parent_bucket_id: BucketId, + docs: &[crate::DocId], + agg_data: &mut AggregationsSegmentCtx, + ) -> crate::Result<()> { + agg_data.column_block_accessor.fetch_block_with_missing( + docs, + &self.accessor, + self.missing_value_for_accessor, + ); + let cardinality = self.buckets[parent_bucket_id as usize] + .as_mut() + .ok_or_else(|| { + crate::TantivyError::InternalError( + "collection should not happen after finalization".to_string(), + ) + })?; + let col_block_accessor = &agg_data.column_block_accessor; + if let Some(compact_space_accessor) = self.compact_space_accessor.as_ref() { + for val in col_block_accessor.iter_vals() { + let val: u128 = compact_space_accessor.compact_to_u128(val as u32); + cardinality.insert(val); + } + } else { + for val in col_block_accessor.iter_vals() { + cardinality.insert(val); + } + } + Ok(()) + } + + fn prepare_max_bucket( + &mut self, + max_bucket: BucketId, + _agg_data: &AggregationsSegmentCtx, + ) -> crate::Result<()> { + if max_bucket as usize >= self.buckets.len() { + let column_type = self.column_type; + self.buckets.resize_with(max_bucket as usize + 1, || { + Some(CardinalityCollector::new(column_type as u8)) + }); + } + Ok(()) + } + + fn compute_metric_value( + &self, + bucket_id: BucketId, + sub_agg_name: &str, + sub_agg_property: &str, + agg_data: &AggregationsSegmentCtx, + ) -> Option { + let req_data = &agg_data.get_cardinality_req_data(self.accessor_idx); + if req_data.name != sub_agg_name || !sub_agg_property.is_empty() { + return None; + } + let cardinality = self.buckets.get(bucket_id as usize)?.as_ref()?; + Some(cardinality.sketch.estimate().trunc()) + } +} + +#[cfg(test)] +mod tests { + use std::net::IpAddr; + use std::str::FromStr; + + use crate::aggregation::agg_req::Aggregations; + use crate::aggregation::tests::exec_request; + use crate::schema::{IntoIpv6Addr, Schema, FAST}; + use crate::Index; + + #[test] + fn cardinality_aggregation_u64() -> crate::Result<()> { + let mut schema_builder = Schema::builder(); + let id_field = schema_builder.add_u64_field("id", FAST); + let index = Index::create_in_ram(schema_builder.build()); + { + let mut writer = index.writer_for_tests()?; + writer.add_document(doc!(id_field => 1u64))?; + writer.add_document(doc!(id_field => 2u64))?; + writer.add_document(doc!(id_field => 3u64))?; + writer.add_document(doc!())?; + writer.commit()?; + } + + let agg_req: Aggregations = serde_json::from_value(json!({ + "cardinality": { + "cardinality": { + "field": "id", + "missing": 0u64 + }, + } + })) + .unwrap(); + + let res = exec_request(agg_req, &index)?; + assert_eq!(res["cardinality"]["value"], 4.0); + + Ok(()) + } + + #[test] + fn cardinality_aggregation_ip_addr() -> crate::Result<()> { + let mut schema_builder = Schema::builder(); + let field = schema_builder.add_ip_addr_field("ip_field", FAST); + let index = Index::create_in_ram(schema_builder.build()); + { + let mut writer = index.writer_for_tests()?; + // IpV6 loopback + writer.add_document(doc!(field=>IpAddr::from_str("::1").unwrap().into_ipv6_addr()))?; + writer.add_document(doc!(field=>IpAddr::from_str("::1").unwrap().into_ipv6_addr()))?; + // IpV4 + writer.add_document( + doc!(field=>IpAddr::from_str("127.0.0.1").unwrap().into_ipv6_addr()), + )?; + writer.commit()?; + } + + let agg_req: Aggregations = serde_json::from_value(json!({ + "cardinality": { + "cardinality": { + "field": "ip_field" + }, + } + })) + .unwrap(); + + let res = exec_request(agg_req, &index)?; + assert_eq!(res["cardinality"]["value"], 2.0); + + Ok(()) + } +} diff --git a/src/aggregation/metric/cardinality/str_collector.rs b/src/aggregation/metric/cardinality/str_collector.rs new file mode 100644 index 000000000..0ff5af74e --- /dev/null +++ b/src/aggregation/metric/cardinality/str_collector.rs @@ -0,0 +1,397 @@ +//! Segment collector for `cardinality` over a `ColumnType::Str` column. +//! +//! Strings are dictionary encoded, and resolving a term ordinal to its bytes +//! is the expensive part. So instead of hashing values during collection, the +//! collector accumulates *term ordinals* per bucket and, at finalization, +//! builds one shared term_ord -> coupon cache for every bucket at once. The +//! coupons are then appended into a [`CardinalityCollector`], which is the +//! same type the numeric collector produces — merging across segments and +//! across column kinds stays in the parent module. + +use std::fmt::Debug; +use std::io; + +use columnar::{Column, ColumnType, Dictionary}; +use datasketches::hll::Coupon; +use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet}; + +use super::term_ord_accumulator::TermOrdAccumulator; +use super::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::*; + +/// A CouponCache is here to cache the mapping term ordinal -> coupon (see above). +/// The idea is that we do not want to fetch terms associated to several term ordinals, +/// several times due to the fact that we have several buckets. +enum CouponCache { + Dense { + coupon_map: Vec, + missing_coupon_opt: Option, + }, + Sparse { + coupon_map: FxHashMap, + missing_coupon_opt: Option, + }, +} + +impl CouponCache { + fn new( + term_ords: Vec, + coupons: Vec, + missing_coupon_opt: Option, + ) -> CouponCache { + let num_terms = term_ords.len(); + assert_eq!(num_terms, coupons.len()); + if term_ords.is_empty() { + return CouponCache::Dense { + coupon_map: Vec::new(), + missing_coupon_opt, + }; + } + let highest_term_ord = term_ords.last().copied().unwrap_or(0u64); + // We prefer the dense implementation, if it is not too wasteful. + // There are two cases for which we can use it. + // 1- if the data is small. + // 2- if the data is not necessarily small, but due to a high occupancy ratio, the RAM usage + // is not that much bigger than if we had used a HashSet. (occupancy ratio + extra + // metadata ~ x2.25) + let should_use_dense = + highest_term_ord < 1_000_000u64 || highest_term_ord < num_terms as u64 * 3u64; + if should_use_dense { + // We don't really care about the value here. We will populate all the values we will + // read anyway. + let uninitialized_coupon = Coupon::from_hash(0); + let mut coupon_map: Vec = + vec![uninitialized_coupon; highest_term_ord as usize + 1]; + + for (term_ord, coupon) in term_ords.into_iter().zip(coupons) { + coupon_map[term_ord as usize] = coupon; + } + CouponCache::Dense { + coupon_map, + missing_coupon_opt, + } + } else { + let coupon_map: FxHashMap = term_ords.into_iter().zip(coupons).collect(); + CouponCache::Sparse { + coupon_map, + missing_coupon_opt, + } + } + } +} + +/// Segment collector for `cardinality` over a `ColumnType::Str` column. +/// +/// Hidden contract: the column passed at construction must be a str column +/// whose values are term ordinals of the associated dictionary, and +/// `max_term_ord_inclusive` must be `accessor.max_value()`. The missing +/// sentinel `accessor.max_value() + 1` may additionally be inserted when +/// `missing_value_for_accessor` is set, hence accumulators size for +/// `max_term_ord_inclusive + 1`. +pub(crate) struct SegmentStrCardinalityCollector { + /// Buckets are Some(_) until they get consumed by + /// `add_intermediate_aggregation_result`. + buckets: Vec>, + accessor_idx: usize, + /// The column accessor to access the fast field values (term ordinals). + accessor: Column, + /// The missing value normalized to the internal u64 representation of the field type. + missing_value_for_accessor: Option, + /// Lazily built at finalization time, shared by every bucket. + coupon_cache: Option, + /// Largest term_ord that may be inserted into a bucket, i.e. + /// `accessor.max_value()`. + max_term_ord_inclusive: u64, +} + +impl Debug for SegmentStrCardinalityCollector { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + f.debug_struct("SegmentStrCardinalityCollector") + .field("num_buckets", &self.buckets.len()) + .field( + "missing_value_for_accessor", + &self.missing_value_for_accessor, + ) + .finish() + } +} + +/// Builds a coupon cache from the given buckets, dictionary, and optional missing value. +/// Returns a mapping from term_ord to the hash (coupon) of the associated term. +fn build_coupon_cache( + buckets: &[Option], + dictionary: &Dictionary, + missing_value_opt: Option<&Key>, +) -> io::Result { + // Pass 1 computes the capacity hint, pass 2 inserts. + let mut max_bucket_len = 0usize; + for bucket in buckets.iter().flatten() { + max_bucket_len = max_bucket_len.max(bucket.len()); + } + let mut term_ords_set = FxHashSet::with_capacity_and_hasher(max_bucket_len * 2, FxBuildHasher); + for bucket in buckets.iter().flatten() { + term_ords_set.extend(bucket.iter_ords()); + } + let mut term_ords: Vec = term_ords_set.into_iter().collect(); + term_ords.sort_unstable(); + + term_ords.pop_if(|highest_term_ord| *highest_term_ord >= dictionary.num_terms() as u64); + + let mut coupons: Vec = Vec::with_capacity(term_ords.len()); + let all_term_ords_found: bool = + dictionary.sorted_ords_to_term_cb(&term_ords, |term_bytes| { + let coupon: Coupon = Coupon::from_hash(term_bytes); + coupons.push(coupon); + })?; + assert!(all_term_ords_found); + + // Regardless of whether or not there is effectively a missing value in one of the buckets, + // we populate the cache with the missing key too (if any). + let missing_coupon_opt: Option = missing_value_opt.map(|missing_key| { + if let Key::Str(missing_value_str) = missing_key { + Coupon::from_hash(missing_value_str.as_bytes()) + } else { + // See https://github.com/quickwit-oss/tantivy/issues/2891 + // A missing key with a type different from Str will not work as intended + // for the moment. + // + // Right now this is just a partial workaround. + Coupon::from_hash("__tantivy_missing_non_str__".as_bytes()) + } + }); + Ok(CouponCache::new(term_ords, coupons, missing_coupon_opt)) +} + +fn append_to_sketch( + term_ords: &impl TermOrdAccumulator, + coupon_cache: &CouponCache, + sketch: &mut CardinalityCollector, +) { + match coupon_cache { + CouponCache::Dense { + coupon_map, + missing_coupon_opt, + } => { + if let Some(missing_coupon) = missing_coupon_opt { + for term_ord in term_ords.iter_ords() { + let coupon: Coupon = coupon_map + .get(term_ord as usize) + .copied() + .unwrap_or(*missing_coupon); + sketch.insert_coupon(coupon); + } + } else { + for term_ord in term_ords.iter_ords() { + if let Some(coupon) = coupon_map.get(term_ord as usize).copied() { + sketch.insert_coupon(coupon); + } + } + } + } + CouponCache::Sparse { + coupon_map, + missing_coupon_opt, + } => { + for term_ord in term_ords.iter_ords() { + if let Some(coupon) = coupon_map.get(&term_ord).copied().or(*missing_coupon_opt) { + sketch.insert_coupon(coupon); + } + } + } + } +} + +impl SegmentStrCardinalityCollector { + pub fn from_req( + accessor_idx: usize, + accessor: Column, + missing_value_for_accessor: Option, + max_term_ord_inclusive: u64, + ) -> Self { + Self { + buckets: Vec::new(), + accessor_idx, + accessor, + missing_value_for_accessor, + coupon_cache: None, + max_term_ord_inclusive, + } + } +} + +impl SegmentAggregationCollector + for SegmentStrCardinalityCollector +{ + fn add_intermediate_aggregation_result( + &mut self, + agg_data: &AggregationsSegmentCtx, + results: &mut IntermediateAggregationResults, + 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 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(), + )); + }; + // Strings are dictionary encoded. Fetching the terms associated to strings + // is expensive. For this reason, we do that once for all buckets and cache the results + // here. + // + // The cache maps a term_ord to the hash of the associated term. The missing value + // sentinel will be associated to the hash of the missing value if any. + if self.coupon_cache.is_none() { + self.coupon_cache = Some(build_coupon_cache( + &self.buckets, + str_dict_column.dictionary(), + req_data.req.missing.as_ref(), + )?); + } + let name = req_data.name.to_string(); + // take the bucket in buckets and replace it with a new empty one + let Some(term_ords) = self.buckets[bucket_id as usize].take() else { + return Err(crate::TantivyError::InternalError( + "the same bucket should not be finalized twice.".to_string(), + )); + }; + let mut cardinality = CardinalityCollector::new(ColumnType::Str as u8); + if let Some(coupon_cache) = self.coupon_cache.as_ref() { + append_to_sketch(&term_ords, coupon_cache, &mut cardinality); + } + results.push( + name, + IntermediateAggregationResult::Metric(IntermediateMetricResult::Cardinality( + cardinality, + )), + )?; + + Ok(()) + } + + fn collect( + &mut self, + parent_bucket_id: BucketId, + docs: &[crate::DocId], + agg_data: &mut AggregationsSegmentCtx, + ) -> crate::Result<()> { + agg_data.column_block_accessor.fetch_block_with_missing( + docs, + &self.accessor, + self.missing_value_for_accessor, + ); + let Some(term_ords) = self.buckets[parent_bucket_id as usize].as_mut() else { + return Err(crate::TantivyError::InternalError( + "collection should not happen after finalization".to_string(), + )); + }; + // Promotion check runs on the pre-block state: the first call + // sees an empty set (no-op), and the last block of inserts + // doesn't trigger a promotion of a set we won't grow further. + // The trait dispatches once per block (via `extend_from_iter`) + // for adaptive variants and inlines to a tight loop for the + // BitSet path. + term_ords.maybe_compact(); + term_ords.extend_from_iter(agg_data.column_block_accessor.iter_vals()); + Ok(()) + } + + fn prepare_max_bucket( + &mut self, + max_bucket: BucketId, + _agg_data: &AggregationsSegmentCtx, + ) -> crate::Result<()> { + if max_bucket as usize >= self.buckets.len() { + let max_term_ord_inclusive = self.max_term_ord_inclusive; + self.buckets.resize_with(max_bucket as usize + 1, || { + Some(S::new(max_term_ord_inclusive)) + }); + } + Ok(()) + } + + fn compute_metric_value( + &self, + bucket_id: BucketId, + sub_agg_name: &str, + sub_agg_property: &str, + agg_data: &AggregationsSegmentCtx, + ) -> Option { + let req_data = &agg_data.get_cardinality_req_data(self.accessor_idx); + if req_data.name != sub_agg_name || !sub_agg_property.is_empty() { + return None; + } + // The sketch isn't built until finalization; the term_ord set's len is + // the exact distinct count. + let term_ords = self.buckets.get(bucket_id as usize)?.as_ref()?; + Some(term_ords.len() as f64) + } +} + +#[cfg(test)] +mod tests { + use crate::aggregation::agg_req::Aggregations; + use crate::aggregation::tests::{exec_request, get_test_index_from_terms}; + use crate::schema::{Schema, FAST, STRING}; + use crate::Index; + + /// Build a single-segment string-cardinality index with 32 unique terms. + /// `column.max_value() = 31` is well below `BITSET_MAX_TERM_ORD`, + /// so the bucket exercises the `BitSet` path end to end. + #[test] + fn cardinality_aggregation_test_str_bitset() -> crate::Result<()> { + let terms: Vec = (0..32).map(|i| format!("term_{i}")).collect(); + let term_refs: Vec> = terms.iter().map(|t| vec![t.as_str()]).collect::>(); + // single segment so we have a single dictionary of 32 terms. + let index = get_test_index_from_terms(true, &term_refs)?; + + let agg_req: Aggregations = serde_json::from_value(json!({ + "cardinality": { + "cardinality": { "field": "string_id" } + }, + })) + .unwrap(); + + let res = exec_request(agg_req, &index)?; + assert_eq!(res["cardinality"]["value"], 32.0); + Ok(()) + } + + /// `BitSet` path with a `missing` parameter: the column-level missing + /// sentinel (`column.max_value() + 1`) flows into the bitset, the + /// dict lookup filter at finalization drops it, and the missing + /// coupon is applied separately. + #[test] + fn cardinality_aggregation_test_str_bitset_with_missing() { + let mut schema_builder = Schema::builder(); + let name_field = schema_builder.add_text_field("name", STRING | FAST); + let index = Index::create_in_ram(schema_builder.build()); + let mut writer = index.writer_for_tests().unwrap(); + for i in 0..16 { + let term = format!("t{i:02}"); + writer.add_document(doc!(name_field => term)).unwrap(); + } + // One empty doc, exercising the missing sentinel. + writer.add_document(doc!()).unwrap(); + writer.commit().unwrap(); + + let agg_req: Aggregations = serde_json::from_value(json!({ + "cardinality": { + "cardinality": { + "field": "name", + "missing": "MISSING_SENTINEL_KEY", + } + }, + })) + .unwrap(); + + let res = exec_request(agg_req, &index).unwrap(); + // 16 distinct real terms + 1 distinct "missing" value = 17. + assert_eq!(res["cardinality"]["value"], 17.0); + } +} diff --git a/src/aggregation/metric/cardinality/term_ord_accumulator.rs b/src/aggregation/metric/cardinality/term_ord_accumulator.rs new file mode 100644 index 000000000..f35d6148c --- /dev/null +++ b/src/aggregation/metric/cardinality/term_ord_accumulator.rs @@ -0,0 +1,343 @@ +//! Per-bucket term-ordinal accumulators used by the str cardinality +//! collector. +//! +//! A str cardinality bucket needs a *set of term ordinals*, and the right +//! representation depends on how large the segment's dictionary slice is: +//! +//! * [`BitSet`] (from `common`): used when `column.max_value()` is small (< +//! [`BITSET_MAX_TERM_ORD`]). Pre-allocated, no promotion machinery. +//! * [`TermOrdSet`]: adaptive. Starts as an `FxHashSet` (cheap when few ords are seen) and +//! promotes to a [`PagedBitset`] once occupancy crosses the density threshold. +//! +//! Both are exposed through the [`TermOrdAccumulator`] trait so that +//! `SegmentStrCardinalityCollector` can be generic over the choice and the hot +//! `collect()` loop monomorphizes to a direct call (no enum dispatch per +//! insert). + +use common::{BitSet, TinySet}; +use rustc_hash::FxHashSet; + +/// Promote FxHashSet -> PagedBitset at ~3% density (`len * 32 > +/// dict_num_terms`). Past this point the bitset (~`dict_num_terms / 7.5` +/// bytes) is smaller than the hashset (~10 B/entry minimum) and avoids +/// the per-insert hash. +const PROMOTION_RATIO: u64 = 32; + +// ================================================================= +// PagedBitset: a sparse bitset indexed by term_ord. +// +// Used as the dense alternative to FxHashSet once a string +// cardinality bucket has accumulated enough unique term ordinals. +// Memory is bounded to (touched pages) * (page bytes), not +// (max_term_ord / 8). +// +// Page geometry mirrors `PagedTermMap` in `term_agg.rs`: 1024 ords +// per page, lazy `Vec>>` directory. +// ================================================================= +const BITSET_PAGE_SHIFT: u32 = 10; +const BITSET_PAGE_BITS: u64 = 1u64 << BITSET_PAGE_SHIFT; // 1024 +const BITSET_PAGE_MASK: u64 = BITSET_PAGE_BITS - 1; +const BITSET_WORDS_PER_PAGE: usize = (BITSET_PAGE_BITS / 64) as usize; // 16 + +#[derive(Clone)] +struct PagedBitsetPage { + words: [TinySet; BITSET_WORDS_PER_PAGE], +} + +impl PagedBitsetPage { + fn new() -> Self { + Self { + words: [TinySet::empty(); BITSET_WORDS_PER_PAGE], + } + } +} + +pub(crate) struct PagedBitset { + pages: Vec>>, + /// Cached number of set bits, maintained on insert. + count: u64, +} + +impl PagedBitset { + /// Allocates a directory big enough to hold ords up to and including + /// `max_term_ord`. Pages are allocated lazily on first set. + fn with_max_term_ord(max_term_ord: u64) -> Self { + let max_page_idx = (max_term_ord >> BITSET_PAGE_SHIFT) as usize; + let num_pages = max_page_idx + 1; + Self { + pages: vec![None; num_pages], + count: 0, + } + } + + #[inline] + fn insert(&mut self, term_ord: u64) { + let page_idx = (term_ord >> BITSET_PAGE_SHIFT) as usize; + let intra = term_ord & BITSET_PAGE_MASK; + let word_idx = (intra >> 6) as usize; + let bit_idx = (intra & 63) as u32; + + let page = match &mut self.pages[page_idx] { + Some(p) => p, + None => { + self.pages[page_idx] = Some(Box::new(PagedBitsetPage::new())); + self.pages[page_idx].as_mut().unwrap() + } + }; + if page.words[word_idx].insert_mut(bit_idx) { + self.count += 1; + } + } + + /// Number of set bits. O(1). + #[inline] + fn len(&self) -> u64 { + self.count + } + + /// Iterate set ords in ascending order. + fn iter_sorted(&self) -> impl Iterator + '_ { + self.pages + .iter() + .enumerate() + .filter_map(|(page_idx, page_opt)| page_opt.as_ref().map(|p| (page_idx, p))) + .flat_map(|(page_idx, page)| { + let page_base_ord = (page_idx as u64) << BITSET_PAGE_SHIFT; + page.words + .iter() + .enumerate() + .flat_map(move |(word_idx, &word)| { + let word_base_ord = page_base_ord + (word_idx as u64) * 64; + word.into_iter() + .map(move |bit| word_base_ord + u64::from(bit)) + }) + }) + } +} + +/// Threshold below which we use `BitSet` instead of `TermOrdSet`. +/// +/// Both `BitSet` and `FxHashSet` have the same 32-byte struct, so the comparison is heap only: +/// * `BitSet` at T=256: 5 `TinySet` words covering 258 bits (with the missing-value sentinel) = +/// 40 bytes. +/// * `FxHashSet` after one insert: 4-bucket hashbrown table ≈ 56 bytes +pub(crate) const BITSET_MAX_TERM_ORD: u64 = 256; + +// ================================================================= +// TermOrdAccumulator: per-bucket abstraction over the entries set. +// +// Implementations: +// - `BitSet` (from `common`): used when `column.max_value()` is small (< BITSET_MAX_TERM_ORD). +// Pre-allocated, no promotion. +// - `TermOrdSet`: adaptive, starts as FxHashSet and promotes to a paged bitset when occupancy +// crosses the density threshold (only if promotion is enabled — typically gated on top-level +// aggregation). +// +// The trait lets `SegmentStrCardinalityCollector` be generic over the choice +// so the hot collect() loop monomorphizes to a direct call (no enum +// dispatch per insert). +// ================================================================= +pub(crate) trait TermOrdAccumulator: Sized { + /// Construct an empty accumulator. + /// `max_term_ord_inclusive` is the largest term_ord that may be + /// inserted (used to size pre-allocated bitsets and the dense bitset + /// on promotion). + fn new(max_term_ord_inclusive: u64) -> Self; + fn insert(&mut self, term_ord: u64); + fn extend_from_iter(&mut self, ords: impl IntoIterator); + /// Hook called once per ingested block. Adaptive impls use this to + /// decide on sparse->dense promotion. + fn maybe_compact(&mut self) {} + fn len(&self) -> usize; + fn iter_ords(&self) -> impl Iterator + '_; +} + +impl TermOrdAccumulator for BitSet { + #[inline] + fn new(max_term_ord_inclusive: u64) -> Self { + // `BitSet::with_max_value(M)` accepts ords in [0, M). + // We need ords up to and including `max_term_ord_inclusive`, plus + // the missing-value sentinel `column.max_value() + 1`. + BitSet::with_max_value((max_term_ord_inclusive + 2) as u32) + } + #[inline] + fn insert(&mut self, term_ord: u64) { + BitSet::insert(self, term_ord as u32); + } + #[inline] + fn len(&self) -> usize { + BitSet::len(self) + } + fn iter_ords(&self) -> impl Iterator + '_ { + // `BitSet` itself doesn't expose iteration, but + // `BitSet::tinyset(bucket)` does. Walk per-bucket and yield each + // set bit. The capacity is `max_value()`; iterating to + // `div_ceil(64)` covers every possible ord exactly once. + let num_buckets = self.max_value().div_ceil(64); + (0..num_buckets).flat_map(move |bucket| { + let chunk_base = u64::from(bucket) * 64; + self.tinyset(bucket) + .into_iter() + .map(move |bit| chunk_base + u64::from(bit)) + }) + } + #[inline(never)] //< required to not have a perf regression + fn extend_from_iter(&mut self, ords: impl IntoIterator) { + for ord in ords { + ::insert(self, ord); + } + } +} + +// TermOrdSet: adaptive sparse->dense accumulator. +// +// Starts as an HashSet (cheap when few ords are seen). When occupancy +// crosses `len * PROMOTION_RATIO > max_term_ord_inclusive`, drains into +// a `PagedBitset` and continues dense. +pub(crate) struct TermOrdSet { + inner: TermOrdSetInner, + /// Largest term_ord that may be inserted. Used for both sizing the + /// dense bitset on promotion and as the promotion-threshold reference. + max_term_ord_inclusive: u64, +} + +enum TermOrdSetInner { + Sparse(FxHashSet), + Dense(PagedBitset), +} + +impl TermOrdAccumulator for TermOrdSet { + fn new(max_term_ord_inclusive: u64) -> Self { + Self { + inner: TermOrdSetInner::Sparse(FxHashSet::default()), + max_term_ord_inclusive, + } + } + + #[inline] + fn insert(&mut self, term_ord: u64) { + match &mut self.inner { + TermOrdSetInner::Sparse(set) => { + set.insert(term_ord); + } + TermOrdSetInner::Dense(bitset) => bitset.insert(term_ord), + } + } + + fn extend_from_iter(&mut self, ords: impl IntoIterator) { + match &mut self.inner { + TermOrdSetInner::Sparse(set) => { + set.extend(ords); + } + TermOrdSetInner::Dense(bitset) => { + for ord in ords { + bitset.insert(ord); + } + } + } + } + + fn maybe_compact(&mut self) { + let TermOrdSetInner::Sparse(set) = &mut self.inner else { + return; + }; + if set.len() as u64 * PROMOTION_RATIO <= self.max_term_ord_inclusive { + return; + } + let mut bitset = PagedBitset::with_max_term_ord(self.max_term_ord_inclusive + 1); + let set = std::mem::take(set); + for ord in set { + bitset.insert(ord); + } + self.inner = TermOrdSetInner::Dense(bitset); + } + + fn len(&self) -> usize { + match &self.inner { + TermOrdSetInner::Sparse(set) => set.len(), + TermOrdSetInner::Dense(bitset) => bitset.len() as usize, + } + } + + fn iter_ords(&self) -> impl Iterator + '_ { + match &self.inner { + TermOrdSetInner::Sparse(set) => itertools::Either::Left(set.iter().copied()), + TermOrdSetInner::Dense(bitset) => itertools::Either::Right(bitset.iter_sorted()), + } + } +} + +#[cfg(test)] +mod tests { + use common::BitSet; + + use super::{PagedBitset, TermOrdAccumulator, TermOrdSet, PROMOTION_RATIO}; + + /// Unit-test the PagedBitset itself: cross-page inserts produce sorted + /// iteration, len() matches the inserted set, and duplicates are + /// idempotent. + #[test] + fn paged_bitset_basic() { + // Span several pages: BITSET_PAGE_BITS = 1024, so ords > 1024 land + // on the second page, > 2048 on the third, etc. + let ords = [0u64, 1, 63, 64, 1023, 1024, 1025, 4096, 4097, 9999, 10_000]; + let max_ord = *ords.iter().max().unwrap(); + let mut bitset = PagedBitset::with_max_term_ord(max_ord); + for &ord in &ords { + bitset.insert(ord); + // Idempotent: inserting again must not increase count. + bitset.insert(ord); + } + assert_eq!(bitset.len(), ords.len() as u64); + let collected: Vec = bitset.iter_sorted().collect(); + let mut expected: Vec = ords.to_vec(); + expected.sort_unstable(); + assert_eq!(collected, expected); + } + + /// Unit-test `TermOrdSet`: starts Sparse, promotes to Dense on + /// `maybe_compact` once the density threshold is crossed, and + /// `iter_ords()` yields the same set in either state. Ords spanning + /// multiple paged-bitset pages exercise the Dense iter ordering. + #[test] + fn term_ord_set_promotes_on_maybe_compact() { + // Pick max so promotion needs few inserts: len * RATIO > max with + // RATIO=32 and max=64 trips at len=3 (3*32=96 > 64). + let max_term_ord = 64u64; + let mut set = ::new(max_term_ord); + // Two inserts: should stay Sparse after maybe_compact (2 * RATIO = 64, not > 64). + set.insert(0); + set.insert(7); + set.maybe_compact(); + assert_eq!(set.len(), 2); + + // Third insert promotes on next maybe_compact. + set.insert(20); + assert_eq!(set.len(), 3); + // Sanity check: at len=3, 3 * PROMOTION_RATIO = 96 > 64. + assert!(3u64 * PROMOTION_RATIO > max_term_ord); + set.maybe_compact(); + + // Post-promotion: extending continues to work. + set.insert(15); + set.insert(15); // dup + assert_eq!(set.len(), 4); + + let mut collected: Vec = set.iter_ords().collect(); + collected.sort_unstable(); + assert_eq!(collected, vec![0, 7, 15, 20]); + } + + /// Unit-test the `BitSet` impl of `TermOrdAccumulator`: insert, + /// dedup, and iter_ords order. + #[test] + fn bitset_accumulator_basic() { + let mut set = ::new(255); + for ord in [0u64, 1, 63, 64, 65, 128, 200, 200, 0] { + ::insert(&mut set, ord); + } + assert_eq!(::len(&set), 7); + let collected: Vec = set.iter_ords().collect(); + assert_eq!(collected, vec![0, 1, 63, 64, 65, 128, 200]); + } +}