sesteves commented on code in PR #20968: URL: https://github.com/apache/datafusion/pull/20968#discussion_r2946382122
########## datafusion/functions-aggregate/src/approx_top_k.rs: ########## @@ -0,0 +1,1377 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Approximate top-k aggregate function using the Filtered Space-Saving algorithm. +//! +//! This implements a distributed-friendly approximate top-k aggregation using +//! the Filtered Space-Saving algorithm. The algorithm maintains a fixed-size summary +//! of counters plus an alpha map (filter) that remembers evicted items' frequencies. +//! +//! Usage: `approx_top_k(column, k)` +//! - `column`: The column to find the most frequent values from +//! - `k`: The number of top elements to track (required, literal integer) +//! +//! Returns: `List<Struct { value: <input_type>, count: UInt64 }>` ordered by count descending. +//! +//! Algorithm references: +//! - Filtered Space-Saving: <http://www.l2f.inesc-id.pt/~fmmb/wiki/uploads/Work/misnis.ref0a.pdf> +//! - Parallel Space Saving: <https://arxiv.org/pdf/1401.0702.pdf> +//! - Space-Saving: Metwally, Agrawal, El Abbadi. "Efficient Computation of Frequent +//! and Top-k Elements in Data Streams" (ICDT 2005) + +use std::any::Any; +use std::collections::HashMap; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, BinaryArray, Float32Array, Float64Array, Int8Array, Int16Array, + Int32Array, Int64Array, LargeStringArray, ListArray, StringArray, StructArray, + UInt8Array, UInt16Array, UInt32Array, UInt64Array, +}; +use arrow::buffer::OffsetBuffer; +use arrow::datatypes::{DataType, Field, FieldRef, Fields}; +use datafusion_common::{Result, ScalarValue}; +use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion_expr::utils::format_state_name; +use datafusion_expr::{ + Accumulator, AggregateUDFImpl, Documentation, Signature, TypeSignature, Volatility, +}; +use datafusion_macros::user_doc; + +make_udaf_expr_and_func!( + ApproxTopK, + approx_top_k, + "Returns the approximate most frequent (top-k) values and their counts using the Filtered Space-Saving algorithm.", + approx_top_k_udaf +); + +// --------------------------------------------------------------------------- +// Algorithm constants +// --------------------------------------------------------------------------- + +/// Suggested constant from the paper "Finding top-k elements in data streams", +/// chap 6, equation (24). Determines the size of the alpha map relative to the capacity. +const ALPHA_MAP_ELEMENTS_PER_COUNTER: usize = 6; + +/// Limit the max alpha value to avoid overflow with merges or weighted additions. +const MAX_ALPHA_VALUE: u64 = u32::MAX as u64; + +/// Maximum allowed value for k in `approx_top_k(column, k)`. +const APPROX_TOP_K_MAX_K: usize = 10_000; + +/// Capacity multiplier for internal tracking (matches ClickHouse's default). +/// +/// We track more items internally than k to improve accuracy. +/// If user asks for top-5, we internally track top `5 * 3 = 15` items. +/// Memory impact: ~100 bytes per counter, so top-100 uses ~30 KB per accumulator. +const CAPACITY_MULTIPLIER: usize = 3; + +// --------------------------------------------------------------------------- +// SpaceSavingSummary (core algorithm) +// --------------------------------------------------------------------------- + +/// Counter entry in the Filtered Space-Saving summary. +/// +/// Each entry tracks an item, its estimated count, and the error bound. +/// The algorithm guarantees that the true count lies within `[count - error, count]`. +#[derive(Debug, Clone)] +struct Counter { + /// The serialized bytes representing the tracked item. + item: Vec<u8>, + /// FNV-1a hash of the item (cached to avoid recomputation). + hash: u64, + /// The estimated frequency count (may overestimate due to eviction handling). + count: u64, + /// The maximum possible overestimation (error bound). + error: u64, +} + +impl Counter { + /// Compare counters for sorting: higher `(count - error)` wins, + /// then higher `count` breaks ties. + fn is_greater_than(&self, other: &Counter) -> bool { + let self_lb = self.count.saturating_sub(self.error); + let other_lb = other.count.saturating_sub(other.error); + (self_lb > other_lb) || (self_lb == other_lb && self.count > other.count) + } + + /// Ordering for top-k selection: highest-ranked counters sort first. + fn cmp_by_rank(&self, other: &Counter) -> std::cmp::Ordering { + if other.is_greater_than(self) { + std::cmp::Ordering::Greater + } else if self.is_greater_than(other) { + std::cmp::Ordering::Less + } else { + std::cmp::Ordering::Equal + } + } +} + +/// Filtered Space-Saving algorithm summary for approximate top-k / heavy hitters. +/// +/// Uses a hash map for O(1) counter lookups and maintains an alpha map (filter) +/// to remember evicted items' frequencies. +#[derive(Debug, Clone)] +struct SpaceSavingSummary { + counters: Vec<Counter>, + counter_map: HashMap<Vec<u8>, usize>, + alpha_map: Vec<u64>, + requested_capacity: usize, + /// Internal target capacity to avoid frequent truncations. + /// Set to `max(64, requested_capacity * 2)`. + target_capacity: usize, +} + +impl SpaceSavingSummary { + fn compute_alpha_map_size(capacity: usize) -> usize { + (capacity * ALPHA_MAP_ELEMENTS_PER_COUNTER).next_power_of_two() + } + + /// FNV-1a hash for item bytes. Review Comment: The implementation now uses fixed-seed `QualityRandomState` (`foldhash::quality::FixedState`) for both counter lookup and alpha-bucket indexing. A fixed seed is required so independently produced partial states assign values to the same buckets and remain mergeable. ########## datafusion/functions-aggregate/src/approx_top_k.rs: ########## @@ -0,0 +1,1475 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Approximate top-k aggregate function using the Filtered Space-Saving algorithm. +//! +//! This implements a distributed-friendly approximate top-k aggregation using +//! the Filtered Space-Saving algorithm. The algorithm maintains a fixed-size summary +//! of counters plus an alpha map (filter) that remembers evicted items' frequencies. +//! +//! Usage: `approx_top_k(column, k)` +//! - `column`: The column to find the most frequent values from +//! - `k`: The number of top elements to track (required, literal integer) +//! +//! Returns: `List<Struct { value: <input_type>, count: UInt64 }>` ordered by count descending. +//! +//! Algorithm references: +//! - Filtered Space-Saving: <http://www.l2f.inesc-id.pt/~fmmb/wiki/uploads/Work/misnis.ref0a.pdf> +//! - Parallel Space Saving: <https://arxiv.org/pdf/1401.0702.pdf> +//! - Space-Saving: Metwally, Agrawal, El Abbadi. "Efficient Computation of Frequent +//! and Top-k Elements in Data Streams" (ICDT 2005) + +use std::cmp::{Ordering, max, min}; +use std::mem::size_of; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, BinaryArray, Date32Array, Date64Array, Float32Array, Float64Array, + Int8Array, Int16Array, Int32Array, Int64Array, LargeBinaryArray, LargeStringArray, + ListArray, StringArray, StructArray, TimestampMicrosecondArray, + TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, + UInt8Array, UInt16Array, UInt32Array, UInt64Array, +}; +use arrow::buffer::OffsetBuffer; +use arrow::datatypes::{DataType, Field, FieldRef, Fields, TimeUnit}; +use hashbrown::HashTable; + +use datafusion_common::{Result, ScalarValue, exec_err, plan_err}; +use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion_expr::utils::format_state_name; +use datafusion_expr::{ + Accumulator, AggregateUDFImpl, Documentation, Signature, TIMEZONE_WILDCARD, + TypeSignature, Volatility, +}; +use datafusion_macros::user_doc; + +make_udaf_expr_and_func!( + ApproxTopK, + approx_top_k, + "Returns the approximate most frequent (top-k) values and their counts using the Filtered Space-Saving algorithm.", + approx_top_k_udaf +); + +// --------------------------------------------------------------------------- +// Algorithm constants +// --------------------------------------------------------------------------- + +/// Suggested constant from the paper "Finding top-k elements in data streams", +/// chap 6, equation (24). Determines the size of the alpha map relative to the capacity. +const ALPHA_MAP_ELEMENTS_PER_COUNTER: usize = 6; + +/// Limit the max alpha value to avoid overflow with merges or weighted additions. +const MAX_ALPHA_VALUE: u64 = u32::MAX as u64; + +/// Maximum allowed value for k in `approx_top_k(column, k)`. +const APPROX_TOP_K_MAX_K: usize = 10_000; + +/// Capacity multiplier for internal tracking (matches ClickHouse's default). +/// +/// We track more items internally than k to improve accuracy. +/// If user asks for top-5, we internally track top `5 * 3 = 15` items. +/// Memory impact: ~100 bytes per counter; target_capacity is 2x the tracked +/// counter count, so top-100 uses ~60 KB per accumulator. +const CAPACITY_MULTIPLIER: usize = 3; + +// --------------------------------------------------------------------------- +// SpaceSavingSummary (core algorithm) +// --------------------------------------------------------------------------- + +/// Counter entry in the Filtered Space-Saving summary. +/// +/// Each entry tracks an item, its estimated count, and the error bound. +/// The algorithm guarantees that the true count lies within `[count - error, count]`. +#[derive(Debug, Clone)] +struct Counter { + /// The serialized bytes representing the tracked item. + item: Vec<u8>, + /// FNV-1a hash of the item (cached to avoid recomputation). + hash: u64, + /// The estimated frequency count (may overestimate due to eviction handling). + count: u64, + /// The maximum possible overestimation (error bound). + error: u64, +} + +impl Counter { + /// Compare counters for sorting: higher `(count - error)` wins, + /// then higher `count` breaks ties. + fn is_greater_than(&self, other: &Counter) -> bool { + let self_lb = self.count.saturating_sub(self.error); + let other_lb = other.count.saturating_sub(other.error); + (self_lb > other_lb) || (self_lb == other_lb && self.count > other.count) + } + + /// Ordering for top-k selection: highest-ranked counters sort first. + fn cmp_by_rank(&self, other: &Counter) -> Ordering { + if other.is_greater_than(self) { + Ordering::Greater + } else if self.is_greater_than(other) { + Ordering::Less + } else { + Ordering::Equal + } + } +} + +/// Filtered Space-Saving algorithm summary for approximate top-k / heavy hitters. +/// +/// Uses a [`HashTable`] that stores `(hash, index)` tuples for O(1) counter +/// lookups without duplicating the key bytes. The actual item data lives in +/// `counters[index].item`. An alpha map (filter) remembers evicted items' +/// frequencies. +#[derive(Debug, Clone)] +struct SpaceSavingSummary { + counters: Vec<Counter>, + /// Maps `(cached_hash, counter_index)`. Lookups use the cached hash for + /// the fast path and fall back to byte equality via the `counters` vec. + counter_map: HashTable<(u64, usize)>, + alpha_map: Vec<u64>, + requested_capacity: usize, + /// Internal target capacity to avoid frequent truncations. + /// Set to `max(64, requested_capacity * 2)`. + target_capacity: usize, + /// Running total of heap bytes owned by counter item `Vec`s. + /// Updated on push / evict / clone so that `size()` is O(1). + item_heap_bytes: usize, +} + +impl SpaceSavingSummary { + fn compute_alpha_map_size(capacity: usize) -> usize { + (capacity * ALPHA_MAP_ELEMENTS_PER_COUNTER).next_power_of_two() + } + + /// FNV-1a hash for item bytes. + fn hash_item(item: &[u8]) -> u64 { + let mut hash: u64 = 0xcbf29ce484222325; + for &byte in item { + hash ^= byte as u64; + hash = hash.wrapping_mul(0x100000001b3); + } + hash + } + + fn new(capacity: usize) -> Self { + Self { + counters: Vec::new(), + counter_map: HashTable::new(), + alpha_map: Vec::new(), + requested_capacity: 0, + target_capacity: 0, + item_heap_bytes: 0, + } + .resized(capacity) + } + + fn resized(mut self, new_capacity: usize) -> Self { + if self.requested_capacity != new_capacity { + debug_assert!(self.counters.is_empty()); + let alpha_map_size = Self::compute_alpha_map_size(new_capacity); + self.alpha_map = vec![0u64; alpha_map_size]; + self.requested_capacity = new_capacity; + self.target_capacity = max(64, new_capacity.saturating_mul(2)); + self.counters.reserve(self.target_capacity); + } + self + } + + fn is_empty(&self) -> bool { + self.counters.is_empty() + } + + #[cfg(test)] + fn len(&self) -> usize { + self.counters.len() + } + + #[cfg(test)] + fn capacity(&self) -> usize { + self.requested_capacity + } + + /// Find a counter by item bytes, using the pre-computed hash for fast + /// lookup and falling back to byte equality for collision resolution. + fn find_counter_idx(&self, item: &[u8], hash: u64) -> Option<usize> { + self.counter_map + .find(hash, |&(h, idx)| { + h == hash && self.counters[idx].item == item + }) + .map(|&(_, idx)| idx) + } + + #[cfg(test)] + fn find_counter(&self, item: &[u8]) -> Option<&Counter> { + let hash = Self::hash_item(item); + self.find_counter_idx(item, hash) + .map(|idx| &self.counters[idx]) + } + + /// Add an item with increment 1. + fn add(&mut self, item: &[u8]) { + self.insert(item, 1, 0); + } + + /// Core insertion algorithm from Filtered Space-Saving. + fn insert(&mut self, item: &[u8], increment: u64, error: u64) { + let hash = Self::hash_item(item); + + // Fast path: item already tracked. + if let Some(idx) = self.find_counter_idx(item, hash) { + self.counters[idx].count = self.counters[idx].count.saturating_add(increment); + self.counters[idx].error = self.counters[idx].error.saturating_add(error); + return; + } + + // Below capacity: add directly. + if self.counters.len() < self.requested_capacity { + self.push_counter(item.to_vec(), hash, increment, error); + return; + } + + // At capacity: use alpha map for historical frequency. + let alpha_mask = self.alpha_map.len() - 1; + let alpha_idx = (hash as usize) & alpha_mask; + let alpha = self.alpha_map[alpha_idx]; + + self.push_counter( + item.to_vec(), + hash, + alpha.saturating_add(increment), + alpha.saturating_add(error), + ); + } + + fn push_counter(&mut self, item: Vec<u8>, hash: u64, count: u64, error: u64) { + let idx = self.counters.len(); + self.item_heap_bytes += item.capacity(); + self.counter_map + .insert_unique(hash, (hash, idx), |&(h, _)| h); + self.counters.push(Counter { + item, + hash, + count, + error, + }); + self.truncate_if_needed(false); + } + + /// Truncate counters when `target_capacity` is reached, + /// updating the alpha map with evicted items' true counts. + fn truncate_if_needed(&mut self, force_rebuild: bool) { + let need_truncate = self.counters.len() >= self.target_capacity; + + if need_truncate { + let k = self.requested_capacity; + if k < self.counters.len() { + self.counters + .select_nth_unstable_by(k - 1, |a, b| a.cmp_by_rank(b)); + + let alpha_mask = self.alpha_map.len() - 1; + for counter in self.counters.drain(k..) { + let alpha_idx = (counter.hash as usize) & alpha_mask; + let true_count = counter.count.saturating_sub(counter.error); + self.alpha_map[alpha_idx] = min( + self.alpha_map[alpha_idx].saturating_add(true_count), + MAX_ALPHA_VALUE, + ); + self.item_heap_bytes -= counter.item.capacity(); + } + } + } + + if need_truncate || force_rebuild { + self.rebuild_counter_map(); + } + } + + /// Rebuild the `counter_map` from the current `counters` vec. + fn rebuild_counter_map(&mut self) { + self.counter_map.clear(); + for (idx, counter) in self.counters.iter().enumerate() { + self.counter_map.insert_unique( + counter.hash, + (counter.hash, idx), + |&(h, _)| h, + ); + } + } + + #[cfg(test)] + fn get(&self, item: &[u8]) -> Option<(u64, u64)> { + self.find_counter(item).map(|c| (c.count, c.error)) + } + + /// Get the top-k items sorted by (count - error) descending. + fn top_k(&self, k: usize) -> Vec<(&[u8], u64, u64)> { + if k == 0 || self.counters.is_empty() { + return Vec::new(); + } + + let mut sorted: Vec<_> = self.counters.iter().collect(); + let return_size = min(sorted.len(), k); + + if return_size < sorted.len() { + sorted.select_nth_unstable_by(return_size - 1, |a, b| a.cmp_by_rank(b)); + sorted.truncate(return_size); + } + + sorted.sort_by(|a, b| a.cmp_by_rank(b)); + + sorted + .into_iter() + .map(|c| (c.item.as_slice(), c.count, c.error)) + .collect() + } + + /// Merge another summary into this one using the Parallel Space-Saving + /// reduce-and-combine algorithm from <https://arxiv.org/pdf/1401.0702.pdf>, + /// matching ClickHouse's `SpaceSaving::merge()` implementation. + fn merge(&mut self, other: &SpaceSavingSummary) { + if other.is_empty() { + return; + } + + if self.is_empty() { + self.counters.clone_from(&other.counters); + self.counter_map.clone_from(&other.counter_map); + self.alpha_map.clone_from(&other.alpha_map); + self.requested_capacity = other.requested_capacity; + self.target_capacity = other.target_capacity; + // Recompute from cloned vecs since clone may allocate exact-size + // (capacity == len), which can differ from the original's capacity. + self.item_heap_bytes = self.counters.iter().map(|c| c.item.capacity()).sum(); + return; + } + + // Compute m1/m2: the minimum counter count in each summary. + // Per the Parallel Space-Saving paper (Theorem 1), if a summary has + // reached capacity, items not in its counter list could have had at + // most min(counter.count) frequency. This is the merge correction. + let m1 = if self.counters.len() >= self.requested_capacity { + self.counters.iter().map(|c| c.count).min().unwrap_or(0) + } else { + 0 + }; + let m2 = if other.counters.len() >= other.requested_capacity { + other.counters.iter().map(|c| c.count).min().unwrap_or(0) + } else { + 0 + }; + + // Step 1: Bump all self counters by m2 (upper bound of what they + // could have counted in the other partition). + if m2 > 0 { + for counter in &mut self.counters { + counter.count = counter.count.saturating_add(m2); + counter.error = counter.error.saturating_add(m2); + } + } + + // Step 2: Merge other's counters into self. + for other_counter in &other.counters { + if let Some(idx) = + self.find_counter_idx(&other_counter.item, other_counter.hash) + { + // Item exists in both: add other's count, subtract the m2 we + // already added in step 1 (guaranteed non-negative). + self.counters[idx].count = self.counters[idx] + .count + .saturating_add(other_counter.count.saturating_sub(m2)); + self.counters[idx].error = self.counters[idx] + .error + .saturating_add(other_counter.error.saturating_sub(m2)); + } else { + // Item only in other: add with m1 (upper bound of what it + // could have counted in self's partition). + let item = other_counter.item.clone(); + self.item_heap_bytes += item.capacity(); + self.counters.push(Counter { + item, + hash: other_counter.hash, + count: other_counter.count.saturating_add(m1), + error: other_counter.error.saturating_add(m1), + }); + } + } + + // Merge alpha maps element-wise. + if self.alpha_map.len() == other.alpha_map.len() { + for (i, &other_alpha) in other.alpha_map.iter().enumerate() { + self.alpha_map[i] = min( + self.alpha_map[i].saturating_add(other_alpha), + MAX_ALPHA_VALUE, + ); + } + } + + self.truncate_if_needed(true); + } + + /// Serialize the summary to bytes (matches ClickHouse's `write()` format). + /// + /// Only the top `requested_capacity` counters are written. The alpha map + /// carries evicted frequency information for the coordinator merge. + fn serialize(&mut self) -> Vec<u8> { + self.truncate_if_needed(false); Review Comment: With the original Filtered Space-Saving update, every streaming update leaves the summary at or below its requested capacity. `serialize()` is now non-mutating and does not call truncation; reduction is only needed after combining candidates during a distributed merge. ########## datafusion/functions-aggregate/src/approx_top_k.rs: ########## @@ -0,0 +1,1377 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Approximate top-k aggregate function using the Filtered Space-Saving algorithm. +//! +//! This implements a distributed-friendly approximate top-k aggregation using +//! the Filtered Space-Saving algorithm. The algorithm maintains a fixed-size summary +//! of counters plus an alpha map (filter) that remembers evicted items' frequencies. +//! +//! Usage: `approx_top_k(column, k)` +//! - `column`: The column to find the most frequent values from +//! - `k`: The number of top elements to track (required, literal integer) +//! +//! Returns: `List<Struct { value: <input_type>, count: UInt64 }>` ordered by count descending. +//! +//! Algorithm references: +//! - Filtered Space-Saving: <http://www.l2f.inesc-id.pt/~fmmb/wiki/uploads/Work/misnis.ref0a.pdf> +//! - Parallel Space Saving: <https://arxiv.org/pdf/1401.0702.pdf> +//! - Space-Saving: Metwally, Agrawal, El Abbadi. "Efficient Computation of Frequent +//! and Top-k Elements in Data Streams" (ICDT 2005) + +use std::any::Any; +use std::collections::HashMap; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, BinaryArray, Float32Array, Float64Array, Int8Array, Int16Array, + Int32Array, Int64Array, LargeStringArray, ListArray, StringArray, StructArray, + UInt8Array, UInt16Array, UInt32Array, UInt64Array, +}; +use arrow::buffer::OffsetBuffer; +use arrow::datatypes::{DataType, Field, FieldRef, Fields}; +use datafusion_common::{Result, ScalarValue}; +use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion_expr::utils::format_state_name; +use datafusion_expr::{ + Accumulator, AggregateUDFImpl, Documentation, Signature, TypeSignature, Volatility, +}; +use datafusion_macros::user_doc; + +make_udaf_expr_and_func!( + ApproxTopK, + approx_top_k, + "Returns the approximate most frequent (top-k) values and their counts using the Filtered Space-Saving algorithm.", + approx_top_k_udaf +); + +// --------------------------------------------------------------------------- +// Algorithm constants +// --------------------------------------------------------------------------- + +/// Suggested constant from the paper "Finding top-k elements in data streams", +/// chap 6, equation (24). Determines the size of the alpha map relative to the capacity. +const ALPHA_MAP_ELEMENTS_PER_COUNTER: usize = 6; + +/// Limit the max alpha value to avoid overflow with merges or weighted additions. +const MAX_ALPHA_VALUE: u64 = u32::MAX as u64; + +/// Maximum allowed value for k in `approx_top_k(column, k)`. +const APPROX_TOP_K_MAX_K: usize = 10_000; + +/// Capacity multiplier for internal tracking (matches ClickHouse's default). +/// +/// We track more items internally than k to improve accuracy. +/// If user asks for top-5, we internally track top `5 * 3 = 15` items. +/// Memory impact: ~100 bytes per counter, so top-100 uses ~30 KB per accumulator. +const CAPACITY_MULTIPLIER: usize = 3; + +// --------------------------------------------------------------------------- +// SpaceSavingSummary (core algorithm) +// --------------------------------------------------------------------------- + +/// Counter entry in the Filtered Space-Saving summary. +/// +/// Each entry tracks an item, its estimated count, and the error bound. +/// The algorithm guarantees that the true count lies within `[count - error, count]`. +#[derive(Debug, Clone)] +struct Counter { + /// The serialized bytes representing the tracked item. + item: Vec<u8>, + /// FNV-1a hash of the item (cached to avoid recomputation). + hash: u64, + /// The estimated frequency count (may overestimate due to eviction handling). + count: u64, + /// The maximum possible overestimation (error bound). + error: u64, +} + +impl Counter { + /// Compare counters for sorting: higher `(count - error)` wins, + /// then higher `count` breaks ties. + fn is_greater_than(&self, other: &Counter) -> bool { + let self_lb = self.count.saturating_sub(self.error); + let other_lb = other.count.saturating_sub(other.error); + (self_lb > other_lb) || (self_lb == other_lb && self.count > other.count) + } + + /// Ordering for top-k selection: highest-ranked counters sort first. + fn cmp_by_rank(&self, other: &Counter) -> std::cmp::Ordering { + if other.is_greater_than(self) { + std::cmp::Ordering::Greater + } else if self.is_greater_than(other) { + std::cmp::Ordering::Less + } else { + std::cmp::Ordering::Equal + } + } +} + +/// Filtered Space-Saving algorithm summary for approximate top-k / heavy hitters. +/// +/// Uses a hash map for O(1) counter lookups and maintains an alpha map (filter) +/// to remember evicted items' frequencies. +#[derive(Debug, Clone)] +struct SpaceSavingSummary { + counters: Vec<Counter>, + counter_map: HashMap<Vec<u8>, usize>, + alpha_map: Vec<u64>, + requested_capacity: usize, + /// Internal target capacity to avoid frequent truncations. + /// Set to `max(64, requested_capacity * 2)`. + target_capacity: usize, +} + +impl SpaceSavingSummary { + fn compute_alpha_map_size(capacity: usize) -> usize { + (capacity * ALPHA_MAP_ELEMENTS_PER_COUNTER).next_power_of_two() + } + + /// FNV-1a hash for item bytes. + fn hash_item(item: &[u8]) -> u64 { + let mut hash: u64 = 0xcbf29ce484222325; + for &byte in item { + hash ^= byte as u64; + hash = hash.wrapping_mul(0x100000001b3); + } + hash + } + + fn new(capacity: usize) -> Self { + Self { + counters: Vec::new(), + counter_map: HashMap::new(), + alpha_map: Vec::new(), + requested_capacity: 0, + target_capacity: 0, + } + .resized(capacity) + } + + fn resized(mut self, new_capacity: usize) -> Self { + if self.requested_capacity != new_capacity { + debug_assert!(self.counters.is_empty()); + let alpha_map_size = Self::compute_alpha_map_size(new_capacity); + self.alpha_map = vec![0u64; alpha_map_size]; + self.requested_capacity = new_capacity; + self.target_capacity = std::cmp::max(64, new_capacity * 2); + self.counters.reserve(self.target_capacity); + } + self + } + + fn is_empty(&self) -> bool { + self.counters.is_empty() + } + + #[cfg(test)] + fn len(&self) -> usize { + self.counters.len() + } + + #[cfg(test)] + fn capacity(&self) -> usize { + self.requested_capacity + } + + fn find_counter_mut(&mut self, item: &[u8]) -> Option<&mut Counter> { + self.counter_map + .get(item) + .copied() + .map(|idx| &mut self.counters[idx]) + } + + #[cfg(test)] + fn find_counter(&self, item: &[u8]) -> Option<&Counter> { + self.counter_map.get(item).map(|&idx| &self.counters[idx]) + } + + /// Add an item with increment 1. + fn add(&mut self, item: &[u8]) { + self.insert(item, 1, 0); + } + + /// Core insertion algorithm from Filtered Space-Saving. + fn insert(&mut self, item: &[u8], increment: u64, error: u64) { + let hash = Self::hash_item(item); + + // Fast path: item already tracked. + if let Some(counter) = self.find_counter_mut(item) { + counter.count += increment; + counter.error += error; + return; + } + + // Below capacity: add directly. + if self.counters.len() < self.requested_capacity { + self.push_counter(item.to_vec(), hash, increment, error); + return; + } + + // At capacity: use alpha map for historical frequency. + let alpha_mask = self.alpha_map.len() - 1; + let alpha_idx = (hash as usize) & alpha_mask; + let alpha = self.alpha_map[alpha_idx]; + + self.push_counter(item.to_vec(), hash, alpha + increment, alpha + error); + } + + fn push_counter(&mut self, item: Vec<u8>, hash: u64, count: u64, error: u64) { + let idx = self.counters.len(); + self.counter_map.insert(item.clone(), idx); + self.counters.push(Counter { + item, + hash, + count, + error, + }); + self.truncate_if_needed(false); + } + + /// Truncate counters when `target_capacity` is reached, + /// updating the alpha map with evicted items' true counts. + fn truncate_if_needed(&mut self, force_rebuild: bool) { + let need_truncate = self.counters.len() >= self.target_capacity; + + if need_truncate { + let k = self.requested_capacity; + if k < self.counters.len() { + self.counters + .select_nth_unstable_by(k - 1, |a, b| a.cmp_by_rank(b)); + + let alpha_mask = self.alpha_map.len() - 1; + for counter in self.counters.drain(k..) { + let alpha_idx = (counter.hash as usize) & alpha_mask; + let true_count = counter.count.saturating_sub(counter.error); + self.alpha_map[alpha_idx] = std::cmp::min( + self.alpha_map[alpha_idx] + true_count, + MAX_ALPHA_VALUE, + ); + } + } + } + + if need_truncate || force_rebuild { + self.counter_map.clear(); + for (idx, counter) in self.counters.iter().enumerate() { + self.counter_map.insert(counter.item.clone(), idx); + } + } + } + + #[cfg(test)] + fn get(&self, item: &[u8]) -> Option<(u64, u64)> { + self.find_counter(item).map(|c| (c.count, c.error)) + } + + /// Get the top-k items sorted by count descending. + fn top_k(&self, k: usize) -> Vec<(&[u8], u64, u64)> { + if k == 0 || self.counters.is_empty() { + return Vec::new(); + } + + let mut sorted: Vec<_> = self.counters.iter().collect(); + let return_size = std::cmp::min(sorted.len(), k); + + if return_size < sorted.len() { + sorted.select_nth_unstable_by(return_size - 1, |a, b| a.cmp_by_rank(b)); + sorted.truncate(return_size); + } + + sorted.sort_by(|a, b| a.cmp_by_rank(b)); + + sorted + .into_iter() + .map(|c| (c.item.as_slice(), c.count, c.error)) + .collect() + } + + /// Merge another summary into this one. + fn merge(&mut self, other: &SpaceSavingSummary) { + if other.is_empty() { + return; + } + + if self.is_empty() { + self.counters.clone_from(&other.counters); + self.counter_map.clone_from(&other.counter_map); + self.alpha_map.clone_from(&other.alpha_map); + self.requested_capacity = other.requested_capacity; + self.target_capacity = other.target_capacity; + return; + } + + for other_counter in &other.counters { + if let Some(idx) = self.counter_map.get(&other_counter.item).copied() { + self.counters[idx].count += other_counter.count; + self.counters[idx].error += other_counter.error; + } else { + self.counters.push(Counter { + item: other_counter.item.clone(), + hash: other_counter.hash, + count: other_counter.count, + error: other_counter.error, + }); + } + } + + // Merge alpha maps element-wise. Sizes should always match because the + // planner guarantees the same k (and thus the same capacity/alpha map size) + // across all partitions. If they differ due to a bug, we skip the merge + // which only degrades accuracy without affecting correctness. + if self.alpha_map.len() == other.alpha_map.len() { + for (i, &other_alpha) in other.alpha_map.iter().enumerate() { + self.alpha_map[i] = + std::cmp::min(self.alpha_map[i] + other_alpha, MAX_ALPHA_VALUE); + } + } + + self.truncate_if_needed(true); + } + + /// Serialize the summary to bytes. + fn to_bytes(&self) -> Vec<u8> { + let counters_to_write: Vec<_> = { + let mut sorted: Vec<_> = self.counters.iter().collect(); Review Comment: This changed with the original Filtered Space-Saving update. Streaming updates now reject or immediately replace the minimum counter, so the summary never exceeds its requested capacity. Serialization is non-mutating and no longer truncates. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
