You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "sesteves (via GitHub)" <gi...@apache.org> on 2026/03/16 17:17:50 UTC
[PR] feat: add approx_top_k aggregate function [datafusion]
sesteves opened a new pull request, #20968:
URL: https://github.com/apache/datafusion/pull/20968
## Which issue does this PR close?
Closes #20967.
## Rationale
Finding the most frequently occurring values in a column is a very common analytical pattern — top error codes, most popular products, most active users, frequent URL paths, etc. The exact approach (`GROUP BY value ORDER BY COUNT(*) DESC LIMIT k`) requires materializing all distinct groups, which is memory-intensive and slow on high-cardinality columns.
`approx_top_k` solves this with a streaming approximation algorithm that operates in bounded memory, making it practical for large-scale analytics.
This function is already available in ClickHouse (`topK`) and via extensions in PostgreSQL and Druid, and is a natural addition to DataFusion's existing family of approximate aggregate functions (`approx_distinct`, `approx_median`, `approx_percentile_cont`).
## What changes are included in this PR?
### Core implementation (`datafusion/functions-aggregate/src/approx_top_k.rs`)
- **`SpaceSavingSummary`** — Implements the Filtered Space-Saving algorithm (Metwally et al., 2005):
- Fixed-size counter summary with eviction of minimum-count items when full
- **Alpha map** for improved accuracy — tracks recently evicted items and filters low-frequency noise before it enters the main summary (same approach used by ClickHouse)
- `CAPACITY_MULTIPLIER = 3` (matching ClickHouse's default) — the summary tracks `k × 3` counters internally
- Merge support for parallel/distributed execution
- Binary serialization/deserialization for intermediate state transfer
- **`ApproxTopK`** — `AggregateUDFImpl` with `#[user_doc]` documentation, supporting any hashable scalar type
- **`ApproxTopKAccumulator`** — Accumulator with `update_batch`, `merge_batch`, `evaluate` (returns `List(Struct({value: T, count: UInt64}))` ordered by count descending), and `state` (serialized summary + data type)
- **9 unit tests** covering basic operations, eviction, alpha map, merge, serialization, accumulator update/evaluate, merge_batch, and a distributed merge simulation
### Registration (`datafusion/functions-aggregate/src/lib.rs`)
- Added module, `expr_fn` re-export, and registration in `all_default_aggregate_functions()`
### SQL logic tests (`datafusion/sqllogictest/test_files/approx_top_k.slt`)
- 8 tests: basic string/int, GROUP BY, k=1, NULL handling, error cases (invalid k), WHERE clause filtering
### Documentation
- `docs/source/user-guide/sql/aggregate_functions.md` — TOC entry + full documentation section
- `docs/source/user-guide/expressions.md` — quick-reference table entry
### Integration tests
- **Proto roundtrip** (`datafusion/proto/tests/cases/roundtrip_logical_plan.rs`) — `approx_top_k(col("a"), lit(3))`
- **DataFrame** (`datafusion/core/tests/dataframe/dataframe_functions.rs`) — `test_fn_approx_top_k` with `insta::assert_snapshot!`
## Are these changes tested?
Yes:
- 9 unit tests in `approx_top_k.rs` (algorithm correctness, serialization, accumulator lifecycle, distributed merge)
- 8 SQL logic tests in `approx_top_k.slt` (end-to-end SQL behavior including edge cases and errors)
- 1 DataFrame integration test with snapshot assertion
- 1 proto roundtrip test
## Are there any user-facing changes?
Yes — new `approx_top_k` aggregate function available in SQL and the DataFrame API:
```sql
-- Top 5 most frequent values (k defaults to 5)
SELECT approx_top_k(url_path) FROM access_logs;
-- Top 3 per group
SELECT region, approx_top_k(product_id, 3) FROM sales GROUP BY region;
```
```rust
// DataFrame API
use datafusion::functions_aggregate::approx_top_k::approx_top_k;
df.aggregate(vec![], vec![approx_top_k(col("url_path"), lit(5))])?;
```
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "martin-g (via GitHub)" <gi...@apache.org>.
martin-g commented on PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#issuecomment-4154573219
Unfortunately, after the squash+force push I will have to review the whole PR again ...
Please don't squash your commits in a PR!
It just makes incremental PR reviews hard/impossible.
The merge with squash them anyway, with separate commit message for all commits.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r3123517113
##########
datafusion/functions-aggregate/src/approx_top_k.rs:
##########
@@ -0,0 +1,1421 @@
+// 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, 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};
+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; 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) -> 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>,
Review Comment:
`Vec<u8>`: this is quite wasteful to store in the map.
we can store the hash inline in the map using the `HashTable` API instead.
`HashTable<u64, usize>`
We can remove the hash from the `counters`
and look up the `item` using the `usize` pointer using `counter_map.find`.
(Look up the pattern in other places like aggregation / join ...)
This saves many allocations / expensive rehashing / memory usage.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r3728960623
##########
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);
+
+ // If there are still more counters than requested_capacity (because
+ // target_capacity wasn't reached), partition out the top ones and
+ // fold the tail into the alpha map before serializing.
+ let k = self.requested_capacity;
+ if k > 0 && 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();
+ }
+ self.rebuild_counter_map();
+ }
+
+ let num_counters = self.counters.len();
+
+ // Pre-compute total size for a single allocation.
+ let counter_bytes: usize =
+ self.counters.iter().map(|c| 4 + c.item.len() + 16).sum();
+ let total = 16 + counter_bytes + 8 + self.alpha_map.len() * 8;
+ let mut bytes = Vec::with_capacity(total);
+
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(num_counters as u64).to_le_bytes());
+
+ for counter in &self.counters {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
+ for &alpha in &self.alpha_map {
+ bytes.extend_from_slice(&alpha.to_le_bytes());
+ }
+
+ bytes
+ }
+
+ /// Deserialize a summary from bytes.
+ fn from_bytes(bytes: &[u8]) -> Result<Self> {
+ if bytes.len() < 16 {
+ return exec_err!("Invalid Space-Saving summary bytes: too short");
+ }
+
+ let requested_capacity =
+ u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
+ let num_counters = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+
+ // Validate against reasonable upper bounds to prevent OOM from
+ // malformed state.
+ if requested_capacity == 0 {
+ return exec_err!("Invalid Space-Saving summary: requested_capacity is 0");
+ }
+ let max_capacity = APPROX_TOP_K_MAX_K * CAPACITY_MULTIPLIER;
+ if requested_capacity > max_capacity {
+ return exec_err!(
+ "Invalid Space-Saving summary: requested_capacity {requested_capacity} \
+ exceeds maximum {max_capacity}"
+ );
+ }
+ if num_counters > requested_capacity {
+ return exec_err!(
+ "Invalid Space-Saving summary: num_counters {num_counters} exceeds \
+ requested_capacity {requested_capacity}"
+ );
+ }
+ // Each counter needs at least 20 bytes (4 len + 0 item + 8 count + 8 error).
+ let max_possible = (bytes.len().saturating_sub(16)) / 20;
+ if num_counters > max_possible {
+ return exec_err!(
+ "Invalid Space-Saving summary: num_counters {num_counters} exceeds \
+ what fits in {} bytes",
+ bytes.len()
+ );
+ }
+
+ let mut counters = Vec::with_capacity(num_counters);
+ let mut counter_map = HashTable::with_capacity(num_counters);
+ let mut item_heap_bytes: usize = 0;
+ let mut offset: usize = 16;
+
+ for idx in 0..num_counters {
+ if offset.checked_add(4).is_none_or(|end| end > bytes.len()) {
+ return exec_err!(
+ "Invalid Space-Saving summary bytes: truncated item length"
+ );
+ }
+ let item_len =
+ u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
+ as usize;
+ offset += 4;
+
+ let needed = item_len.checked_add(16);
+ if needed
+ .is_none_or(|n| offset.checked_add(n).is_none_or(|end| end > bytes.len()))
+ {
+ return exec_err!(
+ "Invalid Space-Saving summary bytes: truncated counter"
+ );
+ }
+
+ let item = bytes[offset..offset + item_len].to_vec();
+ offset += item_len;
+
+ let count = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let error = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let hash = Self::hash_item(&item);
+ item_heap_bytes += item.capacity();
+ counter_map.insert_unique(hash, (hash, idx), |&(h, _)| h);
+ counters.push(Counter {
+ item,
+ hash,
+ count,
+ error,
+ });
+ }
+
+ if offset.checked_add(8).is_none_or(|end| end > bytes.len()) {
+ return exec_err!(
+ "Invalid Space-Saving summary bytes: missing alpha map size"
+ );
+ }
+ let alpha_map_size =
+ u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap()) as usize;
+ offset += 8;
+
+ // Validate alpha_map_size is a power of two (required for bitmask indexing).
+ if alpha_map_size == 0 || !alpha_map_size.is_power_of_two() {
+ return exec_err!(
+ "Invalid Space-Saving summary: alpha_map_size {alpha_map_size} \
+ is not a positive power of two"
+ );
+ }
+
+ let alpha_bytes = alpha_map_size
+ .checked_mul(8)
+ .and_then(|n| offset.checked_add(n));
+ if alpha_bytes.is_none_or(|end| end > bytes.len()) {
+ return exec_err!("Invalid Space-Saving summary bytes: truncated alpha map");
+ }
+
+ let mut alpha_map = Vec::with_capacity(alpha_map_size);
+ for _ in 0..alpha_map_size {
+ let alpha = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+ alpha_map.push(alpha);
+ }
+
+ let target_capacity = max(64, requested_capacity.saturating_mul(2));
+
+ Ok(Self {
+ counters,
+ counter_map,
+ alpha_map,
+ requested_capacity,
+ target_capacity,
+ item_heap_bytes,
+ })
+ }
+
+ /// Approximate size in bytes of this summary. O(1) thanks to
+ /// incremental `item_heap_bytes` tracking.
+ fn size(&self) -> usize {
+ size_of::<Self>()
+ + self.counters.capacity() * size_of::<Counter>()
+ + self.item_heap_bytes
+ // HashTable<(u64, usize)>: each bucket stores (hash, idx) + 1 control byte.
+ + self.counter_map.capacity()
+ * (size_of::<(u64, usize)>() + size_of::<u8>())
+ + self.alpha_map.capacity() * size_of::<u64>()
+ }
+}
+
+// ---------------------------------------------------------------------------
+// ApproxTopK UDAF struct
+// ---------------------------------------------------------------------------
+
+/// Approximate top-k UDAF using the Filtered Space-Saving algorithm.
+#[user_doc(
+ doc_section(label = "Approximate Functions"),
+ description = r#"Returns the approximate most frequent (top-k) values with their estimated counts, using the Filtered Space-Saving algorithm. The returned counts are upper-bound estimates; the true frequency lies in `[count - error, count]`. NULL values are skipped; an empty or all-NULL input returns an empty list `[]`. For float columns, -0.0 and +0.0 are treated as distinct values, and different NaN representations are tracked separately."#,
+ syntax_example = "approx_top_k(expression, k)",
+ sql_example = r#"```sql
+> SELECT approx_top_k(column_name, 3) FROM table_name;
++-----------------------------------------------------------------------------+
+| approx_top_k(column_name,Int64(3)) |
++-----------------------------------------------------------------------------+
+| [{value: foo, count: 3}, {value: bar, count: 2}, {value: baz, count: 1}] |
++-----------------------------------------------------------------------------+
+```"#,
+ standard_argument(name = "expression",),
+ argument(
+ name = "k",
+ description = "The number of top elements to return. Must be a literal integer between 1 and 10,000."
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ApproxTopK {
+ signature: Signature,
+}
+
+impl Default for ApproxTopK {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ApproxTopK {
+ pub fn new() -> Self {
+ // Supported value types for the first argument.
+ let value_types = &[
+ DataType::Utf8,
+ DataType::LargeUtf8,
+ DataType::Binary,
+ DataType::LargeBinary,
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ DataType::Float32,
+ DataType::Float64,
+ DataType::Date32,
+ DataType::Date64,
+ DataType::Timestamp(TimeUnit::Second, None),
+ DataType::Timestamp(TimeUnit::Millisecond, None),
+ DataType::Timestamp(TimeUnit::Microsecond, None),
+ DataType::Timestamp(TimeUnit::Nanosecond, None),
+ DataType::Timestamp(TimeUnit::Second, Some(TIMEZONE_WILDCARD.into())),
+ DataType::Timestamp(TimeUnit::Millisecond, Some(TIMEZONE_WILDCARD.into())),
+ DataType::Timestamp(TimeUnit::Microsecond, Some(TIMEZONE_WILDCARD.into())),
+ DataType::Timestamp(TimeUnit::Nanosecond, Some(TIMEZONE_WILDCARD.into())),
+ ];
+ // k must be a literal integer; accept any integer type for convenience.
+ let k_types = &[
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ ];
+
+ let mut variants = Vec::with_capacity(value_types.len() * k_types.len());
+ for vt in value_types {
+ for kt in k_types {
+ variants.push(TypeSignature::Exact(vec![vt.clone(), kt.clone()]));
+ }
+ }
+
+ Self {
+ signature: Signature::one_of(variants, Volatility::Immutable),
+ }
+ }
+}
+
+impl AggregateUDFImpl for ApproxTopK {
+ fn name(&self) -> &str {
+ "approx_top_k"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ let value_type = if !arg_types.is_empty() {
+ match &arg_types[0] {
+ // Large variants are narrowed: the summary stores in-memory
+ // byte slices, not large column offsets, so i32 offsets suffice.
+ DataType::LargeUtf8 => DataType::Utf8,
+ DataType::LargeBinary => DataType::Binary,
+ other => other.clone(),
+ }
+ } else {
+ DataType::Utf8
+ };
+
+ let struct_fields = Fields::from(vec![
+ Field::new("value", value_type, true),
+ Field::new("count", DataType::UInt64, false),
+ ]);
+ Ok(DataType::List(Arc::new(Field::new(
+ "item",
+ DataType::Struct(struct_fields),
+ true,
+ ))))
+ }
+
+ fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
+ Ok(vec![Arc::new(Field::new(
+ format_state_name(args.name, "summary"),
+ DataType::Binary,
+ true,
+ ))])
+ }
+
+ fn accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+ if args.exprs.len() < 2 {
+ return plan_err!("approx_top_k requires two arguments: column and k");
+ }
+
+ let k_expr = &args.exprs[1];
+ let k = k_expr
+ .downcast_ref::<datafusion_physical_expr::expressions::Literal>()
+ .and_then(|lit| match lit.value() {
+ // Guard against negative values before casting to usize to
+ // avoid wrapping (e.g. -1i32 as usize → u64::MAX).
+ // Zero is allowed through and caught by the bounds check below.
+ ScalarValue::Int8(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int16(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int32(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int64(Some(v)) if *v >= 0 => Some(*v as usize),
Review Comment:
Zero is rejected, but I kept literal conversion separate from the shared range validation so signed and unsigned literals follow the same 1..=10000 validation and error path. Tests cover zero, negative values, 10000, and 10001.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r2944482840
##########
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:
Can we use `foldhash` as well (which I just merged)?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#issuecomment-4192502400
> Unfortunately, after the squash+force push I will have to review the whole PR again ... Please don't squash your commits in a PR! It just makes incremental PR reviews hard/impossible. The merge with squash them anyway, with separate commit message for all commits.
Sorry about that, the squash slipped through accidentally!
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r3728937729
##########
datafusion/functions-aggregate/src/approx_top_k.rs:
##########
@@ -0,0 +1,1421 @@
+// 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, 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};
+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; 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) -> 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>,
Review Comment:
I kept the cached hash on each counter because eviction uses it to update the alpha bucket, and index rebuilding can then avoid rehashing item bytes. The HashTable<(u64, usize)> stores no key bytes, so keys are not duplicated or separately allocated.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r2944480541
##########
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;
Review Comment:
Should this be configurable?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#issuecomment-4297707365
@sesteves
starts to look good!
Can you also remove the testing submodule?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#issuecomment-5268411294
Thank you for the thorough review @timsaucer . I have addressed the feedback and aligned the implementation closer to the original paper rather than clickhouse variant. We have already been running this in our production environment since we need to provide top-k with constant memory utilization regardless of dataset size.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#issuecomment-4517389753
@Dandandan @martin-g when should I expect the next round of reviews on this? thanks!
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r3123517113
##########
datafusion/functions-aggregate/src/approx_top_k.rs:
##########
@@ -0,0 +1,1421 @@
+// 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, 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};
+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; 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) -> 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>,
Review Comment:
`Vec<u8>`: this is quite wasteful to store in the map.
we can store the hash inline in the map using the `HashTable` API instead.
`HashTable<u64, usize>`
and look up the `item` using the `usize` pointer using `counter_map.find`.
(Look up the pattern in other places like aggregation / join ...)
This saves many allocations / rehashing.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r3729705937
##########
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),
+ );
Review Comment:
Yes. The implementation now follows original Filtered Space-Saving: while `alpha + 1` is below the minimum tracked count, the item is rejected and alpha is incremented. Once it reaches the threshold, the minimum counter is replaced using alpha as the estimate and error. Counters are ranked by estimated count, then error. Tests cover filtering, admission, and same-bucket replacement.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "martin-g (via GitHub)" <gi...@apache.org>.
martin-g commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r2979897987
##########
datafusion/core/tests/dataframe/dataframe_functions.rs:
##########
@@ -409,6 +411,27 @@ async fn test_fn_approx_median() -> Result<()> {
Ok(())
}
+#[tokio::test]
+async fn test_fn_approx_top_k() -> Result<()> {
+ // Column b has values [1, 10, 10, 100] -- 10 appears twice, others once.
+ let expr = approx_top_k(vec![col("b"), lit(2)]);
+
+ let df = create_test_table().await?;
+ let batches = df.aggregate(vec![], vec![expr]).unwrap().collect().await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +-----------------------------------------------+
+ | approx_top_k(test.b,Int32(2)) |
+ +-----------------------------------------------+
+ | [{value: 10, count: 2}, {value: 1, count: 1}] |
Review Comment:
Is the second result (`{value: 1, count: 1}`) stable here ?
`value: 100` is Equal - both are `count=1, error=0`.
##########
datafusion/functions-aggregate/src/approx_top_k.rs:
##########
@@ -0,0 +1,1381 @@
+// 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.
Review Comment:
```suggestion
/// Get the top-k items sorted by (count - error) descending.
```
##########
datafusion/functions-aggregate/src/approx_top_k.rs:
##########
@@ -0,0 +1,1381 @@
+// 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 serialize(&mut self) -> Vec<u8> {
+ // Ensure counters are truncated and alpha map is up to date before
+ // serializing, in case to_bytes is called without a prior truncation.
+ self.truncate_if_needed(false);
+
+ let counters_to_write: Vec<_> = {
+ let mut sorted: Vec<_> = self.counters.iter().collect();
+ let return_size = std::cmp::min(sorted.len(), self.requested_capacity);
+ if return_size > 0 && return_size < sorted.len() {
+ // After select_nth_unstable_by, the top-k counters are in
+ // sorted[..return_size] but in arbitrary order. This is fine
+ // since deserialization doesn't depend on counter ordering.
+ sorted.select_nth_unstable_by(return_size - 1, |a, b| a.cmp_by_rank(b));
+ }
+ sorted.truncate(return_size);
+ sorted
+ };
+
+ let mut bytes = Vec::new();
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(counters_to_write.len() as u64).to_le_bytes());
+
+ for counter in counters_to_write {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
+ for &alpha in &self.alpha_map {
+ bytes.extend_from_slice(&alpha.to_le_bytes());
+ }
+
+ bytes
+ }
+
+ /// Deserialize a summary from bytes.
+ fn from_bytes(bytes: &[u8]) -> Result<Self> {
+ if bytes.len() < 16 {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: too short".to_string(),
+ ));
+ }
+
+ let requested_capacity =
+ u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
+ let num_counters = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+
+ let mut counters = Vec::with_capacity(num_counters);
+ let mut counter_map = HashMap::with_capacity(num_counters);
+ let mut offset = 16;
+
+ for idx in 0..num_counters {
+ if offset + 4 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated item length"
+ .to_string(),
+ ));
+ }
+ let item_len =
+ u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
+ as usize;
+ offset += 4;
+
+ if offset + item_len + 16 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated counter".to_string(),
+ ));
+ }
+
+ let item = bytes[offset..offset + item_len].to_vec();
+ offset += item_len;
+
+ let count = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let error = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let hash = Self::hash_item(&item);
+ counter_map.insert(item.clone(), idx);
+ counters.push(Counter {
+ item,
+ hash,
+ count,
+ error,
+ });
+ }
+
+ if offset + 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: missing alpha map size".to_string(),
+ ));
+ }
+ let alpha_map_size =
+ u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap()) as usize;
+ offset += 8;
+
+ if offset + alpha_map_size * 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated alpha map".to_string(),
+ ));
+ }
+
+ let mut alpha_map = Vec::with_capacity(alpha_map_size);
+ for _ in 0..alpha_map_size {
+ let alpha = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+ alpha_map.push(alpha);
+ }
+
+ let target_capacity = std::cmp::max(64, requested_capacity * 2);
+
+ Ok(Self {
+ counters,
+ counter_map,
+ alpha_map,
+ requested_capacity,
+ target_capacity,
+ })
+ }
+
+ /// Approximate size in bytes of this summary.
+ fn size(&self) -> usize {
+ size_of::<Self>()
+ + self
+ .counters
+ .iter()
+ .map(|c| size_of::<Counter>() + c.item.len())
+ .sum::<usize>()
Review Comment:
This is not the correct way to calculate the size of a collection.
For measuring the size of a Vec you need: `coll.capacity() * size_of::<CollItem>()` (as you do for `alpha_map`).
For map you'd need: `map.capacity() * (size_of::<Key>()) + size_of::<Value>())` (as you do for `counters_map`).
So, it should be something like:
```rust
fn size(&self) -> usize {
let item_bytes: usize = self.counters.iter().map(|c| c.item.capacity() * size_of::<u8>()).sum();
size_of::<Self>()
+ (self.counters.capacity() * size_of::<Counter>())
// The keys are cloned from `counters` so their heap allocation is measured separately.
+ (self.counter_map.capacity() * (size_of::<(Vec<u8>, usize)>() + size_of::<u8>()))
+ (self.alpha_map.capacity() * size_of::<u64>())
+ item_bytes * 2 // once for counters and once for counters_map's keys
}
```
##########
datafusion/functions-aggregate/src/approx_top_k.rs:
##########
@@ -0,0 +1,1381 @@
+// 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.
Review Comment:
https://github.com/apache/datafusion/pull/20968/changes#diff-ab75150c7bffcb234e6dbbcefcc0c6093cfb1cfc12a1c6572f64519c9bdb0134R171 actually does `capacity * 2`, so it is 60KB, not 30KB.
##########
datafusion/functions-aggregate/src/approx_top_k.rs:
##########
@@ -0,0 +1,1381 @@
+// 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 serialize(&mut self) -> Vec<u8> {
+ // Ensure counters are truncated and alpha map is up to date before
+ // serializing, in case to_bytes is called without a prior truncation.
+ self.truncate_if_needed(false);
+
+ let counters_to_write: Vec<_> = {
+ let mut sorted: Vec<_> = self.counters.iter().collect();
+ let return_size = std::cmp::min(sorted.len(), self.requested_capacity);
+ if return_size > 0 && return_size < sorted.len() {
+ // After select_nth_unstable_by, the top-k counters are in
+ // sorted[..return_size] but in arbitrary order. This is fine
+ // since deserialization doesn't depend on counter ordering.
+ sorted.select_nth_unstable_by(return_size - 1, |a, b| a.cmp_by_rank(b));
+ }
+ sorted.truncate(return_size);
+ sorted
+ };
+
+ let mut bytes = Vec::new();
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(counters_to_write.len() as u64).to_le_bytes());
+
+ for counter in counters_to_write {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
+ for &alpha in &self.alpha_map {
+ bytes.extend_from_slice(&alpha.to_le_bytes());
+ }
+
+ bytes
+ }
+
+ /// Deserialize a summary from bytes.
+ fn from_bytes(bytes: &[u8]) -> Result<Self> {
+ if bytes.len() < 16 {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: too short".to_string(),
+ ));
+ }
+
+ let requested_capacity =
+ u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
+ let num_counters = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+
+ let mut counters = Vec::with_capacity(num_counters);
+ let mut counter_map = HashMap::with_capacity(num_counters);
+ let mut offset = 16;
+
+ for idx in 0..num_counters {
+ if offset + 4 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated item length"
+ .to_string(),
+ ));
+ }
+ let item_len =
+ u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
+ as usize;
+ offset += 4;
+
+ if offset + item_len + 16 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated counter".to_string(),
+ ));
+ }
+
+ let item = bytes[offset..offset + item_len].to_vec();
+ offset += item_len;
+
+ let count = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let error = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let hash = Self::hash_item(&item);
+ counter_map.insert(item.clone(), idx);
+ counters.push(Counter {
+ item,
+ hash,
+ count,
+ error,
+ });
+ }
+
+ if offset + 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: missing alpha map size".to_string(),
+ ));
+ }
+ let alpha_map_size =
+ u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap()) as usize;
+ offset += 8;
+
+ if offset + alpha_map_size * 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated alpha map".to_string(),
+ ));
+ }
+
+ let mut alpha_map = Vec::with_capacity(alpha_map_size);
+ for _ in 0..alpha_map_size {
+ let alpha = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+ alpha_map.push(alpha);
+ }
+
+ let target_capacity = std::cmp::max(64, requested_capacity * 2);
+
+ Ok(Self {
+ counters,
+ counter_map,
+ alpha_map,
+ requested_capacity,
+ target_capacity,
+ })
+ }
+
+ /// Approximate size in bytes of this summary.
+ fn size(&self) -> usize {
+ size_of::<Self>()
+ + self
+ .counters
+ .iter()
+ .map(|c| size_of::<Counter>() + c.item.len())
+ .sum::<usize>()
+ // HashMap overhead: each bucket has a key-value pair plus control byte metadata.
+ + self.counter_map.capacity()
+ * (size_of::<(Vec<u8>, usize)>() + size_of::<u8>())
+ + self.alpha_map.capacity() * size_of::<u64>()
+ }
+}
+
+// ---------------------------------------------------------------------------
+// ApproxTopK UDAF struct
+// ---------------------------------------------------------------------------
+
+/// Approximate top-k UDAF using the Filtered Space-Saving algorithm.
+#[user_doc(
+ doc_section(label = "Approximate Functions"),
+ description = "Returns the approximate most frequent (top-k) values and their counts using the Filtered Space-Saving algorithm. Note: for float columns, -0.0 and +0.0 are treated as distinct values, and different NaN representations are tracked separately.",
+ syntax_example = "approx_top_k(expression, k)",
+ sql_example = r#"```sql
+> SELECT approx_top_k(column_name, 3) FROM table_name;
++-------------------------------------------+
+| approx_top_k(column_name, 3) |
++-------------------------------------------+
+| [{value: foo, count: 3}, {value: bar, count: 2}, {value: baz, count: 1}] |
++-------------------------------------------+
+```"#,
+ standard_argument(name = "expression",),
+ argument(
+ name = "k",
+ description = "The number of top elements to return. Must be a literal integer between 1 and 10,000."
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ApproxTopK {
+ signature: Signature,
+}
+
+impl Default for ApproxTopK {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ApproxTopK {
+ pub fn new() -> Self {
+ // Supported value types for the first argument.
+ let value_types = &[
+ DataType::Utf8,
+ DataType::LargeUtf8,
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ DataType::Float32,
+ DataType::Float64,
+ ];
+ // k must be a literal integer; accept any integer type for convenience.
+ let k_types = &[
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ ];
+
+ let mut variants = Vec::with_capacity(value_types.len() * k_types.len());
+ for vt in value_types {
+ for kt in k_types {
+ variants.push(TypeSignature::Exact(vec![vt.clone(), kt.clone()]));
+ }
+ }
+
+ Self {
+ signature: Signature::one_of(variants, Volatility::Immutable),
+ }
+ }
+}
+
+impl AggregateUDFImpl for ApproxTopK {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "approx_top_k"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ let value_type = if !arg_types.is_empty() {
+ match &arg_types[0] {
+ // LargeUtf8 is narrowed to Utf8: the summary stores in-memory
+ // byte slices, not large column offsets, so i32 offsets suffice.
+ DataType::Utf8 | DataType::LargeUtf8 => DataType::Utf8,
+ DataType::Int8 => DataType::Int8,
+ DataType::Int16 => DataType::Int16,
+ DataType::Int32 => DataType::Int32,
+ DataType::Int64 => DataType::Int64,
+ DataType::UInt8 => DataType::UInt8,
+ DataType::UInt16 => DataType::UInt16,
+ DataType::UInt32 => DataType::UInt32,
+ DataType::UInt64 => DataType::UInt64,
+ DataType::Float32 => DataType::Float32,
+ DataType::Float64 => DataType::Float64,
+ _ => DataType::Utf8,
+ }
+ } else {
+ DataType::Utf8
+ };
+
+ let struct_fields = Fields::from(vec![
+ Field::new("value", value_type, true),
+ Field::new("count", DataType::UInt64, false),
+ ]);
+ Ok(DataType::List(Arc::new(Field::new(
+ "item",
+ DataType::Struct(struct_fields),
+ true,
+ ))))
+ }
+
+ fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
+ Ok(vec![
+ Arc::new(Field::new(
+ format_state_name(args.name, "summary"),
+ DataType::Binary,
+ true,
+ )),
+ Arc::new(Field::new(
+ format_state_name(args.name, "k"),
+ DataType::UInt64,
+ true,
+ )),
+ Arc::new(Field::new(
+ format_state_name(args.name, "data_type"),
+ DataType::Utf8,
+ true,
+ )),
+ ])
+ }
+
+ fn accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+ if args.exprs.len() < 2 {
+ return Err(datafusion_common::DataFusionError::Plan(
+ "approx_top_k requires two arguments: column and k".to_string(),
+ ));
+ }
+
+ let k_expr = &args.exprs[1];
+ let k = k_expr
+ .as_any()
+ .downcast_ref::<datafusion_physical_expr::expressions::Literal>()
+ .and_then(|lit| match lit.value() {
+ // Guard against negative values before casting to usize to
+ // avoid wrapping (e.g. -1i32 as usize → u64::MAX).
+ // Zero is allowed through and caught by the bounds check below.
+ ScalarValue::Int8(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int16(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int32(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int64(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::UInt8(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt16(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt32(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt64(Some(v)) => Some(*v as usize),
+ _ => None,
+ })
+ .ok_or_else(|| {
+ datafusion_common::DataFusionError::Plan(
+ "approx_top_k requires k to be a positive literal integer"
+ .to_string(),
+ )
+ })?;
+
+ if k == 0 || k > APPROX_TOP_K_MAX_K {
+ return Err(datafusion_common::DataFusionError::Plan(format!(
+ "approx_top_k requires k to be between 1 and {APPROX_TOP_K_MAX_K}, got {k}"
+ )));
+ }
+
+ let data_type = args.expr_fields[0].data_type().clone();
+ Ok(Box::new(ApproxTopKAccumulator::new_with_data_type(
+ k,
+ Some(data_type),
+ )))
+ }
+
+ fn documentation(&self) -> Option<&Documentation> {
+ self.doc()
+ }
+}
+
+// ---------------------------------------------------------------------------
+// Accumulator
+// ---------------------------------------------------------------------------
+
+/// Accumulator for `approx_top_k` using the Filtered Space-Saving algorithm.
+#[derive(Debug)]
+struct ApproxTopKAccumulator {
+ summary: SpaceSavingSummary,
+ k: usize,
+ /// The data type of the input column. `None` until the first batch is seen.
+ input_data_type: Option<DataType>,
+}
+
+impl ApproxTopKAccumulator {
+ fn new_with_data_type(k: usize, input_data_type: Option<DataType>) -> Self {
+ let capacity = k * CAPACITY_MULTIPLIER;
+ Self {
+ summary: SpaceSavingSummary::new(capacity),
+ k,
+ input_data_type,
+ }
+ }
+
+ /// Build the value array for the result based on the input data type.
+ fn build_value_array(&self, top_items: &[(&[u8], u64, u64)]) -> Result<ArrayRef> {
+ match self.input_data_type.as_ref().unwrap_or(&DataType::Utf8) {
+ DataType::Utf8 | DataType::LargeUtf8 => {
+ let values: Vec<Option<String>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| String::from_utf8(bytes.to_vec()).ok())
+ .collect();
+ Ok(Arc::new(StringArray::from(values)) as ArrayRef)
+ }
+ DataType::Int8 => {
+ let values: Vec<Option<i8>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 1]>::try_from(*bytes).ok().map(i8::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(Int8Array::from(values)) as ArrayRef)
+ }
+ DataType::Int16 => {
+ let values: Vec<Option<i16>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 2]>::try_from(*bytes).ok().map(i16::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(Int16Array::from(values)) as ArrayRef)
+ }
+ DataType::Int32 => {
+ let values: Vec<Option<i32>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 4]>::try_from(*bytes).ok().map(i32::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(Int32Array::from(values)) as ArrayRef)
+ }
+ DataType::Int64 => {
+ let values: Vec<Option<i64>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 8]>::try_from(*bytes).ok().map(i64::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(Int64Array::from(values)) as ArrayRef)
+ }
+ DataType::UInt8 => {
+ let values: Vec<Option<u8>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 1]>::try_from(*bytes).ok().map(u8::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(UInt8Array::from(values)) as ArrayRef)
+ }
+ DataType::UInt16 => {
+ let values: Vec<Option<u16>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 2]>::try_from(*bytes).ok().map(u16::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(UInt16Array::from(values)) as ArrayRef)
+ }
+ DataType::UInt32 => {
+ let values: Vec<Option<u32>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 4]>::try_from(*bytes).ok().map(u32::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(UInt32Array::from(values)) as ArrayRef)
+ }
+ DataType::UInt64 => {
+ let values: Vec<Option<u64>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 8]>::try_from(*bytes).ok().map(u64::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(UInt64Array::from(values)) as ArrayRef)
+ }
+ DataType::Float32 => {
+ let values: Vec<Option<f32>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 4]>::try_from(*bytes).ok().map(f32::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(Float32Array::from(values)) as ArrayRef)
+ }
+ DataType::Float64 => {
+ let values: Vec<Option<f64>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 8]>::try_from(*bytes).ok().map(f64::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(Float64Array::from(values)) as ArrayRef)
+ }
+ _ => {
+ let values: Vec<Option<String>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| String::from_utf8(bytes.to_vec()).ok())
+ .collect();
+ Ok(Arc::new(StringArray::from(values)) as ArrayRef)
+ }
+ }
+ }
+
+ /// Get the output data type for the value field.
+ /// Falls back to `Utf8` if the input data type was never captured (empty input).
+ fn output_value_data_type(&self) -> DataType {
+ match self.input_data_type.as_ref().unwrap_or(&DataType::Utf8) {
+ DataType::Utf8 | DataType::LargeUtf8 => DataType::Utf8,
+ DataType::Int8 => DataType::Int8,
+ DataType::Int16 => DataType::Int16,
+ DataType::Int32 => DataType::Int32,
+ DataType::Int64 => DataType::Int64,
+ DataType::UInt8 => DataType::UInt8,
+ DataType::UInt16 => DataType::UInt16,
+ DataType::UInt32 => DataType::UInt32,
+ DataType::UInt64 => DataType::UInt64,
+ DataType::Float32 => DataType::Float32,
+ DataType::Float64 => DataType::Float64,
+ _ => DataType::Utf8,
+ }
+ }
+
+ fn data_type_to_string(dt: &DataType) -> String {
+ match dt {
+ DataType::Utf8 => "utf8".to_string(),
Review Comment:
Any reason not to use DataType's Display impl + DataType Parser ?
https://github.com/apache/arrow-rs/blob/70445c54ff510e1b9f8d6e6782b92a1ce45fbb28/arrow-schema/src/datatype_parse.rs
##########
datafusion/functions-aggregate/src/approx_top_k.rs:
##########
@@ -0,0 +1,1381 @@
+// 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 serialize(&mut self) -> Vec<u8> {
+ // Ensure counters are truncated and alpha map is up to date before
+ // serializing, in case to_bytes is called without a prior truncation.
+ self.truncate_if_needed(false);
+
+ let counters_to_write: Vec<_> = {
+ let mut sorted: Vec<_> = self.counters.iter().collect();
+ let return_size = std::cmp::min(sorted.len(), self.requested_capacity);
+ if return_size > 0 && return_size < sorted.len() {
+ // After select_nth_unstable_by, the top-k counters are in
+ // sorted[..return_size] but in arbitrary order. This is fine
+ // since deserialization doesn't depend on counter ordering.
+ sorted.select_nth_unstable_by(return_size - 1, |a, b| a.cmp_by_rank(b));
+ }
+ sorted.truncate(return_size);
+ sorted
+ };
+
+ let mut bytes = Vec::new();
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(counters_to_write.len() as u64).to_le_bytes());
+
+ for counter in counters_to_write {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
+ for &alpha in &self.alpha_map {
+ bytes.extend_from_slice(&alpha.to_le_bytes());
+ }
+
+ bytes
+ }
+
+ /// Deserialize a summary from bytes.
+ fn from_bytes(bytes: &[u8]) -> Result<Self> {
+ if bytes.len() < 16 {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: too short".to_string(),
+ ));
+ }
+
+ let requested_capacity =
+ u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
+ let num_counters = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+
+ let mut counters = Vec::with_capacity(num_counters);
+ let mut counter_map = HashMap::with_capacity(num_counters);
+ let mut offset = 16;
+
+ for idx in 0..num_counters {
+ if offset + 4 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated item length"
+ .to_string(),
+ ));
+ }
+ let item_len =
+ u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
+ as usize;
+ offset += 4;
+
+ if offset + item_len + 16 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated counter".to_string(),
+ ));
+ }
+
+ let item = bytes[offset..offset + item_len].to_vec();
+ offset += item_len;
+
+ let count = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let error = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let hash = Self::hash_item(&item);
+ counter_map.insert(item.clone(), idx);
+ counters.push(Counter {
+ item,
+ hash,
+ count,
+ error,
+ });
+ }
+
+ if offset + 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: missing alpha map size".to_string(),
+ ));
+ }
+ let alpha_map_size =
+ u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap()) as usize;
+ offset += 8;
+
+ if offset + alpha_map_size * 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated alpha map".to_string(),
+ ));
+ }
+
+ let mut alpha_map = Vec::with_capacity(alpha_map_size);
+ for _ in 0..alpha_map_size {
+ let alpha = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+ alpha_map.push(alpha);
+ }
+
+ let target_capacity = std::cmp::max(64, requested_capacity * 2);
+
+ Ok(Self {
+ counters,
+ counter_map,
+ alpha_map,
+ requested_capacity,
+ target_capacity,
+ })
+ }
+
+ /// Approximate size in bytes of this summary.
+ fn size(&self) -> usize {
+ size_of::<Self>()
+ + self
+ .counters
+ .iter()
+ .map(|c| size_of::<Counter>() + c.item.len())
+ .sum::<usize>()
+ // HashMap overhead: each bucket has a key-value pair plus control byte metadata.
+ + self.counter_map.capacity()
+ * (size_of::<(Vec<u8>, usize)>() + size_of::<u8>())
+ + self.alpha_map.capacity() * size_of::<u64>()
+ }
+}
+
+// ---------------------------------------------------------------------------
+// ApproxTopK UDAF struct
+// ---------------------------------------------------------------------------
+
+/// Approximate top-k UDAF using the Filtered Space-Saving algorithm.
+#[user_doc(
+ doc_section(label = "Approximate Functions"),
+ description = "Returns the approximate most frequent (top-k) values and their counts using the Filtered Space-Saving algorithm. Note: for float columns, -0.0 and +0.0 are treated as distinct values, and different NaN representations are tracked separately.",
+ syntax_example = "approx_top_k(expression, k)",
+ sql_example = r#"```sql
+> SELECT approx_top_k(column_name, 3) FROM table_name;
++-------------------------------------------+
+| approx_top_k(column_name, 3) |
++-------------------------------------------+
+| [{value: foo, count: 3}, {value: bar, count: 2}, {value: baz, count: 1}] |
++-------------------------------------------+
+```"#,
+ standard_argument(name = "expression",),
+ argument(
+ name = "k",
+ description = "The number of top elements to return. Must be a literal integer between 1 and 10,000."
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ApproxTopK {
+ signature: Signature,
+}
+
+impl Default for ApproxTopK {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ApproxTopK {
+ pub fn new() -> Self {
+ // Supported value types for the first argument.
+ let value_types = &[
+ DataType::Utf8,
+ DataType::LargeUtf8,
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ DataType::Float32,
+ DataType::Float64,
+ ];
+ // k must be a literal integer; accept any integer type for convenience.
+ let k_types = &[
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ ];
+
+ let mut variants = Vec::with_capacity(value_types.len() * k_types.len());
+ for vt in value_types {
+ for kt in k_types {
+ variants.push(TypeSignature::Exact(vec![vt.clone(), kt.clone()]));
+ }
+ }
+
+ Self {
+ signature: Signature::one_of(variants, Volatility::Immutable),
+ }
+ }
+}
+
+impl AggregateUDFImpl for ApproxTopK {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "approx_top_k"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ let value_type = if !arg_types.is_empty() {
+ match &arg_types[0] {
+ // LargeUtf8 is narrowed to Utf8: the summary stores in-memory
+ // byte slices, not large column offsets, so i32 offsets suffice.
+ DataType::Utf8 | DataType::LargeUtf8 => DataType::Utf8,
+ DataType::Int8 => DataType::Int8,
+ DataType::Int16 => DataType::Int16,
+ DataType::Int32 => DataType::Int32,
+ DataType::Int64 => DataType::Int64,
+ DataType::UInt8 => DataType::UInt8,
+ DataType::UInt16 => DataType::UInt16,
+ DataType::UInt32 => DataType::UInt32,
+ DataType::UInt64 => DataType::UInt64,
+ DataType::Float32 => DataType::Float32,
+ DataType::Float64 => DataType::Float64,
+ _ => DataType::Utf8,
+ }
+ } else {
+ DataType::Utf8
+ };
+
+ let struct_fields = Fields::from(vec![
+ Field::new("value", value_type, true),
+ Field::new("count", DataType::UInt64, false),
+ ]);
+ Ok(DataType::List(Arc::new(Field::new(
+ "item",
+ DataType::Struct(struct_fields),
+ true,
+ ))))
+ }
+
+ fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
+ Ok(vec![
+ Arc::new(Field::new(
+ format_state_name(args.name, "summary"),
+ DataType::Binary,
+ true,
+ )),
+ Arc::new(Field::new(
+ format_state_name(args.name, "k"),
+ DataType::UInt64,
+ true,
+ )),
+ Arc::new(Field::new(
+ format_state_name(args.name, "data_type"),
+ DataType::Utf8,
+ true,
+ )),
+ ])
+ }
+
+ fn accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+ if args.exprs.len() < 2 {
+ return Err(datafusion_common::DataFusionError::Plan(
+ "approx_top_k requires two arguments: column and k".to_string(),
+ ));
+ }
+
+ let k_expr = &args.exprs[1];
+ let k = k_expr
+ .as_any()
+ .downcast_ref::<datafusion_physical_expr::expressions::Literal>()
Review Comment:
What if the expression is a cast, e.g. `CAST("2" AS INT)` ? Would it still work ?
##########
datafusion/functions-aggregate/src/approx_top_k.rs:
##########
@@ -0,0 +1,1381 @@
+// 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 serialize(&mut self) -> Vec<u8> {
+ // Ensure counters are truncated and alpha map is up to date before
+ // serializing, in case to_bytes is called without a prior truncation.
+ self.truncate_if_needed(false);
+
+ let counters_to_write: Vec<_> = {
+ let mut sorted: Vec<_> = self.counters.iter().collect();
+ let return_size = std::cmp::min(sorted.len(), self.requested_capacity);
+ if return_size > 0 && return_size < sorted.len() {
+ // After select_nth_unstable_by, the top-k counters are in
+ // sorted[..return_size] but in arbitrary order. This is fine
+ // since deserialization doesn't depend on counter ordering.
+ sorted.select_nth_unstable_by(return_size - 1, |a, b| a.cmp_by_rank(b));
+ }
+ sorted.truncate(return_size);
+ sorted
+ };
+
+ let mut bytes = Vec::new();
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(counters_to_write.len() as u64).to_le_bytes());
+
+ for counter in counters_to_write {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
+ for &alpha in &self.alpha_map {
+ bytes.extend_from_slice(&alpha.to_le_bytes());
+ }
+
+ bytes
+ }
+
+ /// Deserialize a summary from bytes.
+ fn from_bytes(bytes: &[u8]) -> Result<Self> {
+ if bytes.len() < 16 {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: too short".to_string(),
+ ));
+ }
+
+ let requested_capacity =
+ u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
+ let num_counters = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+
+ let mut counters = Vec::with_capacity(num_counters);
+ let mut counter_map = HashMap::with_capacity(num_counters);
+ let mut offset = 16;
+
+ for idx in 0..num_counters {
+ if offset + 4 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated item length"
+ .to_string(),
+ ));
+ }
+ let item_len =
+ u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
+ as usize;
+ offset += 4;
+
+ if offset + item_len + 16 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated counter".to_string(),
+ ));
+ }
+
+ let item = bytes[offset..offset + item_len].to_vec();
+ offset += item_len;
+
+ let count = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let error = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let hash = Self::hash_item(&item);
+ counter_map.insert(item.clone(), idx);
+ counters.push(Counter {
+ item,
+ hash,
+ count,
+ error,
+ });
+ }
+
+ if offset + 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: missing alpha map size".to_string(),
+ ));
+ }
+ let alpha_map_size =
+ u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap()) as usize;
+ offset += 8;
+
+ if offset + alpha_map_size * 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated alpha map".to_string(),
+ ));
+ }
+
+ let mut alpha_map = Vec::with_capacity(alpha_map_size);
+ for _ in 0..alpha_map_size {
+ let alpha = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+ alpha_map.push(alpha);
+ }
+
+ let target_capacity = std::cmp::max(64, requested_capacity * 2);
+
+ Ok(Self {
+ counters,
+ counter_map,
+ alpha_map,
+ requested_capacity,
+ target_capacity,
+ })
+ }
+
+ /// Approximate size in bytes of this summary.
+ fn size(&self) -> usize {
+ size_of::<Self>()
+ + self
+ .counters
+ .iter()
+ .map(|c| size_of::<Counter>() + c.item.len())
+ .sum::<usize>()
+ // HashMap overhead: each bucket has a key-value pair plus control byte metadata.
+ + self.counter_map.capacity()
+ * (size_of::<(Vec<u8>, usize)>() + size_of::<u8>())
+ + self.alpha_map.capacity() * size_of::<u64>()
+ }
+}
+
+// ---------------------------------------------------------------------------
+// ApproxTopK UDAF struct
+// ---------------------------------------------------------------------------
+
+/// Approximate top-k UDAF using the Filtered Space-Saving algorithm.
+#[user_doc(
+ doc_section(label = "Approximate Functions"),
+ description = "Returns the approximate most frequent (top-k) values and their counts using the Filtered Space-Saving algorithm. Note: for float columns, -0.0 and +0.0 are treated as distinct values, and different NaN representations are tracked separately.",
+ syntax_example = "approx_top_k(expression, k)",
+ sql_example = r#"```sql
+> SELECT approx_top_k(column_name, 3) FROM table_name;
++-------------------------------------------+
+| approx_top_k(column_name, 3) |
++-------------------------------------------+
+| [{value: foo, count: 3}, {value: bar, count: 2}, {value: baz, count: 1}] |
++-------------------------------------------+
+```"#,
+ standard_argument(name = "expression",),
+ argument(
+ name = "k",
+ description = "The number of top elements to return. Must be a literal integer between 1 and 10,000."
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ApproxTopK {
+ signature: Signature,
+}
+
+impl Default for ApproxTopK {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ApproxTopK {
+ pub fn new() -> Self {
+ // Supported value types for the first argument.
+ let value_types = &[
+ DataType::Utf8,
+ DataType::LargeUtf8,
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ DataType::Float32,
+ DataType::Float64,
+ ];
+ // k must be a literal integer; accept any integer type for convenience.
+ let k_types = &[
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ ];
+
+ let mut variants = Vec::with_capacity(value_types.len() * k_types.len());
+ for vt in value_types {
+ for kt in k_types {
+ variants.push(TypeSignature::Exact(vec![vt.clone(), kt.clone()]));
+ }
+ }
+
+ Self {
+ signature: Signature::one_of(variants, Volatility::Immutable),
+ }
+ }
+}
+
+impl AggregateUDFImpl for ApproxTopK {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "approx_top_k"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ let value_type = if !arg_types.is_empty() {
+ match &arg_types[0] {
+ // LargeUtf8 is narrowed to Utf8: the summary stores in-memory
+ // byte slices, not large column offsets, so i32 offsets suffice.
+ DataType::Utf8 | DataType::LargeUtf8 => DataType::Utf8,
+ DataType::Int8 => DataType::Int8,
+ DataType::Int16 => DataType::Int16,
+ DataType::Int32 => DataType::Int32,
+ DataType::Int64 => DataType::Int64,
+ DataType::UInt8 => DataType::UInt8,
+ DataType::UInt16 => DataType::UInt16,
+ DataType::UInt32 => DataType::UInt32,
+ DataType::UInt64 => DataType::UInt64,
+ DataType::Float32 => DataType::Float32,
+ DataType::Float64 => DataType::Float64,
+ _ => DataType::Utf8,
+ }
+ } else {
+ DataType::Utf8
+ };
+
+ let struct_fields = Fields::from(vec![
+ Field::new("value", value_type, true),
+ Field::new("count", DataType::UInt64, false),
+ ]);
+ Ok(DataType::List(Arc::new(Field::new(
+ "item",
+ DataType::Struct(struct_fields),
+ true,
+ ))))
+ }
+
+ fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
+ Ok(vec![
+ Arc::new(Field::new(
+ format_state_name(args.name, "summary"),
+ DataType::Binary,
+ true,
+ )),
+ Arc::new(Field::new(
+ format_state_name(args.name, "k"),
+ DataType::UInt64,
+ true,
+ )),
+ Arc::new(Field::new(
+ format_state_name(args.name, "data_type"),
+ DataType::Utf8,
+ true,
+ )),
+ ])
+ }
+
+ fn accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+ if args.exprs.len() < 2 {
+ return Err(datafusion_common::DataFusionError::Plan(
+ "approx_top_k requires two arguments: column and k".to_string(),
+ ));
+ }
+
+ let k_expr = &args.exprs[1];
+ let k = k_expr
+ .as_any()
+ .downcast_ref::<datafusion_physical_expr::expressions::Literal>()
+ .and_then(|lit| match lit.value() {
+ // Guard against negative values before casting to usize to
+ // avoid wrapping (e.g. -1i32 as usize → u64::MAX).
+ // Zero is allowed through and caught by the bounds check below.
+ ScalarValue::Int8(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int16(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int32(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int64(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::UInt8(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt16(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt32(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt64(Some(v)) => Some(*v as usize),
+ _ => None,
+ })
+ .ok_or_else(|| {
+ datafusion_common::DataFusionError::Plan(
+ "approx_top_k requires k to be a positive literal integer"
+ .to_string(),
+ )
+ })?;
+
+ if k == 0 || k > APPROX_TOP_K_MAX_K {
+ return Err(datafusion_common::DataFusionError::Plan(format!(
+ "approx_top_k requires k to be between 1 and {APPROX_TOP_K_MAX_K}, got {k}"
+ )));
+ }
+
+ let data_type = args.expr_fields[0].data_type().clone();
+ Ok(Box::new(ApproxTopKAccumulator::new_with_data_type(
+ k,
+ Some(data_type),
+ )))
+ }
+
+ fn documentation(&self) -> Option<&Documentation> {
+ self.doc()
+ }
+}
+
+// ---------------------------------------------------------------------------
+// Accumulator
+// ---------------------------------------------------------------------------
+
+/// Accumulator for `approx_top_k` using the Filtered Space-Saving algorithm.
+#[derive(Debug)]
+struct ApproxTopKAccumulator {
+ summary: SpaceSavingSummary,
+ k: usize,
+ /// The data type of the input column. `None` until the first batch is seen.
+ input_data_type: Option<DataType>,
+}
+
+impl ApproxTopKAccumulator {
+ fn new_with_data_type(k: usize, input_data_type: Option<DataType>) -> Self {
Review Comment:
`input_data_type` is None only in the tests.
IMO you could drop `Option<>` and simplify the related logic a bit.
##########
datafusion/functions-aggregate/src/approx_top_k.rs:
##########
@@ -0,0 +1,1381 @@
+// 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 serialize(&mut self) -> Vec<u8> {
+ // Ensure counters are truncated and alpha map is up to date before
+ // serializing, in case to_bytes is called without a prior truncation.
+ self.truncate_if_needed(false);
+
+ let counters_to_write: Vec<_> = {
+ let mut sorted: Vec<_> = self.counters.iter().collect();
+ let return_size = std::cmp::min(sorted.len(), self.requested_capacity);
+ if return_size > 0 && return_size < sorted.len() {
+ // After select_nth_unstable_by, the top-k counters are in
+ // sorted[..return_size] but in arbitrary order. This is fine
+ // since deserialization doesn't depend on counter ordering.
+ sorted.select_nth_unstable_by(return_size - 1, |a, b| a.cmp_by_rank(b));
+ }
+ sorted.truncate(return_size);
+ sorted
+ };
+
+ let mut bytes = Vec::new();
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(counters_to_write.len() as u64).to_le_bytes());
+
+ for counter in counters_to_write {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
+ for &alpha in &self.alpha_map {
+ bytes.extend_from_slice(&alpha.to_le_bytes());
+ }
+
+ bytes
+ }
+
+ /// Deserialize a summary from bytes.
+ fn from_bytes(bytes: &[u8]) -> Result<Self> {
+ if bytes.len() < 16 {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: too short".to_string(),
+ ));
+ }
+
+ let requested_capacity =
+ u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
+ let num_counters = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+
+ let mut counters = Vec::with_capacity(num_counters);
+ let mut counter_map = HashMap::with_capacity(num_counters);
+ let mut offset = 16;
+
+ for idx in 0..num_counters {
+ if offset + 4 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated item length"
+ .to_string(),
+ ));
+ }
+ let item_len =
+ u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
+ as usize;
+ offset += 4;
+
+ if offset + item_len + 16 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated counter".to_string(),
+ ));
+ }
+
+ let item = bytes[offset..offset + item_len].to_vec();
+ offset += item_len;
+
+ let count = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let error = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let hash = Self::hash_item(&item);
+ counter_map.insert(item.clone(), idx);
+ counters.push(Counter {
+ item,
+ hash,
+ count,
+ error,
+ });
+ }
+
+ if offset + 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: missing alpha map size".to_string(),
+ ));
+ }
+ let alpha_map_size =
+ u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap()) as usize;
+ offset += 8;
+
+ if offset + alpha_map_size * 8 > bytes.len() {
Review Comment:
Could `alpha_map_size * 8` lead to an overflow here ?
Maybe use a saturating_mul() ?!
##########
datafusion/functions-aggregate/src/approx_top_k.rs:
##########
@@ -0,0 +1,1381 @@
+// 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 serialize(&mut self) -> Vec<u8> {
+ // Ensure counters are truncated and alpha map is up to date before
+ // serializing, in case to_bytes is called without a prior truncation.
+ self.truncate_if_needed(false);
+
+ let counters_to_write: Vec<_> = {
+ let mut sorted: Vec<_> = self.counters.iter().collect();
+ let return_size = std::cmp::min(sorted.len(), self.requested_capacity);
+ if return_size > 0 && return_size < sorted.len() {
+ // After select_nth_unstable_by, the top-k counters are in
+ // sorted[..return_size] but in arbitrary order. This is fine
+ // since deserialization doesn't depend on counter ordering.
+ sorted.select_nth_unstable_by(return_size - 1, |a, b| a.cmp_by_rank(b));
+ }
+ sorted.truncate(return_size);
+ sorted
+ };
+
+ let mut bytes = Vec::new();
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(counters_to_write.len() as u64).to_le_bytes());
+
+ for counter in counters_to_write {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
+ for &alpha in &self.alpha_map {
+ bytes.extend_from_slice(&alpha.to_le_bytes());
+ }
+
+ bytes
+ }
+
+ /// Deserialize a summary from bytes.
+ fn from_bytes(bytes: &[u8]) -> Result<Self> {
+ if bytes.len() < 16 {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: too short".to_string(),
+ ));
+ }
+
+ let requested_capacity =
+ u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
+ let num_counters = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+
+ let mut counters = Vec::with_capacity(num_counters);
+ let mut counter_map = HashMap::with_capacity(num_counters);
+ let mut offset = 16;
+
+ for idx in 0..num_counters {
+ if offset + 4 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated item length"
+ .to_string(),
+ ));
+ }
+ let item_len =
+ u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
+ as usize;
+ offset += 4;
+
+ if offset + item_len + 16 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated counter".to_string(),
+ ));
+ }
+
+ let item = bytes[offset..offset + item_len].to_vec();
+ offset += item_len;
+
+ let count = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let error = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let hash = Self::hash_item(&item);
+ counter_map.insert(item.clone(), idx);
+ counters.push(Counter {
+ item,
+ hash,
+ count,
+ error,
+ });
+ }
+
+ if offset + 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: missing alpha map size".to_string(),
+ ));
+ }
+ let alpha_map_size =
+ u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap()) as usize;
+ offset += 8;
+
+ if offset + alpha_map_size * 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated alpha map".to_string(),
+ ));
+ }
+
+ let mut alpha_map = Vec::with_capacity(alpha_map_size);
+ for _ in 0..alpha_map_size {
+ let alpha = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+ alpha_map.push(alpha);
+ }
+
+ let target_capacity = std::cmp::max(64, requested_capacity * 2);
+
+ Ok(Self {
+ counters,
+ counter_map,
+ alpha_map,
+ requested_capacity,
+ target_capacity,
+ })
+ }
+
+ /// Approximate size in bytes of this summary.
+ fn size(&self) -> usize {
+ size_of::<Self>()
+ + self
+ .counters
+ .iter()
+ .map(|c| size_of::<Counter>() + c.item.len())
+ .sum::<usize>()
+ // HashMap overhead: each bucket has a key-value pair plus control byte metadata.
+ + self.counter_map.capacity()
+ * (size_of::<(Vec<u8>, usize)>() + size_of::<u8>())
+ + self.alpha_map.capacity() * size_of::<u64>()
+ }
+}
+
+// ---------------------------------------------------------------------------
+// ApproxTopK UDAF struct
+// ---------------------------------------------------------------------------
+
+/// Approximate top-k UDAF using the Filtered Space-Saving algorithm.
+#[user_doc(
+ doc_section(label = "Approximate Functions"),
+ description = "Returns the approximate most frequent (top-k) values and their counts using the Filtered Space-Saving algorithm. Note: for float columns, -0.0 and +0.0 are treated as distinct values, and different NaN representations are tracked separately.",
+ syntax_example = "approx_top_k(expression, k)",
+ sql_example = r#"```sql
+> SELECT approx_top_k(column_name, 3) FROM table_name;
++-------------------------------------------+
+| approx_top_k(column_name, 3) |
++-------------------------------------------+
+| [{value: foo, count: 3}, {value: bar, count: 2}, {value: baz, count: 1}] |
++-------------------------------------------+
+```"#,
+ standard_argument(name = "expression",),
+ argument(
+ name = "k",
+ description = "The number of top elements to return. Must be a literal integer between 1 and 10,000."
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ApproxTopK {
+ signature: Signature,
+}
+
+impl Default for ApproxTopK {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ApproxTopK {
+ pub fn new() -> Self {
+ // Supported value types for the first argument.
+ let value_types = &[
+ DataType::Utf8,
+ DataType::LargeUtf8,
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ DataType::Float32,
+ DataType::Float64,
Review Comment:
Are there plans to add more supported types ?
Like dates, times, timestamps, binaryview, ... ?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r3713613691
##########
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
+ }
+ }
Review Comment:
It feels a bit more idiomatic to `impl Ord for Counter` instead.
##########
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.",
Review Comment:
nit: name of the algorithm probably isn't super important for user facing documentation. It is very useful for developers, but no so much for users in my opinion. Not a blocker if you disagree.
##########
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>
Review Comment:
This link didn't work for me and it wasn't easy to find the reference with a quick search.
##########
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
+ }
Review Comment:
From an agent's evaluation, but it looks like a reasonable argument to me:
Eager per-accumulator allocation will hurt on `GROUP BY`. `resized()` allocates `alpha_map = vec![0u64; next_pow2(k*3*6)]` and c`ounters.reserve(max(64, k*6))` at construction time, before a single row arrives. For k=100 that's ~16 KB of alpha map + ~29 KB reserved counters per accumulator; for k=10000 it's ~2 MB + ~2.9 MB. There's no `GroupsAccumulator`, so `GroupsAccumulatorAdapter` creates one of these per group. A high-cardinality grouping key with a singleton group each still pays full price. Lazy-allocate both, and consider growing the alpha map with observed cardinality rather than sizing it from k up front.
##########
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:
It's not obvious to me why we need to call `truncate_if_needed()` on a call to `serialize()`. It seems like every other place where we would mutate the struct has calls to truncate if needed in them. Did I miss one, so that we want to lazily do this only when needed to output? Otherwise we could probably remove the `mut`
##########
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);
+
+ // If there are still more counters than requested_capacity (because
+ // target_capacity wasn't reached), partition out the top ones and
+ // fold the tail into the alpha map before serializing.
+ let k = self.requested_capacity;
+ if k > 0 && 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();
+ }
+ self.rebuild_counter_map();
+ }
+
+ let num_counters = self.counters.len();
+
+ // Pre-compute total size for a single allocation.
+ let counter_bytes: usize =
+ self.counters.iter().map(|c| 4 + c.item.len() + 16).sum();
+ let total = 16 + counter_bytes + 8 + self.alpha_map.len() * 8;
+ let mut bytes = Vec::with_capacity(total);
+
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(num_counters as u64).to_le_bytes());
+
+ for counter in &self.counters {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
+ for &alpha in &self.alpha_map {
+ bytes.extend_from_slice(&alpha.to_le_bytes());
+ }
+
+ bytes
+ }
+
+ /// Deserialize a summary from bytes.
+ fn from_bytes(bytes: &[u8]) -> Result<Self> {
+ if bytes.len() < 16 {
+ return exec_err!("Invalid Space-Saving summary bytes: too short");
+ }
+
+ let requested_capacity =
+ u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
+ let num_counters = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+
+ // Validate against reasonable upper bounds to prevent OOM from
+ // malformed state.
+ if requested_capacity == 0 {
+ return exec_err!("Invalid Space-Saving summary: requested_capacity is 0");
+ }
+ let max_capacity = APPROX_TOP_K_MAX_K * CAPACITY_MULTIPLIER;
+ if requested_capacity > max_capacity {
+ return exec_err!(
+ "Invalid Space-Saving summary: requested_capacity {requested_capacity} \
+ exceeds maximum {max_capacity}"
+ );
+ }
+ if num_counters > requested_capacity {
+ return exec_err!(
+ "Invalid Space-Saving summary: num_counters {num_counters} exceeds \
+ requested_capacity {requested_capacity}"
+ );
+ }
+ // Each counter needs at least 20 bytes (4 len + 0 item + 8 count + 8 error).
+ let max_possible = (bytes.len().saturating_sub(16)) / 20;
+ if num_counters > max_possible {
+ return exec_err!(
+ "Invalid Space-Saving summary: num_counters {num_counters} exceeds \
Review Comment:
nit: Here and throughout there are references to the "Space-Saving". IMO this just the specific implementation used and references to it like here aren't adding value to the user, like when they get this error. I don't expect a consumer of the error will know what to make of "invalid space saving summary" text.
##########
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:
It would be useful to have a one or two line comment in the code about the decision so future people don't wonder the same, as I did.
##########
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.
Review Comment:
From an agent's review:
Documented ordering doesn't match actual ordering. Module docs (L28), #[user_doc], and the PR description all say "ordered by count descending," but `top_k()` sorts by `(count - error, count)`. Once errors are non-zero, the returned count column is not monotonically decreasing — a user sorting visually by count will see it out of order. Pick one: sort output by count, document lower-bound ordering, or expose error as a third struct field so the result is interpretable at all. Right now a user has no way to know a count is inflated.
##########
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);
+
+ // If there are still more counters than requested_capacity (because
+ // target_capacity wasn't reached), partition out the top ones and
+ // fold the tail into the alpha map before serializing.
+ let k = self.requested_capacity;
+ if k > 0 && 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();
+ }
+ self.rebuild_counter_map();
+ }
+
+ let num_counters = self.counters.len();
+
+ // Pre-compute total size for a single allocation.
+ let counter_bytes: usize =
+ self.counters.iter().map(|c| 4 + c.item.len() + 16).sum();
+ let total = 16 + counter_bytes + 8 + self.alpha_map.len() * 8;
+ let mut bytes = Vec::with_capacity(total);
+
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(num_counters as u64).to_le_bytes());
+
+ for counter in &self.counters {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
+ for &alpha in &self.alpha_map {
+ bytes.extend_from_slice(&alpha.to_le_bytes());
+ }
+
+ bytes
+ }
+
+ /// Deserialize a summary from bytes.
+ fn from_bytes(bytes: &[u8]) -> Result<Self> {
+ if bytes.len() < 16 {
+ return exec_err!("Invalid Space-Saving summary bytes: too short");
+ }
+
+ let requested_capacity =
+ u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
+ let num_counters = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+
+ // Validate against reasonable upper bounds to prevent OOM from
+ // malformed state.
+ if requested_capacity == 0 {
+ return exec_err!("Invalid Space-Saving summary: requested_capacity is 0");
+ }
+ let max_capacity = APPROX_TOP_K_MAX_K * CAPACITY_MULTIPLIER;
+ if requested_capacity > max_capacity {
+ return exec_err!(
+ "Invalid Space-Saving summary: requested_capacity {requested_capacity} \
+ exceeds maximum {max_capacity}"
+ );
+ }
+ if num_counters > requested_capacity {
+ return exec_err!(
+ "Invalid Space-Saving summary: num_counters {num_counters} exceeds \
+ requested_capacity {requested_capacity}"
+ );
+ }
+ // Each counter needs at least 20 bytes (4 len + 0 item + 8 count + 8 error).
+ let max_possible = (bytes.len().saturating_sub(16)) / 20;
+ if num_counters > max_possible {
+ return exec_err!(
+ "Invalid Space-Saving summary: num_counters {num_counters} exceeds \
+ what fits in {} bytes",
+ bytes.len()
+ );
+ }
+
+ let mut counters = Vec::with_capacity(num_counters);
+ let mut counter_map = HashTable::with_capacity(num_counters);
+ let mut item_heap_bytes: usize = 0;
+ let mut offset: usize = 16;
+
+ for idx in 0..num_counters {
+ if offset.checked_add(4).is_none_or(|end| end > bytes.len()) {
+ return exec_err!(
+ "Invalid Space-Saving summary bytes: truncated item length"
+ );
+ }
+ let item_len =
+ u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
+ as usize;
+ offset += 4;
+
+ let needed = item_len.checked_add(16);
+ if needed
+ .is_none_or(|n| offset.checked_add(n).is_none_or(|end| end > bytes.len()))
+ {
+ return exec_err!(
+ "Invalid Space-Saving summary bytes: truncated counter"
+ );
+ }
+
+ let item = bytes[offset..offset + item_len].to_vec();
+ offset += item_len;
+
+ let count = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let error = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let hash = Self::hash_item(&item);
+ item_heap_bytes += item.capacity();
+ counter_map.insert_unique(hash, (hash, idx), |&(h, _)| h);
+ counters.push(Counter {
+ item,
+ hash,
+ count,
+ error,
+ });
+ }
+
+ if offset.checked_add(8).is_none_or(|end| end > bytes.len()) {
+ return exec_err!(
+ "Invalid Space-Saving summary bytes: missing alpha map size"
+ );
+ }
+ let alpha_map_size =
+ u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap()) as usize;
+ offset += 8;
+
+ // Validate alpha_map_size is a power of two (required for bitmask indexing).
+ if alpha_map_size == 0 || !alpha_map_size.is_power_of_two() {
+ return exec_err!(
+ "Invalid Space-Saving summary: alpha_map_size {alpha_map_size} \
+ is not a positive power of two"
+ );
+ }
+
+ let alpha_bytes = alpha_map_size
+ .checked_mul(8)
+ .and_then(|n| offset.checked_add(n));
+ if alpha_bytes.is_none_or(|end| end > bytes.len()) {
+ return exec_err!("Invalid Space-Saving summary bytes: truncated alpha map");
+ }
+
+ let mut alpha_map = Vec::with_capacity(alpha_map_size);
+ for _ in 0..alpha_map_size {
+ let alpha = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+ alpha_map.push(alpha);
+ }
+
+ let target_capacity = max(64, requested_capacity.saturating_mul(2));
+
+ Ok(Self {
+ counters,
+ counter_map,
+ alpha_map,
+ requested_capacity,
+ target_capacity,
+ item_heap_bytes,
+ })
+ }
+
+ /// Approximate size in bytes of this summary. O(1) thanks to
+ /// incremental `item_heap_bytes` tracking.
+ fn size(&self) -> usize {
+ size_of::<Self>()
+ + self.counters.capacity() * size_of::<Counter>()
+ + self.item_heap_bytes
+ // HashTable<(u64, usize)>: each bucket stores (hash, idx) + 1 control byte.
+ + self.counter_map.capacity()
+ * (size_of::<(u64, usize)>() + size_of::<u8>())
+ + self.alpha_map.capacity() * size_of::<u64>()
+ }
+}
+
+// ---------------------------------------------------------------------------
+// ApproxTopK UDAF struct
+// ---------------------------------------------------------------------------
+
+/// Approximate top-k UDAF using the Filtered Space-Saving algorithm.
+#[user_doc(
+ doc_section(label = "Approximate Functions"),
+ description = r#"Returns the approximate most frequent (top-k) values with their estimated counts, using the Filtered Space-Saving algorithm. The returned counts are upper-bound estimates; the true frequency lies in `[count - error, count]`. NULL values are skipped; an empty or all-NULL input returns an empty list `[]`. For float columns, -0.0 and +0.0 are treated as distinct values, and different NaN representations are tracked separately."#,
+ syntax_example = "approx_top_k(expression, k)",
+ sql_example = r#"```sql
+> SELECT approx_top_k(column_name, 3) FROM table_name;
++-----------------------------------------------------------------------------+
+| approx_top_k(column_name,Int64(3)) |
++-----------------------------------------------------------------------------+
+| [{value: foo, count: 3}, {value: bar, count: 2}, {value: baz, count: 1}] |
++-----------------------------------------------------------------------------+
+```"#,
+ standard_argument(name = "expression",),
+ argument(
+ name = "k",
+ description = "The number of top elements to return. Must be a literal integer between 1 and 10,000."
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ApproxTopK {
+ signature: Signature,
+}
+
+impl Default for ApproxTopK {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ApproxTopK {
+ pub fn new() -> Self {
+ // Supported value types for the first argument.
+ let value_types = &[
+ DataType::Utf8,
+ DataType::LargeUtf8,
+ DataType::Binary,
+ DataType::LargeBinary,
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ DataType::Float32,
+ DataType::Float64,
+ DataType::Date32,
+ DataType::Date64,
+ DataType::Timestamp(TimeUnit::Second, None),
+ DataType::Timestamp(TimeUnit::Millisecond, None),
+ DataType::Timestamp(TimeUnit::Microsecond, None),
+ DataType::Timestamp(TimeUnit::Nanosecond, None),
+ DataType::Timestamp(TimeUnit::Second, Some(TIMEZONE_WILDCARD.into())),
+ DataType::Timestamp(TimeUnit::Millisecond, Some(TIMEZONE_WILDCARD.into())),
+ DataType::Timestamp(TimeUnit::Microsecond, Some(TIMEZONE_WILDCARD.into())),
+ DataType::Timestamp(TimeUnit::Nanosecond, Some(TIMEZONE_WILDCARD.into())),
+ ];
+ // k must be a literal integer; accept any integer type for convenience.
+ let k_types = &[
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ ];
+
+ let mut variants = Vec::with_capacity(value_types.len() * k_types.len());
+ for vt in value_types {
+ for kt in k_types {
+ variants.push(TypeSignature::Exact(vec![vt.clone(), kt.clone()]));
+ }
+ }
+
+ Self {
+ signature: Signature::one_of(variants, Volatility::Immutable),
+ }
+ }
+}
+
+impl AggregateUDFImpl for ApproxTopK {
+ fn name(&self) -> &str {
+ "approx_top_k"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ let value_type = if !arg_types.is_empty() {
+ match &arg_types[0] {
+ // Large variants are narrowed: the summary stores in-memory
+ // byte slices, not large column offsets, so i32 offsets suffice.
+ DataType::LargeUtf8 => DataType::Utf8,
+ DataType::LargeBinary => DataType::Binary,
+ other => other.clone(),
+ }
+ } else {
+ DataType::Utf8
+ };
+
+ let struct_fields = Fields::from(vec![
+ Field::new("value", value_type, true),
+ Field::new("count", DataType::UInt64, false),
+ ]);
+ Ok(DataType::List(Arc::new(Field::new(
+ "item",
+ DataType::Struct(struct_fields),
+ true,
+ ))))
+ }
+
+ fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
+ Ok(vec![Arc::new(Field::new(
+ format_state_name(args.name, "summary"),
+ DataType::Binary,
+ true,
+ ))])
+ }
+
+ fn accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+ if args.exprs.len() < 2 {
+ return plan_err!("approx_top_k requires two arguments: column and k");
+ }
+
+ let k_expr = &args.exprs[1];
+ let k = k_expr
+ .downcast_ref::<datafusion_physical_expr::expressions::Literal>()
+ .and_then(|lit| match lit.value() {
+ // Guard against negative values before casting to usize to
+ // avoid wrapping (e.g. -1i32 as usize → u64::MAX).
+ // Zero is allowed through and caught by the bounds check below.
+ ScalarValue::Int8(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int16(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int32(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int64(Some(v)) if *v >= 0 => Some(*v as usize),
Review Comment:
Should these all be `*v > 0` ? I don't know what top k of 0 values would mean.
##########
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);
+
+ // If there are still more counters than requested_capacity (because
+ // target_capacity wasn't reached), partition out the top ones and
+ // fold the tail into the alpha map before serializing.
+ let k = self.requested_capacity;
+ if k > 0 && 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();
+ }
+ self.rebuild_counter_map();
+ }
+
+ let num_counters = self.counters.len();
+
+ // Pre-compute total size for a single allocation.
+ let counter_bytes: usize =
+ self.counters.iter().map(|c| 4 + c.item.len() + 16).sum();
+ let total = 16 + counter_bytes + 8 + self.alpha_map.len() * 8;
+ let mut bytes = Vec::with_capacity(total);
+
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(num_counters as u64).to_le_bytes());
+
+ for counter in &self.counters {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
+ for &alpha in &self.alpha_map {
+ bytes.extend_from_slice(&alpha.to_le_bytes());
+ }
+
+ bytes
+ }
+
+ /// Deserialize a summary from bytes.
+ fn from_bytes(bytes: &[u8]) -> Result<Self> {
+ if bytes.len() < 16 {
+ return exec_err!("Invalid Space-Saving summary bytes: too short");
+ }
+
+ let requested_capacity =
+ u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
+ let num_counters = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+
+ // Validate against reasonable upper bounds to prevent OOM from
+ // malformed state.
+ if requested_capacity == 0 {
+ return exec_err!("Invalid Space-Saving summary: requested_capacity is 0");
+ }
+ let max_capacity = APPROX_TOP_K_MAX_K * CAPACITY_MULTIPLIER;
+ if requested_capacity > max_capacity {
+ return exec_err!(
+ "Invalid Space-Saving summary: requested_capacity {requested_capacity} \
+ exceeds maximum {max_capacity}"
+ );
+ }
+ if num_counters > requested_capacity {
+ return exec_err!(
+ "Invalid Space-Saving summary: num_counters {num_counters} exceeds \
+ requested_capacity {requested_capacity}"
+ );
+ }
+ // Each counter needs at least 20 bytes (4 len + 0 item + 8 count + 8 error).
+ let max_possible = (bytes.len().saturating_sub(16)) / 20;
+ if num_counters > max_possible {
+ return exec_err!(
+ "Invalid Space-Saving summary: num_counters {num_counters} exceeds \
+ what fits in {} bytes",
+ bytes.len()
+ );
+ }
+
+ let mut counters = Vec::with_capacity(num_counters);
+ let mut counter_map = HashTable::with_capacity(num_counters);
+ let mut item_heap_bytes: usize = 0;
+ let mut offset: usize = 16;
+
+ for idx in 0..num_counters {
+ if offset.checked_add(4).is_none_or(|end| end > bytes.len()) {
+ return exec_err!(
+ "Invalid Space-Saving summary bytes: truncated item length"
+ );
+ }
+ let item_len =
+ u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
+ as usize;
+ offset += 4;
+
+ let needed = item_len.checked_add(16);
+ if needed
+ .is_none_or(|n| offset.checked_add(n).is_none_or(|end| end > bytes.len()))
+ {
+ return exec_err!(
+ "Invalid Space-Saving summary bytes: truncated counter"
+ );
+ }
+
+ let item = bytes[offset..offset + item_len].to_vec();
+ offset += item_len;
+
+ let count = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let error = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let hash = Self::hash_item(&item);
+ item_heap_bytes += item.capacity();
+ counter_map.insert_unique(hash, (hash, idx), |&(h, _)| h);
+ counters.push(Counter {
+ item,
+ hash,
+ count,
+ error,
+ });
+ }
+
+ if offset.checked_add(8).is_none_or(|end| end > bytes.len()) {
+ return exec_err!(
+ "Invalid Space-Saving summary bytes: missing alpha map size"
+ );
+ }
+ let alpha_map_size =
+ u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap()) as usize;
+ offset += 8;
+
+ // Validate alpha_map_size is a power of two (required for bitmask indexing).
+ if alpha_map_size == 0 || !alpha_map_size.is_power_of_two() {
+ return exec_err!(
+ "Invalid Space-Saving summary: alpha_map_size {alpha_map_size} \
+ is not a positive power of two"
+ );
+ }
+
+ let alpha_bytes = alpha_map_size
+ .checked_mul(8)
+ .and_then(|n| offset.checked_add(n));
+ if alpha_bytes.is_none_or(|end| end > bytes.len()) {
+ return exec_err!("Invalid Space-Saving summary bytes: truncated alpha map");
+ }
+
+ let mut alpha_map = Vec::with_capacity(alpha_map_size);
+ for _ in 0..alpha_map_size {
+ let alpha = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+ alpha_map.push(alpha);
+ }
+
+ let target_capacity = max(64, requested_capacity.saturating_mul(2));
+
+ Ok(Self {
+ counters,
+ counter_map,
+ alpha_map,
+ requested_capacity,
+ target_capacity,
+ item_heap_bytes,
+ })
+ }
+
+ /// Approximate size in bytes of this summary. O(1) thanks to
+ /// incremental `item_heap_bytes` tracking.
+ fn size(&self) -> usize {
+ size_of::<Self>()
+ + self.counters.capacity() * size_of::<Counter>()
+ + self.item_heap_bytes
+ // HashTable<(u64, usize)>: each bucket stores (hash, idx) + 1 control byte.
+ + self.counter_map.capacity()
+ * (size_of::<(u64, usize)>() + size_of::<u8>())
+ + self.alpha_map.capacity() * size_of::<u64>()
+ }
+}
+
+// ---------------------------------------------------------------------------
+// ApproxTopK UDAF struct
+// ---------------------------------------------------------------------------
+
+/// Approximate top-k UDAF using the Filtered Space-Saving algorithm.
+#[user_doc(
+ doc_section(label = "Approximate Functions"),
+ description = r#"Returns the approximate most frequent (top-k) values with their estimated counts, using the Filtered Space-Saving algorithm. The returned counts are upper-bound estimates; the true frequency lies in `[count - error, count]`. NULL values are skipped; an empty or all-NULL input returns an empty list `[]`. For float columns, -0.0 and +0.0 are treated as distinct values, and different NaN representations are tracked separately."#,
+ syntax_example = "approx_top_k(expression, k)",
+ sql_example = r#"```sql
+> SELECT approx_top_k(column_name, 3) FROM table_name;
++-----------------------------------------------------------------------------+
+| approx_top_k(column_name,Int64(3)) |
++-----------------------------------------------------------------------------+
+| [{value: foo, count: 3}, {value: bar, count: 2}, {value: baz, count: 1}] |
++-----------------------------------------------------------------------------+
+```"#,
+ standard_argument(name = "expression",),
+ argument(
+ name = "k",
+ description = "The number of top elements to return. Must be a literal integer between 1 and 10,000."
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ApproxTopK {
+ signature: Signature,
+}
+
+impl Default for ApproxTopK {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ApproxTopK {
+ pub fn new() -> Self {
+ // Supported value types for the first argument.
+ let value_types = &[
+ DataType::Utf8,
+ DataType::LargeUtf8,
+ DataType::Binary,
+ DataType::LargeBinary,
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ DataType::Float32,
+ DataType::Float64,
+ DataType::Date32,
+ DataType::Date64,
+ DataType::Timestamp(TimeUnit::Second, None),
+ DataType::Timestamp(TimeUnit::Millisecond, None),
+ DataType::Timestamp(TimeUnit::Microsecond, None),
+ DataType::Timestamp(TimeUnit::Nanosecond, None),
+ DataType::Timestamp(TimeUnit::Second, Some(TIMEZONE_WILDCARD.into())),
+ DataType::Timestamp(TimeUnit::Millisecond, Some(TIMEZONE_WILDCARD.into())),
+ DataType::Timestamp(TimeUnit::Microsecond, Some(TIMEZONE_WILDCARD.into())),
+ DataType::Timestamp(TimeUnit::Nanosecond, Some(TIMEZONE_WILDCARD.into())),
+ ];
+ // k must be a literal integer; accept any integer type for convenience.
+ let k_types = &[
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ ];
+
+ let mut variants = Vec::with_capacity(value_types.len() * k_types.len());
+ for vt in value_types {
+ for kt in k_types {
+ variants.push(TypeSignature::Exact(vec![vt.clone(), kt.clone()]));
+ }
+ }
+
+ Self {
+ signature: Signature::one_of(variants, Volatility::Immutable),
+ }
+ }
+}
+
+impl AggregateUDFImpl for ApproxTopK {
+ fn name(&self) -> &str {
+ "approx_top_k"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ let value_type = if !arg_types.is_empty() {
+ match &arg_types[0] {
+ // Large variants are narrowed: the summary stores in-memory
+ // byte slices, not large column offsets, so i32 offsets suffice.
+ DataType::LargeUtf8 => DataType::Utf8,
+ DataType::LargeBinary => DataType::Binary,
+ other => other.clone(),
+ }
+ } else {
+ DataType::Utf8
+ };
+
+ let struct_fields = Fields::from(vec![
+ Field::new("value", value_type, true),
+ Field::new("count", DataType::UInt64, false),
+ ]);
+ Ok(DataType::List(Arc::new(Field::new(
+ "item",
+ DataType::Struct(struct_fields),
+ true,
+ ))))
+ }
+
+ fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
+ Ok(vec![Arc::new(Field::new(
+ format_state_name(args.name, "summary"),
+ DataType::Binary,
+ true,
+ ))])
+ }
+
+ fn accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+ if args.exprs.len() < 2 {
+ return plan_err!("approx_top_k requires two arguments: column and k");
+ }
+
+ let k_expr = &args.exprs[1];
+ let k = k_expr
+ .downcast_ref::<datafusion_physical_expr::expressions::Literal>()
+ .and_then(|lit| match lit.value() {
+ // Guard against negative values before casting to usize to
+ // avoid wrapping (e.g. -1i32 as usize → u64::MAX).
+ // Zero is allowed through and caught by the bounds check below.
+ ScalarValue::Int8(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int16(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int32(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int64(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::UInt8(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt16(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt32(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt64(Some(v)) => Some(*v as usize),
+ _ => None,
+ });
+
+ let Some(k) = k else {
+ return plan_err!(
+ "approx_top_k requires k to be a positive literal integer \
+ between 1 and 10000"
+ );
+ };
+
+ if k == 0 || k > APPROX_TOP_K_MAX_K {
+ return plan_err!(
+ "approx_top_k requires k to be between 1 and {APPROX_TOP_K_MAX_K}, got {k}"
+ );
+ }
+
+ let data_type = args.expr_fields[0].data_type().clone();
+ Ok(Box::new(ApproxTopKAccumulator::new_with_data_type(
+ k, data_type,
+ )))
+ }
+
+ fn documentation(&self) -> Option<&Documentation> {
+ self.doc()
+ }
+}
+
+// ---------------------------------------------------------------------------
+// Accumulator
+// ---------------------------------------------------------------------------
+
+/// Accumulator for `approx_top_k` using the Filtered Space-Saving algorithm.
+#[derive(Debug)]
+struct ApproxTopKAccumulator {
+ summary: SpaceSavingSummary,
+ k: usize,
+ /// The data type of the input column.
+ input_data_type: DataType,
+}
+
+impl ApproxTopKAccumulator {
+ fn new_with_data_type(k: usize, input_data_type: DataType) -> Self {
+ let capacity = k * CAPACITY_MULTIPLIER;
+ Self {
+ summary: SpaceSavingSummary::new(capacity),
+ k,
+ input_data_type,
+ }
+ }
+
+ /// Build the value array for the result based on the input data type.
+ fn build_value_array(&self, top_items: &[(&[u8], u64, u64)]) -> Result<ArrayRef> {
+ match &self.input_data_type {
+ DataType::Utf8 | DataType::LargeUtf8 => {
Review Comment:
Do we have Utf8View support?
##########
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);
+
+ // If there are still more counters than requested_capacity (because
+ // target_capacity wasn't reached), partition out the top ones and
+ // fold the tail into the alpha map before serializing.
+ let k = self.requested_capacity;
+ if k > 0 && 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();
+ }
+ self.rebuild_counter_map();
+ }
+
+ let num_counters = self.counters.len();
+
+ // Pre-compute total size for a single allocation.
+ let counter_bytes: usize =
+ self.counters.iter().map(|c| 4 + c.item.len() + 16).sum();
+ let total = 16 + counter_bytes + 8 + self.alpha_map.len() * 8;
+ let mut bytes = Vec::with_capacity(total);
+
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(num_counters as u64).to_le_bytes());
+
+ for counter in &self.counters {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
Review Comment:
From an agent's review:
The serialized state carries the entire alpha map, and `merge()` never reads it. `serialize()` writes `alpha_map.len() * 8` bytes; `merge()` only sums alpha maps element-wise (L411–418) and never consults them for admission or ranking. Alphas only affect `insert()`, which the final-stage accumulator never calls. So every partial state pays ≥16 KB (2 MB at max k) of shuffle bytes, per group per partition, for data that has no effect on the result. Either drop it from the wire format or explain what path consumes it.
##########
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),
+ );
Review Comment:
From an agent's review:
Does the alpha map actually change retention? This is the one I'd most want the author to answer. At capacity, `insert()` admits unconditionally with `count = alpha + increment`, `error = alpha + error`. Eviction ranks via `cmp_by_rank`, whose primary key is `count - error` (L112–116) — which for a freshly readmitted item is exactly increment, independent of alpha. So alpha degenerates to a tiebreak on count among items with equal lower bounds; an item with alpha = 10⁶ still loses to any resident item seen twice. ClickHouse's SpaceSaving evicts the minimum by count (the upper bound), which is what makes the alpha boost meaningful and what makes the "filter low-frequency noise before it enters the summary" claim true. Here nothing is filtered — everything is admitted and bulk-evicted at 2× capacity. Either the ranking key should be count, or the docs/PR description shouldn't claim ClickHouse-equivalent FSS semantics. test_space_saving_alpha_map (L1209) actually asserts the
degenerate behavior (count == error + 1) without flagging it.
##########
datafusion/sqllogictest/test_files/approx_top_k.slt:
##########
@@ -0,0 +1,87 @@
+# 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.
+
+###########
+# approx_top_k tests
+###########
Review Comment:
From an agent's review:
Nothing tests the approximation. Every SLT case and all but two unit tests have distinct-count ≤ capacity, so error is always 0 and the function is exact — which means the entire interesting half of this PR is untested. Missing: a case where eviction changes the reported result, a non-zero-error assertion, a multi-partition merge where an item exists in only one partition and has been evicted from the other, and a recall assertion on a skewed (Zipfian) input. The SLT also never sets target_partitions, so the serialize → from_bytes → merge path is never exercised end-to-end from SQL — only through hand-built unit tests. And assert!(evicted_count >= 61) (L1205) is a very loose assertion for the eviction test.
##########
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);
+
+ // If there are still more counters than requested_capacity (because
+ // target_capacity wasn't reached), partition out the top ones and
+ // fold the tail into the alpha map before serializing.
+ let k = self.requested_capacity;
+ if k > 0 && 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();
+ }
+ self.rebuild_counter_map();
+ }
+
+ let num_counters = self.counters.len();
+
+ // Pre-compute total size for a single allocation.
+ let counter_bytes: usize =
+ self.counters.iter().map(|c| 4 + c.item.len() + 16).sum();
+ let total = 16 + counter_bytes + 8 + self.alpha_map.len() * 8;
+ let mut bytes = Vec::with_capacity(total);
+
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(num_counters as u64).to_le_bytes());
+
+ for counter in &self.counters {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
+ for &alpha in &self.alpha_map {
+ bytes.extend_from_slice(&alpha.to_le_bytes());
+ }
+
+ bytes
+ }
+
+ /// Deserialize a summary from bytes.
+ fn from_bytes(bytes: &[u8]) -> Result<Self> {
+ if bytes.len() < 16 {
+ return exec_err!("Invalid Space-Saving summary bytes: too short");
+ }
+
+ let requested_capacity =
+ u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
+ let num_counters = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+
+ // Validate against reasonable upper bounds to prevent OOM from
+ // malformed state.
+ if requested_capacity == 0 {
+ return exec_err!("Invalid Space-Saving summary: requested_capacity is 0");
+ }
+ let max_capacity = APPROX_TOP_K_MAX_K * CAPACITY_MULTIPLIER;
+ if requested_capacity > max_capacity {
+ return exec_err!(
+ "Invalid Space-Saving summary: requested_capacity {requested_capacity} \
+ exceeds maximum {max_capacity}"
+ );
+ }
+ if num_counters > requested_capacity {
+ return exec_err!(
+ "Invalid Space-Saving summary: num_counters {num_counters} exceeds \
+ requested_capacity {requested_capacity}"
+ );
+ }
+ // Each counter needs at least 20 bytes (4 len + 0 item + 8 count + 8 error).
+ let max_possible = (bytes.len().saturating_sub(16)) / 20;
+ if num_counters > max_possible {
+ return exec_err!(
+ "Invalid Space-Saving summary: num_counters {num_counters} exceeds \
+ what fits in {} bytes",
+ bytes.len()
+ );
+ }
+
+ let mut counters = Vec::with_capacity(num_counters);
+ let mut counter_map = HashTable::with_capacity(num_counters);
+ let mut item_heap_bytes: usize = 0;
+ let mut offset: usize = 16;
+
+ for idx in 0..num_counters {
+ if offset.checked_add(4).is_none_or(|end| end > bytes.len()) {
+ return exec_err!(
+ "Invalid Space-Saving summary bytes: truncated item length"
+ );
+ }
+ let item_len =
+ u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
+ as usize;
+ offset += 4;
+
+ let needed = item_len.checked_add(16);
+ if needed
+ .is_none_or(|n| offset.checked_add(n).is_none_or(|end| end > bytes.len()))
+ {
+ return exec_err!(
+ "Invalid Space-Saving summary bytes: truncated counter"
+ );
+ }
+
+ let item = bytes[offset..offset + item_len].to_vec();
+ offset += item_len;
+
+ let count = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let error = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let hash = Self::hash_item(&item);
+ item_heap_bytes += item.capacity();
+ counter_map.insert_unique(hash, (hash, idx), |&(h, _)| h);
+ counters.push(Counter {
+ item,
+ hash,
+ count,
+ error,
+ });
+ }
+
+ if offset.checked_add(8).is_none_or(|end| end > bytes.len()) {
+ return exec_err!(
+ "Invalid Space-Saving summary bytes: missing alpha map size"
+ );
+ }
+ let alpha_map_size =
+ u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap()) as usize;
+ offset += 8;
+
+ // Validate alpha_map_size is a power of two (required for bitmask indexing).
+ if alpha_map_size == 0 || !alpha_map_size.is_power_of_two() {
+ return exec_err!(
+ "Invalid Space-Saving summary: alpha_map_size {alpha_map_size} \
+ is not a positive power of two"
+ );
+ }
+
+ let alpha_bytes = alpha_map_size
+ .checked_mul(8)
+ .and_then(|n| offset.checked_add(n));
+ if alpha_bytes.is_none_or(|end| end > bytes.len()) {
+ return exec_err!("Invalid Space-Saving summary bytes: truncated alpha map");
+ }
+
+ let mut alpha_map = Vec::with_capacity(alpha_map_size);
+ for _ in 0..alpha_map_size {
+ let alpha = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+ alpha_map.push(alpha);
+ }
+
+ let target_capacity = max(64, requested_capacity.saturating_mul(2));
+
+ Ok(Self {
+ counters,
+ counter_map,
+ alpha_map,
+ requested_capacity,
+ target_capacity,
+ item_heap_bytes,
+ })
+ }
+
+ /// Approximate size in bytes of this summary. O(1) thanks to
+ /// incremental `item_heap_bytes` tracking.
+ fn size(&self) -> usize {
+ size_of::<Self>()
+ + self.counters.capacity() * size_of::<Counter>()
+ + self.item_heap_bytes
+ // HashTable<(u64, usize)>: each bucket stores (hash, idx) + 1 control byte.
+ + self.counter_map.capacity()
+ * (size_of::<(u64, usize)>() + size_of::<u8>())
+ + self.alpha_map.capacity() * size_of::<u64>()
+ }
+}
+
+// ---------------------------------------------------------------------------
+// ApproxTopK UDAF struct
+// ---------------------------------------------------------------------------
+
+/// Approximate top-k UDAF using the Filtered Space-Saving algorithm.
+#[user_doc(
+ doc_section(label = "Approximate Functions"),
+ description = r#"Returns the approximate most frequent (top-k) values with their estimated counts, using the Filtered Space-Saving algorithm. The returned counts are upper-bound estimates; the true frequency lies in `[count - error, count]`. NULL values are skipped; an empty or all-NULL input returns an empty list `[]`. For float columns, -0.0 and +0.0 are treated as distinct values, and different NaN representations are tracked separately."#,
+ syntax_example = "approx_top_k(expression, k)",
+ sql_example = r#"```sql
+> SELECT approx_top_k(column_name, 3) FROM table_name;
++-----------------------------------------------------------------------------+
+| approx_top_k(column_name,Int64(3)) |
++-----------------------------------------------------------------------------+
+| [{value: foo, count: 3}, {value: bar, count: 2}, {value: baz, count: 1}] |
++-----------------------------------------------------------------------------+
+```"#,
+ standard_argument(name = "expression",),
+ argument(
+ name = "k",
+ description = "The number of top elements to return. Must be a literal integer between 1 and 10,000."
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ApproxTopK {
+ signature: Signature,
+}
+
+impl Default for ApproxTopK {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ApproxTopK {
+ pub fn new() -> Self {
+ // Supported value types for the first argument.
+ let value_types = &[
+ DataType::Utf8,
+ DataType::LargeUtf8,
+ DataType::Binary,
+ DataType::LargeBinary,
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ DataType::Float32,
+ DataType::Float64,
+ DataType::Date32,
+ DataType::Date64,
+ DataType::Timestamp(TimeUnit::Second, None),
+ DataType::Timestamp(TimeUnit::Millisecond, None),
+ DataType::Timestamp(TimeUnit::Microsecond, None),
+ DataType::Timestamp(TimeUnit::Nanosecond, None),
+ DataType::Timestamp(TimeUnit::Second, Some(TIMEZONE_WILDCARD.into())),
+ DataType::Timestamp(TimeUnit::Millisecond, Some(TIMEZONE_WILDCARD.into())),
+ DataType::Timestamp(TimeUnit::Microsecond, Some(TIMEZONE_WILDCARD.into())),
+ DataType::Timestamp(TimeUnit::Nanosecond, Some(TIMEZONE_WILDCARD.into())),
+ ];
+ // k must be a literal integer; accept any integer type for convenience.
+ let k_types = &[
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ ];
+
+ let mut variants = Vec::with_capacity(value_types.len() * k_types.len());
+ for vt in value_types {
+ for kt in k_types {
+ variants.push(TypeSignature::Exact(vec![vt.clone(), kt.clone()]));
+ }
+ }
+
+ Self {
+ signature: Signature::one_of(variants, Volatility::Immutable),
+ }
+ }
+}
+
+impl AggregateUDFImpl for ApproxTopK {
+ fn name(&self) -> &str {
+ "approx_top_k"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ let value_type = if !arg_types.is_empty() {
+ match &arg_types[0] {
+ // Large variants are narrowed: the summary stores in-memory
+ // byte slices, not large column offsets, so i32 offsets suffice.
+ DataType::LargeUtf8 => DataType::Utf8,
+ DataType::LargeBinary => DataType::Binary,
+ other => other.clone(),
+ }
+ } else {
+ DataType::Utf8
+ };
+
+ let struct_fields = Fields::from(vec![
+ Field::new("value", value_type, true),
+ Field::new("count", DataType::UInt64, false),
+ ]);
+ Ok(DataType::List(Arc::new(Field::new(
+ "item",
+ DataType::Struct(struct_fields),
+ true,
+ ))))
+ }
+
+ fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
+ Ok(vec![Arc::new(Field::new(
+ format_state_name(args.name, "summary"),
+ DataType::Binary,
+ true,
+ ))])
+ }
+
+ fn accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+ if args.exprs.len() < 2 {
+ return plan_err!("approx_top_k requires two arguments: column and k");
+ }
+
+ let k_expr = &args.exprs[1];
+ let k = k_expr
+ .downcast_ref::<datafusion_physical_expr::expressions::Literal>()
+ .and_then(|lit| match lit.value() {
+ // Guard against negative values before casting to usize to
+ // avoid wrapping (e.g. -1i32 as usize → u64::MAX).
+ // Zero is allowed through and caught by the bounds check below.
+ ScalarValue::Int8(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int16(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int32(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int64(Some(v)) if *v >= 0 => Some(*v as usize),
Review Comment:
I see you capture it in the lines below, but could trivially include here instead.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#issuecomment-6128906091
Hi @timsaucer, thanks again for reviewing this PR. I’ve addressed the feedback and updated the implementation. Could you take another look when you have a chance?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "martin-g (via GitHub)" <gi...@apache.org>.
martin-g commented on PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#issuecomment-4192538713
No problem!
It would be good a second person to review the PR anyway.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r3123517113
##########
datafusion/functions-aggregate/src/approx_top_k.rs:
##########
@@ -0,0 +1,1421 @@
+// 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, 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};
+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; 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) -> 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>,
Review Comment:
`Vec<u8>`: this is quite wasteful to store in the map.
we can store the hash inline in the map
`HashTable<u64, usize>`
and look up the `item` using the `usize` pointer using `counter_map.find`.
(Look up the pattern in other places like aggregation / join ...)
This saves many allocations / rehashing.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#issuecomment-4279310521
Hi @martin-g and @Dandandan ! Friendly ping on this when you have a moment. I know you mentioned wanting a second person to review - let me know if there is someone specific I should tag, or if there is anything else I can do from my side to help get this across the finish line.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
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:
Good point! I looked into it - foldhash seems great for batch columnar hashing (42% faster for 8k rows). But here the FNV-1a is only used for alpha map bucket indexing, hashing one item at a time (typically 1-8 bytes for integers, short strings). At that scale there's no meaningful difference, and the 5-line inline FNV avoids adding a dependency from functions-aggregate to foldhash. It also matches what ClickHouse does in their SpaceSaving.h. Happy to switch if you feel strongly though!
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r2946426381
##########
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:
Good catch! Added truncate_if_needed(false) at the start of to_bytes as a defensive measure. In practice truncate_if_needed is called inside add() during update_batch, so by the time state() calls to_bytes the counters should already be truncated - but better safe than sorry, especially since it also ensures the alpha map is fully up to date before serialization.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#issuecomment-5205242709
## [Codecov](https://app.codecov.io/gh/apache/datafusion/pull/20968?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:white_check_mark: All modified and coverable lines are covered by tests.
:white_check_mark: Project coverage is 81.05%. Comparing base ([`3e3a92d`](https://app.codecov.io/gh/apache/datafusion/commit/3e3a92de29ed3d454e72c7bade6328508b6098c6?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`4e651ba`](https://app.codecov.io/gh/apache/datafusion/commit/4e651bacf1acfefc04d07aa1b3e4c8141a9aae89?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 8 commits behind head on main.
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## main #20968 +/- ##
==========================================
+ Coverage 80.91% 81.05% +0.13%
==========================================
Files 1103 1106 +3
Lines 377219 382443 +5224
Branches 377219 382443 +5224
==========================================
+ Hits 305244 309986 +4742
- Misses 53775 54165 +390
- Partials 18200 18292 +92
```
</details>
[:umbrella: View full report in Codecov by Harness](https://app.codecov.io/gh/apache/datafusion/pull/20968?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r2946356469
##########
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;
Review Comment:
The value 6 is empirically tuned to hit the sweet spot where collision rates are low enough for effective filtering without excessive memory overhead.
Clickhouse doesn't expose it either, and also uses the same const from the paper:
https://github.com/ClickHouse/ClickHouse/blob/b2b79b1c6c22e699a7c075fa6bfdb3fa788b66ca/src/Common/SpaceSaving.h#L81
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r2946434166
##########
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();
+ let return_size = std::cmp::min(sorted.len(), self.requested_capacity);
+ if return_size > 0 && return_size < sorted.len() {
+ // After select_nth_unstable_by, the top-k counters are in
+ // sorted[..return_size] but in arbitrary order. This is fine
+ // since deserialization doesn't depend on counter ordering.
+ sorted.select_nth_unstable_by(return_size - 1, |a, b| a.cmp_by_rank(b));
+ }
+ sorted.truncate(return_size);
+ sorted
+ };
+
+ let mut bytes = Vec::new();
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(counters_to_write.len() as u64).to_le_bytes());
+
+ for counter in counters_to_write {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
+ for &alpha in &self.alpha_map {
+ bytes.extend_from_slice(&alpha.to_le_bytes());
+ }
+
+ bytes
+ }
+
+ /// Deserialize a summary from bytes.
+ fn from_bytes(bytes: &[u8]) -> Result<Self> {
+ if bytes.len() < 16 {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: too short".to_string(),
+ ));
+ }
+
+ let requested_capacity =
+ u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
+ let num_counters = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+
+ let mut counters = Vec::with_capacity(num_counters);
+ let mut counter_map = HashMap::with_capacity(num_counters);
+ let mut offset = 16;
+
+ for idx in 0..num_counters {
+ if offset + 4 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated item length"
+ .to_string(),
+ ));
+ }
+ let item_len =
+ u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
+ as usize;
+ offset += 4;
+
+ if offset + item_len + 16 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated counter".to_string(),
+ ));
+ }
+
+ let item = bytes[offset..offset + item_len].to_vec();
+ offset += item_len;
+
+ let count = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let error = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let hash = Self::hash_item(&item);
+ counter_map.insert(item.clone(), idx);
+ counters.push(Counter {
+ item,
+ hash,
+ count,
+ error,
+ });
+ }
+
+ if offset + 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: missing alpha map size".to_string(),
+ ));
+ }
+ let alpha_map_size =
+ u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap()) as usize;
+ offset += 8;
+
+ if offset + alpha_map_size * 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated alpha map".to_string(),
+ ));
+ }
+
+ let mut alpha_map = Vec::with_capacity(alpha_map_size);
+ for _ in 0..alpha_map_size {
+ let alpha = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+ alpha_map.push(alpha);
+ }
+
+ let target_capacity = std::cmp::max(64, requested_capacity * 2);
+
+ Ok(Self {
+ counters,
+ counter_map,
+ alpha_map,
+ requested_capacity,
+ target_capacity,
+ })
+ }
+
+ /// Approximate size in bytes of this summary.
+ fn size(&self) -> usize {
+ size_of::<Self>()
+ + self
+ .counters
+ .iter()
+ .map(|c| size_of::<Counter>() + c.item.len())
+ .sum::<usize>()
+ // HashMap overhead: each bucket has a key-value pair plus control byte metadata.
+ + self.counter_map.capacity()
+ * (size_of::<(Vec<u8>, usize)>() + size_of::<u8>())
+ + self.alpha_map.capacity() * size_of::<u64>()
+ }
+}
+
+// ---------------------------------------------------------------------------
+// ApproxTopK UDAF struct
+// ---------------------------------------------------------------------------
+
+/// Approximate top-k UDAF using the Filtered Space-Saving algorithm.
+#[user_doc(
+ doc_section(label = "Approximate Functions"),
+ description = "Returns the approximate most frequent (top-k) values and their counts using the Filtered Space-Saving algorithm. Note: for float columns, -0.0 and +0.0 are treated as distinct values, and different NaN representations are tracked separately.",
+ syntax_example = "approx_top_k(expression, k)",
+ sql_example = r#"```sql
+> SELECT approx_top_k(column_name, 3) FROM table_name;
++-------------------------------------------+
+| approx_top_k(column_name, 3) |
++-------------------------------------------+
+| [{value: foo, count: 3}, {value: bar, count: 2}, {value: baz, count: 1}] |
++-------------------------------------------+
+```"#,
+ standard_argument(name = "expression",),
+ argument(
+ name = "k",
+ description = "The number of top elements to return. Must be a literal integer between 1 and 10,000."
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ApproxTopK {
+ signature: Signature,
+}
+
+impl Default for ApproxTopK {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ApproxTopK {
+ pub fn new() -> Self {
+ // Supported value types for the first argument.
+ let value_types = &[
+ DataType::Utf8,
+ DataType::LargeUtf8,
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ DataType::Float32,
+ DataType::Float64,
+ ];
+ // k must be a literal integer; accept any integer type for convenience.
+ let k_types = &[
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ ];
+
+ let mut variants = Vec::with_capacity(value_types.len() * k_types.len());
+ for vt in value_types {
+ for kt in k_types {
+ variants.push(TypeSignature::Exact(vec![vt.clone(), kt.clone()]));
+ }
+ }
+
+ Self {
+ signature: Signature::one_of(variants, Volatility::Immutable),
+ }
+ }
+}
+
+impl AggregateUDFImpl for ApproxTopK {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "approx_top_k"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ let value_type = if !arg_types.is_empty() {
+ match &arg_types[0] {
+ // LargeUtf8 is narrowed to Utf8: the summary stores in-memory
+ // byte slices, not large column offsets, so i32 offsets suffice.
+ DataType::Utf8 | DataType::LargeUtf8 => DataType::Utf8,
+ DataType::Int8 => DataType::Int8,
+ DataType::Int16 => DataType::Int16,
+ DataType::Int32 => DataType::Int32,
+ DataType::Int64 => DataType::Int64,
+ DataType::UInt8 => DataType::UInt8,
+ DataType::UInt16 => DataType::UInt16,
+ DataType::UInt32 => DataType::UInt32,
+ DataType::UInt64 => DataType::UInt64,
+ DataType::Float32 => DataType::Float32,
+ DataType::Float64 => DataType::Float64,
+ _ => DataType::Utf8,
+ }
+ } else {
+ DataType::Utf8
+ };
+
+ let struct_fields = Fields::from(vec![
+ Field::new("value", value_type, true),
+ Field::new("count", DataType::UInt64, false),
+ ]);
+ Ok(DataType::List(Arc::new(Field::new(
+ "item",
+ DataType::Struct(struct_fields),
+ true,
+ ))))
+ }
+
+ fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
+ Ok(vec![
+ Arc::new(Field::new(
+ format_state_name(args.name, "summary"),
+ DataType::Binary,
+ true,
+ )),
+ Arc::new(Field::new(
+ format_state_name(args.name, "k"),
+ DataType::UInt64,
+ true,
+ )),
+ Arc::new(Field::new(
+ format_state_name(args.name, "data_type"),
+ DataType::Utf8,
+ true,
+ )),
+ ])
+ }
+
+ fn accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+ if args.exprs.len() < 2 {
+ return Err(datafusion_common::DataFusionError::Plan(
+ "approx_top_k requires two arguments: column and k".to_string(),
+ ));
+ }
+
+ let k_expr = &args.exprs[1];
+ let k = k_expr
+ .as_any()
+ .downcast_ref::<datafusion_physical_expr::expressions::Literal>()
+ .and_then(|lit| match lit.value() {
+ // Guard against negative values before casting to usize to
+ // avoid wrapping (e.g. -1i32 as usize → u64::MAX).
+ // Zero is allowed through and caught by the bounds check below.
+ ScalarValue::Int8(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int16(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int32(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int64(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::UInt8(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt16(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt32(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt64(Some(v)) => Some(*v as usize),
+ _ => None,
+ })
+ .ok_or_else(|| {
+ datafusion_common::DataFusionError::Plan(
+ "approx_top_k requires k to be a positive literal integer"
+ .to_string(),
+ )
+ })?;
+
+ if k == 0 || k > APPROX_TOP_K_MAX_K {
+ return Err(datafusion_common::DataFusionError::Plan(format!(
+ "approx_top_k requires k to be between 1 and {APPROX_TOP_K_MAX_K}, got {k}"
+ )));
+ }
+
+ Ok(Box::new(ApproxTopKAccumulator::new(k)))
Review Comment:
Right, the data type is already known at accumulator creation time, no need to sniff it from the first batch. Updated to use args.expr_fields[0].data_type().clone(). Also removed the now-unused new() constructor since accumulator() always goes through new_with_data_type. Thanks!
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r3123517113
##########
datafusion/functions-aggregate/src/approx_top_k.rs:
##########
@@ -0,0 +1,1421 @@
+// 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, 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};
+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; 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) -> 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>,
Review Comment:
`Vec<u8>`: this is quite wasteful to store in the map.
we can store the hash inline in the map using the `HashTable` API instead.
`HashTable<u64, usize>`
We can remove the hash from the `counters`
and look up the `item` using the `usize` pointer using `counter_map.find`.
(Look up the pattern in other places like aggregation / join ...)
This saves many allocations / expensive rehashing.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r3728969106
##########
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);
+
+ // If there are still more counters than requested_capacity (because
+ // target_capacity wasn't reached), partition out the top ones and
+ // fold the tail into the alpha map before serializing.
+ let k = self.requested_capacity;
+ if k > 0 && 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();
+ }
+ self.rebuild_counter_map();
+ }
+
+ let num_counters = self.counters.len();
+
+ // Pre-compute total size for a single allocation.
+ let counter_bytes: usize =
+ self.counters.iter().map(|c| 4 + c.item.len() + 16).sum();
+ let total = 16 + counter_bytes + 8 + self.alpha_map.len() * 8;
+ let mut bytes = Vec::with_capacity(total);
+
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(num_counters as u64).to_le_bytes());
+
+ for counter in &self.counters {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
Review Comment:
I retained alpha because merge now consumes it: max(alpha) supplies the omission correction for values absent from a full partition summary, and alpha buckets are added to preserve multi-level mergeability. The wire format avoids dense zero-filled maps by selecting sparse entries or a bitmap with packed values.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r3728953750
##########
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:
The mutation is intentional. Inserts may temporarily grow the working set to twice the requested capacity to amortize truncation. Serialization is the boundary where it must be reduced to the requested capacity, with discarded lower bounds folded into alpha.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "martin-g (via GitHub)" <gi...@apache.org>.
martin-g commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r2945320798
##########
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();
+ let return_size = std::cmp::min(sorted.len(), self.requested_capacity);
+ if return_size > 0 && return_size < sorted.len() {
+ // After select_nth_unstable_by, the top-k counters are in
+ // sorted[..return_size] but in arbitrary order. This is fine
+ // since deserialization doesn't depend on counter ordering.
+ sorted.select_nth_unstable_by(return_size - 1, |a, b| a.cmp_by_rank(b));
+ }
+ sorted.truncate(return_size);
+ sorted
+ };
+
+ let mut bytes = Vec::new();
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(counters_to_write.len() as u64).to_le_bytes());
+
+ for counter in counters_to_write {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
+ for &alpha in &self.alpha_map {
+ bytes.extend_from_slice(&alpha.to_le_bytes());
+ }
+
+ bytes
+ }
+
+ /// Deserialize a summary from bytes.
+ fn from_bytes(bytes: &[u8]) -> Result<Self> {
+ if bytes.len() < 16 {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: too short".to_string(),
+ ));
+ }
+
+ let requested_capacity =
+ u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
+ let num_counters = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+
+ let mut counters = Vec::with_capacity(num_counters);
+ let mut counter_map = HashMap::with_capacity(num_counters);
+ let mut offset = 16;
+
+ for idx in 0..num_counters {
+ if offset + 4 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated item length"
+ .to_string(),
+ ));
+ }
+ let item_len =
+ u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
+ as usize;
+ offset += 4;
+
+ if offset + item_len + 16 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated counter".to_string(),
+ ));
+ }
+
+ let item = bytes[offset..offset + item_len].to_vec();
+ offset += item_len;
+
+ let count = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let error = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let hash = Self::hash_item(&item);
+ counter_map.insert(item.clone(), idx);
+ counters.push(Counter {
+ item,
+ hash,
+ count,
+ error,
+ });
+ }
+
+ if offset + 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: missing alpha map size".to_string(),
+ ));
+ }
+ let alpha_map_size =
+ u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap()) as usize;
+ offset += 8;
+
+ if offset + alpha_map_size * 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated alpha map".to_string(),
+ ));
+ }
+
+ let mut alpha_map = Vec::with_capacity(alpha_map_size);
+ for _ in 0..alpha_map_size {
+ let alpha = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+ alpha_map.push(alpha);
+ }
+
+ let target_capacity = std::cmp::max(64, requested_capacity * 2);
+
+ Ok(Self {
+ counters,
+ counter_map,
+ alpha_map,
+ requested_capacity,
+ target_capacity,
+ })
+ }
+
+ /// Approximate size in bytes of this summary.
+ fn size(&self) -> usize {
+ size_of::<Self>()
+ + self
+ .counters
+ .iter()
+ .map(|c| size_of::<Counter>() + c.item.len())
+ .sum::<usize>()
+ // HashMap overhead: each bucket has a key-value pair plus control byte metadata.
+ + self.counter_map.capacity()
+ * (size_of::<(Vec<u8>, usize)>() + size_of::<u8>())
+ + self.alpha_map.capacity() * size_of::<u64>()
+ }
+}
+
+// ---------------------------------------------------------------------------
+// ApproxTopK UDAF struct
+// ---------------------------------------------------------------------------
+
+/// Approximate top-k UDAF using the Filtered Space-Saving algorithm.
+#[user_doc(
+ doc_section(label = "Approximate Functions"),
+ description = "Returns the approximate most frequent (top-k) values and their counts using the Filtered Space-Saving algorithm. Note: for float columns, -0.0 and +0.0 are treated as distinct values, and different NaN representations are tracked separately.",
+ syntax_example = "approx_top_k(expression, k)",
+ sql_example = r#"```sql
+> SELECT approx_top_k(column_name, 3) FROM table_name;
++-------------------------------------------+
+| approx_top_k(column_name, 3) |
++-------------------------------------------+
+| [{value: foo, count: 3}, {value: bar, count: 2}, {value: baz, count: 1}] |
++-------------------------------------------+
+```"#,
+ standard_argument(name = "expression",),
+ argument(
+ name = "k",
+ description = "The number of top elements to return. Must be a literal integer between 1 and 10,000."
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ApproxTopK {
+ signature: Signature,
+}
+
+impl Default for ApproxTopK {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ApproxTopK {
+ pub fn new() -> Self {
+ // Supported value types for the first argument.
+ let value_types = &[
+ DataType::Utf8,
+ DataType::LargeUtf8,
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ DataType::Float32,
+ DataType::Float64,
+ ];
+ // k must be a literal integer; accept any integer type for convenience.
+ let k_types = &[
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ ];
+
+ let mut variants = Vec::with_capacity(value_types.len() * k_types.len());
+ for vt in value_types {
+ for kt in k_types {
+ variants.push(TypeSignature::Exact(vec![vt.clone(), kt.clone()]));
+ }
+ }
+
+ Self {
+ signature: Signature::one_of(variants, Volatility::Immutable),
+ }
+ }
+}
+
+impl AggregateUDFImpl for ApproxTopK {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "approx_top_k"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ let value_type = if !arg_types.is_empty() {
+ match &arg_types[0] {
+ // LargeUtf8 is narrowed to Utf8: the summary stores in-memory
+ // byte slices, not large column offsets, so i32 offsets suffice.
+ DataType::Utf8 | DataType::LargeUtf8 => DataType::Utf8,
+ DataType::Int8 => DataType::Int8,
+ DataType::Int16 => DataType::Int16,
+ DataType::Int32 => DataType::Int32,
+ DataType::Int64 => DataType::Int64,
+ DataType::UInt8 => DataType::UInt8,
+ DataType::UInt16 => DataType::UInt16,
+ DataType::UInt32 => DataType::UInt32,
+ DataType::UInt64 => DataType::UInt64,
+ DataType::Float32 => DataType::Float32,
+ DataType::Float64 => DataType::Float64,
+ _ => DataType::Utf8,
+ }
+ } else {
+ DataType::Utf8
+ };
+
+ let struct_fields = Fields::from(vec![
+ Field::new("value", value_type, true),
+ Field::new("count", DataType::UInt64, false),
+ ]);
+ Ok(DataType::List(Arc::new(Field::new(
+ "item",
+ DataType::Struct(struct_fields),
+ true,
+ ))))
+ }
+
+ fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
+ Ok(vec![
+ Arc::new(Field::new(
+ format_state_name(args.name, "summary"),
+ DataType::Binary,
+ true,
+ )),
+ Arc::new(Field::new(
+ format_state_name(args.name, "k"),
+ DataType::UInt64,
+ true,
+ )),
+ Arc::new(Field::new(
+ format_state_name(args.name, "data_type"),
+ DataType::Utf8,
+ true,
+ )),
+ ])
+ }
+
+ fn accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+ if args.exprs.len() < 2 {
+ return Err(datafusion_common::DataFusionError::Plan(
+ "approx_top_k requires two arguments: column and k".to_string(),
+ ));
+ }
+
+ let k_expr = &args.exprs[1];
+ let k = k_expr
+ .as_any()
+ .downcast_ref::<datafusion_physical_expr::expressions::Literal>()
+ .and_then(|lit| match lit.value() {
+ // Guard against negative values before casting to usize to
+ // avoid wrapping (e.g. -1i32 as usize → u64::MAX).
+ // Zero is allowed through and caught by the bounds check below.
+ ScalarValue::Int8(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int16(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int32(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int64(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::UInt8(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt16(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt32(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt64(Some(v)) => Some(*v as usize),
+ _ => None,
+ })
+ .ok_or_else(|| {
+ datafusion_common::DataFusionError::Plan(
+ "approx_top_k requires k to be a positive literal integer"
+ .to_string(),
+ )
+ })?;
+
+ if k == 0 || k > APPROX_TOP_K_MAX_K {
+ return Err(datafusion_common::DataFusionError::Plan(format!(
+ "approx_top_k requires k to be between 1 and {APPROX_TOP_K_MAX_K}, got {k}"
+ )));
+ }
+
+ Ok(Box::new(ApproxTopKAccumulator::new(k)))
Review Comment:
```suggestion
let data_type = args.expr_fields[0].data_type;
Ok(Box::new(ApproxTopKAccumulator::new_with_data_type(k, Some(data_type))))
```
Would this be more correct ?
##########
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();
+ let return_size = std::cmp::min(sorted.len(), self.requested_capacity);
+ if return_size > 0 && return_size < sorted.len() {
+ // After select_nth_unstable_by, the top-k counters are in
+ // sorted[..return_size] but in arbitrary order. This is fine
+ // since deserialization doesn't depend on counter ordering.
+ sorted.select_nth_unstable_by(return_size - 1, |a, b| a.cmp_by_rank(b));
+ }
+ sorted.truncate(return_size);
+ sorted
+ };
+
+ let mut bytes = Vec::new();
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(counters_to_write.len() as u64).to_le_bytes());
+
+ for counter in counters_to_write {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
+ for &alpha in &self.alpha_map {
+ bytes.extend_from_slice(&alpha.to_le_bytes());
+ }
+
+ bytes
+ }
+
+ /// Deserialize a summary from bytes.
+ fn from_bytes(bytes: &[u8]) -> Result<Self> {
+ if bytes.len() < 16 {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: too short".to_string(),
+ ));
+ }
+
+ let requested_capacity =
+ u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
+ let num_counters = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+
+ let mut counters = Vec::with_capacity(num_counters);
+ let mut counter_map = HashMap::with_capacity(num_counters);
+ let mut offset = 16;
+
+ for idx in 0..num_counters {
+ if offset + 4 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated item length"
+ .to_string(),
+ ));
+ }
+ let item_len =
+ u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
+ as usize;
+ offset += 4;
+
+ if offset + item_len + 16 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated counter".to_string(),
+ ));
+ }
+
+ let item = bytes[offset..offset + item_len].to_vec();
+ offset += item_len;
+
+ let count = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let error = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let hash = Self::hash_item(&item);
+ counter_map.insert(item.clone(), idx);
+ counters.push(Counter {
+ item,
+ hash,
+ count,
+ error,
+ });
+ }
+
+ if offset + 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: missing alpha map size".to_string(),
+ ));
+ }
+ let alpha_map_size =
+ u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap()) as usize;
+ offset += 8;
+
+ if offset + alpha_map_size * 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated alpha map".to_string(),
+ ));
+ }
+
+ let mut alpha_map = Vec::with_capacity(alpha_map_size);
+ for _ in 0..alpha_map_size {
+ let alpha = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+ alpha_map.push(alpha);
+ }
+
+ let target_capacity = std::cmp::max(64, requested_capacity * 2);
+
+ Ok(Self {
+ counters,
+ counter_map,
+ alpha_map,
+ requested_capacity,
+ target_capacity,
+ })
+ }
+
+ /// Approximate size in bytes of this summary.
+ fn size(&self) -> usize {
+ size_of::<Self>()
+ + self
+ .counters
+ .iter()
+ .map(|c| size_of::<Counter>() + c.item.len())
+ .sum::<usize>()
+ // HashMap overhead: each bucket has a key-value pair plus control byte metadata.
+ + self.counter_map.capacity()
+ * (size_of::<(Vec<u8>, usize)>() + size_of::<u8>())
+ + self.alpha_map.capacity() * size_of::<u64>()
+ }
+}
+
+// ---------------------------------------------------------------------------
+// ApproxTopK UDAF struct
+// ---------------------------------------------------------------------------
+
+/// Approximate top-k UDAF using the Filtered Space-Saving algorithm.
+#[user_doc(
+ doc_section(label = "Approximate Functions"),
+ description = "Returns the approximate most frequent (top-k) values and their counts using the Filtered Space-Saving algorithm. Note: for float columns, -0.0 and +0.0 are treated as distinct values, and different NaN representations are tracked separately.",
+ syntax_example = "approx_top_k(expression, k)",
+ sql_example = r#"```sql
+> SELECT approx_top_k(column_name, 3) FROM table_name;
++-------------------------------------------+
+| approx_top_k(column_name, 3) |
++-------------------------------------------+
+| [{value: foo, count: 3}, {value: bar, count: 2}, {value: baz, count: 1}] |
++-------------------------------------------+
+```"#,
+ standard_argument(name = "expression",),
+ argument(
+ name = "k",
+ description = "The number of top elements to return. Must be a literal integer between 1 and 10,000."
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ApproxTopK {
+ signature: Signature,
+}
+
+impl Default for ApproxTopK {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ApproxTopK {
+ pub fn new() -> Self {
+ // Supported value types for the first argument.
+ let value_types = &[
+ DataType::Utf8,
+ DataType::LargeUtf8,
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ DataType::Float32,
+ DataType::Float64,
+ ];
+ // k must be a literal integer; accept any integer type for convenience.
+ let k_types = &[
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ ];
+
+ let mut variants = Vec::with_capacity(value_types.len() * k_types.len());
+ for vt in value_types {
+ for kt in k_types {
+ variants.push(TypeSignature::Exact(vec![vt.clone(), kt.clone()]));
+ }
+ }
+
+ Self {
+ signature: Signature::one_of(variants, Volatility::Immutable),
+ }
+ }
+}
+
+impl AggregateUDFImpl for ApproxTopK {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "approx_top_k"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ let value_type = if !arg_types.is_empty() {
+ match &arg_types[0] {
+ // LargeUtf8 is narrowed to Utf8: the summary stores in-memory
+ // byte slices, not large column offsets, so i32 offsets suffice.
+ DataType::Utf8 | DataType::LargeUtf8 => DataType::Utf8,
+ DataType::Int8 => DataType::Int8,
+ DataType::Int16 => DataType::Int16,
+ DataType::Int32 => DataType::Int32,
+ DataType::Int64 => DataType::Int64,
+ DataType::UInt8 => DataType::UInt8,
+ DataType::UInt16 => DataType::UInt16,
+ DataType::UInt32 => DataType::UInt32,
+ DataType::UInt64 => DataType::UInt64,
+ DataType::Float32 => DataType::Float32,
+ DataType::Float64 => DataType::Float64,
+ _ => DataType::Utf8,
+ }
+ } else {
+ DataType::Utf8
+ };
+
+ let struct_fields = Fields::from(vec![
+ Field::new("value", value_type, true),
+ Field::new("count", DataType::UInt64, false),
+ ]);
+ Ok(DataType::List(Arc::new(Field::new(
+ "item",
+ DataType::Struct(struct_fields),
+ true,
+ ))))
+ }
+
+ fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
+ Ok(vec![
+ Arc::new(Field::new(
+ format_state_name(args.name, "summary"),
+ DataType::Binary,
+ true,
+ )),
+ Arc::new(Field::new(
+ format_state_name(args.name, "k"),
+ DataType::UInt64,
+ true,
+ )),
+ Arc::new(Field::new(
+ format_state_name(args.name, "data_type"),
+ DataType::Utf8,
+ true,
+ )),
+ ])
+ }
+
+ fn accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+ if args.exprs.len() < 2 {
+ return Err(datafusion_common::DataFusionError::Plan(
+ "approx_top_k requires two arguments: column and k".to_string(),
+ ));
+ }
+
+ let k_expr = &args.exprs[1];
+ let k = k_expr
+ .as_any()
+ .downcast_ref::<datafusion_physical_expr::expressions::Literal>()
+ .and_then(|lit| match lit.value() {
+ // Guard against negative values before casting to usize to
+ // avoid wrapping (e.g. -1i32 as usize → u64::MAX).
+ // Zero is allowed through and caught by the bounds check below.
+ ScalarValue::Int8(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int16(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int32(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::Int64(Some(v)) if *v >= 0 => Some(*v as usize),
+ ScalarValue::UInt8(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt16(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt32(Some(v)) => Some(*v as usize),
+ ScalarValue::UInt64(Some(v)) => Some(*v as usize),
+ _ => None,
+ })
+ .ok_or_else(|| {
+ datafusion_common::DataFusionError::Plan(
+ "approx_top_k requires k to be a positive literal integer"
+ .to_string(),
+ )
+ })?;
+
+ if k == 0 || k > APPROX_TOP_K_MAX_K {
+ return Err(datafusion_common::DataFusionError::Plan(format!(
+ "approx_top_k requires k to be between 1 and {APPROX_TOP_K_MAX_K}, got {k}"
+ )));
+ }
+
+ Ok(Box::new(ApproxTopKAccumulator::new(k)))
+ }
+
+ fn documentation(&self) -> Option<&Documentation> {
+ self.doc()
+ }
+}
+
+// ---------------------------------------------------------------------------
+// Accumulator
+// ---------------------------------------------------------------------------
+
+/// Accumulator for `approx_top_k` using the Filtered Space-Saving algorithm.
+#[derive(Debug)]
+struct ApproxTopKAccumulator {
+ summary: SpaceSavingSummary,
+ k: usize,
+ /// The data type of the input column. `None` until the first batch is seen.
+ input_data_type: Option<DataType>,
+}
+
+impl ApproxTopKAccumulator {
+ fn new(k: usize) -> Self {
+ Self::new_with_data_type(k, None)
+ }
+
+ fn new_with_data_type(k: usize, input_data_type: Option<DataType>) -> Self {
+ let capacity = k * CAPACITY_MULTIPLIER;
+ Self {
+ summary: SpaceSavingSummary::new(capacity),
+ k,
+ input_data_type,
+ }
+ }
+
+ /// Build the value array for the result based on the input data type.
+ fn build_value_array(&self, top_items: &[(&[u8], u64, u64)]) -> Result<ArrayRef> {
+ match self.input_data_type.as_ref().unwrap_or(&DataType::Utf8) {
+ DataType::Utf8 | DataType::LargeUtf8 => {
+ let values: Vec<Option<String>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| String::from_utf8(bytes.to_vec()).ok())
+ .collect();
+ Ok(Arc::new(StringArray::from(values)) as ArrayRef)
+ }
+ DataType::Int8 => {
+ let values: Vec<Option<i8>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 1]>::try_from(*bytes).ok().map(i8::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(Int8Array::from(values)) as ArrayRef)
+ }
+ DataType::Int16 => {
+ let values: Vec<Option<i16>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 2]>::try_from(*bytes).ok().map(i16::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(Int16Array::from(values)) as ArrayRef)
+ }
+ DataType::Int32 => {
+ let values: Vec<Option<i32>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 4]>::try_from(*bytes).ok().map(i32::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(Int32Array::from(values)) as ArrayRef)
+ }
+ DataType::Int64 => {
+ let values: Vec<Option<i64>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 8]>::try_from(*bytes).ok().map(i64::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(Int64Array::from(values)) as ArrayRef)
+ }
+ DataType::UInt8 => {
+ let values: Vec<Option<u8>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 1]>::try_from(*bytes).ok().map(u8::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(UInt8Array::from(values)) as ArrayRef)
+ }
+ DataType::UInt16 => {
+ let values: Vec<Option<u16>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 2]>::try_from(*bytes).ok().map(u16::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(UInt16Array::from(values)) as ArrayRef)
+ }
+ DataType::UInt32 => {
+ let values: Vec<Option<u32>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 4]>::try_from(*bytes).ok().map(u32::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(UInt32Array::from(values)) as ArrayRef)
+ }
+ DataType::UInt64 => {
+ let values: Vec<Option<u64>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 8]>::try_from(*bytes).ok().map(u64::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(UInt64Array::from(values)) as ArrayRef)
+ }
+ DataType::Float32 => {
+ let values: Vec<Option<f32>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 4]>::try_from(*bytes).ok().map(f32::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(Float32Array::from(values)) as ArrayRef)
+ }
+ DataType::Float64 => {
+ let values: Vec<Option<f64>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| {
+ <[u8; 8]>::try_from(*bytes).ok().map(f64::from_le_bytes)
+ })
+ .collect();
+ Ok(Arc::new(Float64Array::from(values)) as ArrayRef)
+ }
+ _ => {
+ let values: Vec<Option<String>> = top_items
+ .iter()
+ .map(|(bytes, _, _)| String::from_utf8(bytes.to_vec()).ok())
+ .collect();
+ Ok(Arc::new(StringArray::from(values)) as ArrayRef)
+ }
+ }
+ }
+
+ /// Get the output data type for the value field.
+ /// Falls back to `Utf8` if the input data type was never captured (empty input).
+ fn output_value_data_type(&self) -> DataType {
+ match self.input_data_type.as_ref().unwrap_or(&DataType::Utf8) {
+ DataType::Utf8 | DataType::LargeUtf8 => DataType::Utf8,
+ DataType::Int8 => DataType::Int8,
+ DataType::Int16 => DataType::Int16,
+ DataType::Int32 => DataType::Int32,
+ DataType::Int64 => DataType::Int64,
+ DataType::UInt8 => DataType::UInt8,
+ DataType::UInt16 => DataType::UInt16,
+ DataType::UInt32 => DataType::UInt32,
+ DataType::UInt64 => DataType::UInt64,
+ DataType::Float32 => DataType::Float32,
+ DataType::Float64 => DataType::Float64,
+ _ => DataType::Utf8,
+ }
+ }
+
+ fn data_type_to_string(dt: &DataType) -> String {
+ match dt {
+ DataType::Utf8 => "utf8".to_string(),
+ DataType::LargeUtf8 => "large_utf8".to_string(),
+ DataType::Int8 => "int8".to_string(),
+ DataType::Int16 => "int16".to_string(),
+ DataType::Int32 => "int32".to_string(),
+ DataType::Int64 => "int64".to_string(),
+ DataType::UInt8 => "uint8".to_string(),
+ DataType::UInt16 => "uint16".to_string(),
+ DataType::UInt32 => "uint32".to_string(),
+ DataType::UInt64 => "uint64".to_string(),
+ DataType::Float32 => "float32".to_string(),
+ DataType::Float64 => "float64".to_string(),
+ _ => "utf8".to_string(),
+ }
+ }
+
+ fn string_to_data_type(s: &str) -> DataType {
+ match s {
+ "utf8" => DataType::Utf8,
+ "large_utf8" => DataType::LargeUtf8,
+ "int8" => DataType::Int8,
+ "int16" => DataType::Int16,
+ "int32" => DataType::Int32,
+ "int64" => DataType::Int64,
+ "uint8" => DataType::UInt8,
+ "uint16" => DataType::UInt16,
+ "uint32" => DataType::UInt32,
+ "uint64" => DataType::UInt64,
+ "float32" => DataType::Float32,
+ "float64" => DataType::Float64,
+ _ => DataType::Utf8,
+ }
+ }
+}
+
+impl Accumulator for ApproxTopKAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
+ if values.is_empty() {
+ return Ok(());
+ }
+
+ let data_array = &values[0];
+
+ // Capture the input data type from the first batch if not already set.
+ if self.input_data_type.is_none() {
+ self.input_data_type = Some(data_array.data_type().clone());
+ }
+
+ // Downcast once and iterate directly to avoid per-row ScalarValue allocation.
+ macro_rules! process_array {
+ ($array_type:ty, $data_array:expr) => {{
+ let arr = $data_array.as_any().downcast_ref::<$array_type>().unwrap();
+ for i in 0..arr.len() {
+ if !arr.is_null(i) {
+ self.summary.add(&arr.value(i).to_le_bytes());
+ }
+ }
+ }};
+ }
+
+ match data_array.data_type() {
+ DataType::Utf8 => {
+ let arr = data_array.as_any().downcast_ref::<StringArray>().unwrap();
+ for i in 0..arr.len() {
+ if !arr.is_null(i) {
+ self.summary.add(arr.value(i).as_bytes());
+ }
+ }
+ }
+ DataType::LargeUtf8 => {
+ let arr = data_array
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .unwrap();
+ for i in 0..arr.len() {
+ if !arr.is_null(i) {
+ self.summary.add(arr.value(i).as_bytes());
+ }
+ }
+ }
+ DataType::Int8 => process_array!(Int8Array, data_array),
+ DataType::Int16 => process_array!(Int16Array, data_array),
+ DataType::Int32 => process_array!(Int32Array, data_array),
+ DataType::Int64 => process_array!(Int64Array, data_array),
+ DataType::UInt8 => process_array!(UInt8Array, data_array),
+ DataType::UInt16 => process_array!(UInt16Array, data_array),
+ DataType::UInt32 => process_array!(UInt32Array, data_array),
+ DataType::UInt64 => process_array!(UInt64Array, data_array),
+ // Note: floats are compared by their byte representation, so -0.0 and +0.0
+ // are treated as distinct values, and different NaN bit patterns are tracked
+ // separately. This matches ClickHouse's behavior for topK with floats.
+ DataType::Float32 => process_array!(Float32Array, data_array),
+ DataType::Float64 => process_array!(Float64Array, data_array),
+ other => {
+ return Err(datafusion_common::DataFusionError::Execution(format!(
+ "Unsupported data type for approx_top_k: {other}"
+ )));
+ }
+ }
+
+ Ok(())
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
+ // State layout: [summary (Binary), k (UInt64), data_type (Utf8)].
+ // The `k` field (states[1]) is carried for completeness but not read here
+ // because the planner guarantees all partial accumulators use the same `k`.
+ if states.is_empty() || states[0].is_empty() {
+ return Ok(());
+ }
+
+ let summary_array = states[0]
+ .as_any()
+ .downcast_ref::<BinaryArray>()
+ .ok_or_else(|| {
+ datafusion_common::DataFusionError::Execution(
+ "Expected Binary array for summary state".to_string(),
+ )
+ })?;
+
+ // Read the data type from state if available (third state field).
+ if states.len() > 2
+ && let Some(data_type_array) =
+ states[2].as_any().downcast_ref::<StringArray>()
+ {
+ for i in 0..data_type_array.len() {
+ if !data_type_array.is_null(i) {
+ let data_type_str = data_type_array.value(i);
+ let data_type = Self::string_to_data_type(data_type_str);
+ if self.input_data_type.is_none() {
+ self.input_data_type = Some(data_type);
+ }
+ break;
+ }
+ }
+ }
+
+ for i in 0..summary_array.len() {
+ if summary_array.is_null(i) {
+ continue;
+ }
+ let bytes = summary_array.value(i);
+ let other_summary = SpaceSavingSummary::from_bytes(bytes)?;
+ self.summary.merge(&other_summary);
+ }
+
+ Ok(())
+ }
+
+ fn evaluate(&mut self) -> Result<ScalarValue> {
+ let top_items = self.summary.top_k(self.k);
+
+ let value_data_type = self.output_value_data_type();
+ let struct_fields = Fields::from(vec![
+ Field::new("value", value_data_type, true),
+ Field::new("count", DataType::UInt64, false),
+ ]);
+
+ let value_array = self.build_value_array(&top_items)?;
+ let counts: Vec<u64> = top_items.iter().map(|(_, count, _)| *count).collect();
+ let count_array = Arc::new(UInt64Array::from(counts)) as ArrayRef;
+
+ let struct_array =
+ StructArray::new(struct_fields.clone(), vec![value_array, count_array], None);
+
+ let list_field = Field::new("item", DataType::Struct(struct_fields), true);
+
+ Ok(ScalarValue::List(Arc::new(ListArray::new(
+ Arc::new(list_field),
+ OffsetBuffer::from_lengths([top_items.len()]),
+ Arc::new(struct_array),
+ None,
+ ))))
+ }
+
+ fn state(&mut self) -> Result<Vec<ScalarValue>> {
+ let summary_bytes = self.summary.to_bytes();
+ let data_type_str = Self::data_type_to_string(
+ self.input_data_type.as_ref().unwrap_or(&DataType::Utf8),
+ );
+
+ Ok(vec![
+ ScalarValue::Binary(Some(summary_bytes)),
+ ScalarValue::UInt64(Some(self.k as u64)),
+ ScalarValue::Utf8(Some(data_type_str)),
Review Comment:
```suggestion
match self.input_data_type.as_ref() {
Some(dt) => ScalarValue::Utf8(Some(Self::data_type_to_string(dt))),
None => ScalarValue::Utf8(None),
},
```
i.e. return None for NULL
##########
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:
Shouldn't `self.truncate_if_needed();` be called before using the counters ?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#issuecomment-4154017620
Thanks for the review @Dandandan @martin-g , I believe I've addressed all of your feedback.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#issuecomment-5044575051
> @Dandandan @martin-g when should I expect the next round of reviews on this? thanks!
ping
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r3125082094
##########
datafusion/functions-aggregate/src/approx_top_k.rs:
##########
@@ -0,0 +1,1421 @@
+// 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, 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};
+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; 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) -> 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>,
Review Comment:
Good call! switched to HashTable<(u64, usize)> following the same pattern as group_values/row.rs and join_hash_map.rs. No key bytes cloned anymore, rebuild is zero-allocation, and double-hashing is gone.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#issuecomment-6116437411
Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "jonathanc-n (via GitHub)" <gi...@apache.org>.
jonathanc-n commented on PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#issuecomment-4111325817
cc @Dandandan @martin-g Should be good for another look
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r2987845664
##########
datafusion/functions-aggregate/src/approx_top_k.rs:
##########
@@ -0,0 +1,1381 @@
+// 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 serialize(&mut self) -> Vec<u8> {
+ // Ensure counters are truncated and alpha map is up to date before
+ // serializing, in case to_bytes is called without a prior truncation.
+ self.truncate_if_needed(false);
+
+ let counters_to_write: Vec<_> = {
+ let mut sorted: Vec<_> = self.counters.iter().collect();
+ let return_size = std::cmp::min(sorted.len(), self.requested_capacity);
+ if return_size > 0 && return_size < sorted.len() {
+ // After select_nth_unstable_by, the top-k counters are in
+ // sorted[..return_size] but in arbitrary order. This is fine
+ // since deserialization doesn't depend on counter ordering.
+ sorted.select_nth_unstable_by(return_size - 1, |a, b| a.cmp_by_rank(b));
+ }
+ sorted.truncate(return_size);
+ sorted
+ };
+
+ let mut bytes = Vec::new();
+ bytes.extend_from_slice(&(self.requested_capacity as u64).to_le_bytes());
+ bytes.extend_from_slice(&(counters_to_write.len() as u64).to_le_bytes());
+
+ for counter in counters_to_write {
+ bytes.extend_from_slice(&(counter.item.len() as u32).to_le_bytes());
+ bytes.extend_from_slice(&counter.item);
+ bytes.extend_from_slice(&counter.count.to_le_bytes());
+ bytes.extend_from_slice(&counter.error.to_le_bytes());
+ }
+
+ bytes.extend_from_slice(&(self.alpha_map.len() as u64).to_le_bytes());
+ for &alpha in &self.alpha_map {
+ bytes.extend_from_slice(&alpha.to_le_bytes());
+ }
+
+ bytes
+ }
+
+ /// Deserialize a summary from bytes.
+ fn from_bytes(bytes: &[u8]) -> Result<Self> {
+ if bytes.len() < 16 {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: too short".to_string(),
+ ));
+ }
+
+ let requested_capacity =
+ u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
+ let num_counters = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+
+ let mut counters = Vec::with_capacity(num_counters);
+ let mut counter_map = HashMap::with_capacity(num_counters);
+ let mut offset = 16;
+
+ for idx in 0..num_counters {
+ if offset + 4 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated item length"
+ .to_string(),
+ ));
+ }
+ let item_len =
+ u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
+ as usize;
+ offset += 4;
+
+ if offset + item_len + 16 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated counter".to_string(),
+ ));
+ }
+
+ let item = bytes[offset..offset + item_len].to_vec();
+ offset += item_len;
+
+ let count = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let error = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+
+ let hash = Self::hash_item(&item);
+ counter_map.insert(item.clone(), idx);
+ counters.push(Counter {
+ item,
+ hash,
+ count,
+ error,
+ });
+ }
+
+ if offset + 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: missing alpha map size".to_string(),
+ ));
+ }
+ let alpha_map_size =
+ u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap()) as usize;
+ offset += 8;
+
+ if offset + alpha_map_size * 8 > bytes.len() {
+ return Err(datafusion_common::DataFusionError::Execution(
+ "Invalid Space-Saving summary bytes: truncated alpha map".to_string(),
+ ));
+ }
+
+ let mut alpha_map = Vec::with_capacity(alpha_map_size);
+ for _ in 0..alpha_map_size {
+ let alpha = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
+ offset += 8;
+ alpha_map.push(alpha);
+ }
+
+ let target_capacity = std::cmp::max(64, requested_capacity * 2);
+
+ Ok(Self {
+ counters,
+ counter_map,
+ alpha_map,
+ requested_capacity,
+ target_capacity,
+ })
+ }
+
+ /// Approximate size in bytes of this summary.
+ fn size(&self) -> usize {
+ size_of::<Self>()
+ + self
+ .counters
+ .iter()
+ .map(|c| size_of::<Counter>() + c.item.len())
+ .sum::<usize>()
+ // HashMap overhead: each bucket has a key-value pair plus control byte metadata.
+ + self.counter_map.capacity()
+ * (size_of::<(Vec<u8>, usize)>() + size_of::<u8>())
+ + self.alpha_map.capacity() * size_of::<u64>()
+ }
+}
+
+// ---------------------------------------------------------------------------
+// ApproxTopK UDAF struct
+// ---------------------------------------------------------------------------
+
+/// Approximate top-k UDAF using the Filtered Space-Saving algorithm.
+#[user_doc(
+ doc_section(label = "Approximate Functions"),
+ description = "Returns the approximate most frequent (top-k) values and their counts using the Filtered Space-Saving algorithm. Note: for float columns, -0.0 and +0.0 are treated as distinct values, and different NaN representations are tracked separately.",
+ syntax_example = "approx_top_k(expression, k)",
+ sql_example = r#"```sql
+> SELECT approx_top_k(column_name, 3) FROM table_name;
++-------------------------------------------+
+| approx_top_k(column_name, 3) |
++-------------------------------------------+
+| [{value: foo, count: 3}, {value: bar, count: 2}, {value: baz, count: 1}] |
++-------------------------------------------+
+```"#,
+ standard_argument(name = "expression",),
+ argument(
+ name = "k",
+ description = "The number of top elements to return. Must be a literal integer between 1 and 10,000."
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ApproxTopK {
+ signature: Signature,
+}
+
+impl Default for ApproxTopK {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ApproxTopK {
+ pub fn new() -> Self {
+ // Supported value types for the first argument.
+ let value_types = &[
+ DataType::Utf8,
+ DataType::LargeUtf8,
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ DataType::UInt8,
+ DataType::UInt16,
+ DataType::UInt32,
+ DataType::UInt64,
+ DataType::Float32,
+ DataType::Float64,
Review Comment:
Good point - I've added a few more types.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#discussion_r3123517113
##########
datafusion/functions-aggregate/src/approx_top_k.rs:
##########
@@ -0,0 +1,1421 @@
+// 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, 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};
+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; 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) -> 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>,
Review Comment:
`Vec<u8>`: this is quite wasteful to store in the map.
we can store the hash inline in the map using the `HashTable` API instead.
`HashTable<u64, usize>`
We can remove the hash from the `counters`
and look up the `item` using the `usize` pointer using `counter_map.find`.
(Look up the pattern in other places like aggregation / join ...)
This saves many allocations / expensive rehashing / memory usage (for larger k, lot's of replacements ...)
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "sesteves (via GitHub)" <gi...@apache.org>.
sesteves commented on PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#issuecomment-4386441264
> @sesteves
>
> starts to look good!
>
> Can you also remove the testing submodule commit / update? And fix remaining CI issues?
Done, please lmk if there are still any remaining issues @Dandandan
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add approx_top_k aggregate function [datafusion]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #20968:
URL: https://github.com/apache/datafusion/pull/20968#issuecomment-5041059104
Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org