andygrove commented on code in PR #4801: URL: https://github.com/apache/datafusion-comet/pull/4801#discussion_r3552884824
########## native/spark-expr/src/agg_funcs/quantile_summaries.rs: ########## @@ -0,0 +1,440 @@ +// 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. + +//! A bit-for-bit port of Spark's `QuantileSummaries` (Greenwald-Khanna), the +//! sketch behind `approx_percentile` / `percentile_approx`. Kept free of any +//! DataFusion dependency so it can be unit-tested in isolation. +//! +//! Reference: `org.apache.spark.sql.catalyst.util.QuantileSummaries`. + +use std::collections::VecDeque; + +/// A single sampled statistic: the value, its minimum rank jump `g`, and the +/// maximum span of the rank `delta`. +#[derive(Debug, Clone, PartialEq)] +pub struct Stats { + pub value: f64, + pub g: i64, + pub delta: i64, +} + +#[derive(Debug, Clone)] +pub struct QuantileSummaries { + compress_threshold: usize, + relative_error: f64, + sampled: Vec<Stats>, + count: i64, + compressed: bool, + head_sampled: Vec<f64>, +} + +impl QuantileSummaries { + pub const DEFAULT_COMPRESS_THRESHOLD: usize = 10000; + pub const DEFAULT_HEAD_SIZE: usize = 50000; + + /// Mirrors Spark's `PercentileDigest` ctor which builds a summary with + /// `compressed = true`. + pub fn new(compress_threshold: usize, relative_error: f64) -> Self { + Self { + compress_threshold, + relative_error, + sampled: Vec::new(), + count: 0, + compressed: true, + head_sampled: Vec::new(), + } + } + + pub fn count(&self) -> i64 { + self.count + } + + /// Heap bytes held by this summary (excluding the struct itself). + pub fn heap_size(&self) -> usize { + self.sampled.capacity() * std::mem::size_of::<Stats>() + + self.head_sampled.capacity() * std::mem::size_of::<f64>() + } + + pub fn insert(&mut self, x: f64) { + self.head_sampled.push(x); + self.compressed = false; + if self.head_sampled.len() >= Self::DEFAULT_HEAD_SIZE { + self.with_head_buffer_inserted(); + if self.sampled.len() >= self.compress_threshold { + self.compress(); + } + } + } + + fn with_head_buffer_inserted(&mut self) { + if self.head_sampled.is_empty() { + return; + } + let mut current_count = self.count; + let mut sorted = std::mem::take(&mut self.head_sampled); + // Spark relies on `Array[Double].sorted`; use total ordering so the port + // is deterministic even in the presence of NaN (Spark's typical inputs + // are NaN-free). + sorted.sort_by(|a, b| a.total_cmp(b)); + + let mut new_samples: Vec<Stats> = Vec::new(); + let mut sample_idx = 0usize; + let mut ops_idx = 0usize; + while ops_idx < sorted.len() { + let current_sample = sorted[ops_idx]; + while sample_idx < self.sampled.len() + && self.sampled[sample_idx].value <= current_sample + { + new_samples.push(self.sampled[sample_idx].clone()); + sample_idx += 1; + } + current_count += 1; + let delta = if new_samples.is_empty() + || (sample_idx == self.sampled.len() && ops_idx == sorted.len() - 1) + { + 0 + } else { + // Spark uses `.toInt` here (unlike `.toLong` in merge/compress); we use i64 + // uniformly. They diverge only past ~Int.MAX rows in one accumulator, which + // is not reachable in practice. + (2.0 * self.relative_error * current_count as f64).floor() as i64 + }; + new_samples.push(Stats { + value: current_sample, + g: 1, + delta, + }); + ops_idx += 1; + } + while sample_idx < self.sampled.len() { + new_samples.push(self.sampled[sample_idx].clone()); + sample_idx += 1; + } + self.sampled = new_samples; + self.count = current_count; + } + + pub fn compress(&mut self) { + // Already compressed and the head buffer is empty (insert clears the + // flag whenever it stages a value), so there is nothing to do. This + // mirrors Spark's `PercentileDigest.isCompressed` guard, which also + // compresses at most once. + if self.compressed { + return; + } + self.with_head_buffer_inserted(); + let merge_threshold = 2.0 * self.relative_error * self.count as f64; + self.sampled = Self::compress_immut(&self.sampled, merge_threshold); + self.compressed = true; + } + + fn compress_immut(current_samples: &[Stats], merge_threshold: f64) -> Vec<Stats> { + if current_samples.is_empty() { + return Vec::new(); + } + let mut res: VecDeque<Stats> = VecDeque::new(); + let mut head = current_samples[current_samples.len() - 1].clone(); + // Traverse backward from size-2 down to index 1 (index 0 is preserved + // separately so the minimum is always kept). + let mut i = current_samples.len() as isize - 2; + while i >= 1 { + let sample1 = ¤t_samples[i as usize]; + if ((sample1.g + head.g + head.delta) as f64) < merge_threshold { + head.g += sample1.g; + } else { + res.push_front(head.clone()); + head = sample1.clone(); + } + i -= 1; + } + res.push_front(head.clone()); + let curr_head = ¤t_samples[0]; + if curr_head.value <= head.value && current_samples.len() > 1 { + res.push_front(curr_head.clone()); + } + res.into() + } + + pub fn merge(&self, other: &QuantileSummaries) -> QuantileSummaries { Review Comment: Filed as #4874 to track the consuming-merge and flush double-buffer optimization. -- 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]
