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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
