mbutrovich commented on code in PR #4801: URL: https://github.com/apache/datafusion-comet/pull/4801#discussion_r3552516479
########## 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)); Review Comment: `sort_unstable_by(|a, b| a.total_cmp(b))` would be faster with identical output here. Stability is irrelevant because the sort key is the value itself, so equal keys are bit-equal elements, and `total_cmp` orders `-0.0` / `0.0` / NaN deterministically regardless of stability. Separately: since `float` and `double` are marked `Compatible` with no caveat, is the NaN and `-0.0` ordering worth a test? `total_cmp` and Scala's `Array[Double].sorted` can differ across Scala versions, and Comet still cross-builds 2.12. ########## 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 Review Comment: The comment says Spark uses `.toInt` here and diverges from Comet past `Int.MAX` rows. In the Spark source I checked (`QuantileSummaries.scala:112`, 4.0.0) this line is `math.floor(2 * relativeError * currentCount).toLong`, so Spark uses `.toLong` too and the divergence does not exist. The `i64` choice is correct, but the comment reads as if the two engines differ when they match. Worth correcting or dropping. Could you confirm 3.4 / 3.5 / 4.1 also use `.toLong` (this file rarely changes)? ########## 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 Review Comment: The module header calls this a bit-for-bit port, which discourages any change. It might help to state instead what is actually load-bearing for bit-identity: the 50000-value flush boundary, the integer arithmetic (`floor(2 * relative_error * count)` as `i64`, the `/2` for `target_error`, `ceil(percentile * count)`), the backward compress traversal with the `< merge_threshold` test, the merge interleave and delta adjustment, and the query ceil-rank sweep. The data-structure choices (deque vs vec, clone vs copy, immutable rebuild vs in-place) do not affect results and are free to be optimized. That contract lets the code be idiomatic without risking correctness. ########## 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)] Review Comment: `Stats` is `f64 + i64 + i64`, all `Copy`, so deriving `Copy` would let you drop the roughly eight `.clone()` calls in `with_head_buffer_inserted`, `compress_immut`, and `merge`. No result change, more idiomatic. ########## 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(); Review Comment: `compress_immut` uses `VecDeque` with `push_front`, mirroring Scala's `ListBuffer.prepend`. Idiomatic Rust builds forward into a `Vec` and calls `reverse()` once. The output sequence is identical. ########## spark/src/test/resources/sql-tests/expressions/aggregate/approx_percentile.sql: ########## @@ -0,0 +1,68 @@ +-- 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. + +-- Native approx_percentile via a GK (Greenwald-Khanna) quantile summary port, +-- matching Spark's algorithm and default relative error, so results are +-- bit-identical. Only byte, short, int, long, float, and double inputs are +-- supported. Every query below uses the default `query` mode, which asserts +-- native execution, so an all-fallback run of this file cannot vacuously pass. + +-- scalar percentile over a bigint (long) column +query +SELECT approx_percentile(id, 0.5) FROM range(1000) Review Comment: Coverage is solid. A few paths look untested and would be quick to add: - `int` input (`cast(id AS int)`), the most common case (long, double, float, byte, short are covered). - Extreme percentiles `0.0` and `1.0`, which hit the `percentile <= relative_error` and `percentile >= 1 - relative_error` short-circuits in `query` that nothing else exercises. - Negative values (every case uses `range()`, which is non-negative). - A low-cardinality / duplicate-heavy column (`id % 5`), which exercises the `g` accumulation in `compress_immut`. - Array output with empty input (the `return_array` empty path in `evaluate` is untested at both SQL and Rust level). ########## native/spark-expr/src/agg_funcs/approx_percentile.rs: ########## @@ -0,0 +1,318 @@ +// 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. + +use super::quantile_summaries::QuantileSummaries; +use arrow::array::{Array, ArrayRef, BinaryArray, Float64Array, ListArray}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion::common::utils::SingleRowListArrayBuilder; +use datafusion::common::{downcast_value, Result, ScalarValue}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::Volatility::Immutable; +use datafusion::logical_expr::{Accumulator, AggregateUDFImpl, Signature}; +use datafusion::physical_expr::expressions::format_state_name; +use std::sync::Arc; + +/// Native implementation of Spark's `approx_percentile` / `percentile_approx`, +/// backed by a bit-for-bit `QuantileSummaries` (Greenwald-Khanna) port. The +/// child value is cast to Float64 by the serde; the original `input_type` is +/// carried so results can be cast back to Spark's output type. +#[derive(Debug)] +pub struct ApproxPercentile { + name: String, + signature: Signature, + percentiles: Vec<f64>, + accuracy: i64, + input_type: DataType, + return_array: bool, +} + +impl PartialEq for ApproxPercentile { + fn eq(&self, other: &Self) -> bool { + self.name == other.name + && self.percentiles == other.percentiles + && self.accuracy == other.accuracy + && self.input_type == other.input_type + && self.return_array == other.return_array + } +} +impl Eq for ApproxPercentile {} + +impl std::hash::Hash for ApproxPercentile { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + self.name.hash(state); + self.percentiles + .iter() + .for_each(|p| p.to_bits().hash(state)); + self.accuracy.hash(state); + self.input_type.hash(state); + self.return_array.hash(state); + } +} + +impl ApproxPercentile { + pub fn new( + percentiles: Vec<f64>, + accuracy: i64, + input_type: DataType, + return_array: bool, + ) -> Self { + Self { + name: "approx_percentile".to_string(), + signature: Signature::numeric(1, Immutable), + percentiles, + accuracy, + input_type, + return_array, + } + } +} + +impl AggregateUDFImpl for ApproxPercentile { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + if self.return_array { + Ok(DataType::List(Arc::new(Field::new( + "item", + self.input_type.clone(), + false, + )))) + } else { + Ok(self.input_type.clone()) + } + } + + fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + Ok(Box::new(ApproxPercentileAccumulator::new( + self.percentiles.clone(), + self.accuracy, + self.input_type.clone(), + self.return_array, + ))) + } + + fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + Ok(vec![Arc::new(Field::new( + format_state_name(&self.name, "digest"), + DataType::Binary, + true, + ))]) + } + + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + false + } +} + +#[derive(Debug)] +struct ApproxPercentileAccumulator { + summary: QuantileSummaries, + percentiles: Vec<f64>, + input_type: DataType, + return_array: bool, +} + +impl ApproxPercentileAccumulator { + fn new(percentiles: Vec<f64>, accuracy: i64, input_type: DataType, return_array: bool) -> Self { + let relative_error = 1.0 / accuracy as f64; + Self { + summary: QuantileSummaries::new( + QuantileSummaries::DEFAULT_COMPRESS_THRESHOLD, + relative_error, + ), + percentiles, + input_type, + return_array, + } + } + + /// Cast a double quantile back to Spark's output type. GK always returns an + /// actual inserted value (never an interpolation), so for the supported + /// numeric types this round-trips exactly and is always in range. + fn cast_back(&self, d: f64) -> ScalarValue { + match &self.input_type { + DataType::Int8 => ScalarValue::Int8(Some(d as i8)), + DataType::Int16 => ScalarValue::Int16(Some(d as i16)), + DataType::Int32 => ScalarValue::Int32(Some(d as i32)), + DataType::Int64 => ScalarValue::Int64(Some(d as i64)), + DataType::Float32 => ScalarValue::Float32(Some(d as f32)), + _ => ScalarValue::Float64(Some(d)), + } + } +} + +impl Accumulator for ApproxPercentileAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + let arr = downcast_value!(&values[0], Float64Array); + if arr.null_count() == 0 { + // Fast path: no validity checks needed, iterate the raw values. + for &v in arr.values() { + self.summary.insert(v); + } + } else { + for v in arr.iter().flatten() { + self.summary.insert(v); + } + } + Ok(()) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + let digests = downcast_value!(&states[0], BinaryArray); + self.summary.compress(); + for i in 0..digests.len() { + if digests.is_null(i) { + continue; + } + let peer = QuantileSummaries::from_bytes( + QuantileSummaries::DEFAULT_COMPRESS_THRESHOLD, + digests.value(i), + ); + if self.summary.count() == 0 { + // Empty self: merge would just clone the peer, so move it in. + self.summary = peer; + } else { + self.summary = self.summary.merge(&peer); + } + } + Ok(()) + } + + fn state(&mut self) -> Result<Vec<ScalarValue>> { + self.summary.compress(); + Ok(vec![ScalarValue::Binary(Some(self.summary.to_bytes()))]) + } + + fn evaluate(&mut self) -> Result<ScalarValue> { + self.summary.compress(); + match self.summary.query(&self.percentiles) { Review Comment: Does an empty percentage array reach here, or does Spark's analyzer reject it first? Spark returns null for empty `percentages`, but `query(&[])` returns `Some(vec![])` and would try to build an empty list. If it is unreachable a short comment would help, otherwise a `percentiles.is_empty()` short-circuit to null matches Spark. ########## native/spark-expr/src/agg_funcs/approx_percentile.rs: ########## @@ -0,0 +1,318 @@ +// 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. + +use super::quantile_summaries::QuantileSummaries; +use arrow::array::{Array, ArrayRef, BinaryArray, Float64Array, ListArray}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion::common::utils::SingleRowListArrayBuilder; +use datafusion::common::{downcast_value, Result, ScalarValue}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::Volatility::Immutable; +use datafusion::logical_expr::{Accumulator, AggregateUDFImpl, Signature}; +use datafusion::physical_expr::expressions::format_state_name; +use std::sync::Arc; + +/// Native implementation of Spark's `approx_percentile` / `percentile_approx`, +/// backed by a bit-for-bit `QuantileSummaries` (Greenwald-Khanna) port. The +/// child value is cast to Float64 by the serde; the original `input_type` is +/// carried so results can be cast back to Spark's output type. +#[derive(Debug)] +pub struct ApproxPercentile { + name: String, + signature: Signature, + percentiles: Vec<f64>, + accuracy: i64, + input_type: DataType, + return_array: bool, +} + +impl PartialEq for ApproxPercentile { + fn eq(&self, other: &Self) -> bool { + self.name == other.name + && self.percentiles == other.percentiles + && self.accuracy == other.accuracy + && self.input_type == other.input_type + && self.return_array == other.return_array + } +} +impl Eq for ApproxPercentile {} + +impl std::hash::Hash for ApproxPercentile { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + self.name.hash(state); + self.percentiles + .iter() + .for_each(|p| p.to_bits().hash(state)); + self.accuracy.hash(state); + self.input_type.hash(state); + self.return_array.hash(state); + } +} + +impl ApproxPercentile { + pub fn new( + percentiles: Vec<f64>, + accuracy: i64, + input_type: DataType, + return_array: bool, + ) -> Self { + Self { + name: "approx_percentile".to_string(), + signature: Signature::numeric(1, Immutable), + percentiles, + accuracy, + input_type, + return_array, + } + } +} + +impl AggregateUDFImpl for ApproxPercentile { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + if self.return_array { + Ok(DataType::List(Arc::new(Field::new( + "item", + self.input_type.clone(), + false, + )))) + } else { + Ok(self.input_type.clone()) + } + } + + fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + Ok(Box::new(ApproxPercentileAccumulator::new( + self.percentiles.clone(), + self.accuracy, + self.input_type.clone(), + self.return_array, + ))) + } + + fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + Ok(vec![Arc::new(Field::new( + format_state_name(&self.name, "digest"), + DataType::Binary, + true, + ))]) + } + + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + false + } +} + +#[derive(Debug)] +struct ApproxPercentileAccumulator { + summary: QuantileSummaries, + percentiles: Vec<f64>, + input_type: DataType, + return_array: bool, +} + +impl ApproxPercentileAccumulator { + fn new(percentiles: Vec<f64>, accuracy: i64, input_type: DataType, return_array: bool) -> Self { + let relative_error = 1.0 / accuracy as f64; + Self { + summary: QuantileSummaries::new( + QuantileSummaries::DEFAULT_COMPRESS_THRESHOLD, + relative_error, + ), + percentiles, + input_type, + return_array, + } + } + + /// Cast a double quantile back to Spark's output type. GK always returns an + /// actual inserted value (never an interpolation), so for the supported + /// numeric types this round-trips exactly and is always in range. + fn cast_back(&self, d: f64) -> ScalarValue { + match &self.input_type { + DataType::Int8 => ScalarValue::Int8(Some(d as i8)), + DataType::Int16 => ScalarValue::Int16(Some(d as i16)), + DataType::Int32 => ScalarValue::Int32(Some(d as i32)), + DataType::Int64 => ScalarValue::Int64(Some(d as i64)), + DataType::Float32 => ScalarValue::Float32(Some(d as f32)), + _ => ScalarValue::Float64(Some(d)), Review Comment: The `_ => Float64` arm silently absorbs any type outside the five explicit ones. Since the serde restricts inputs to six types, making `Float64` explicit and using `debug_assert!` / `unreachable!` for the rest would enforce the serde contract instead of masking a future mismatch. ########## native/spark-expr/src/agg_funcs/approx_percentile.rs: ########## @@ -0,0 +1,318 @@ +// 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. + +use super::quantile_summaries::QuantileSummaries; +use arrow::array::{Array, ArrayRef, BinaryArray, Float64Array, ListArray}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion::common::utils::SingleRowListArrayBuilder; +use datafusion::common::{downcast_value, Result, ScalarValue}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::Volatility::Immutable; +use datafusion::logical_expr::{Accumulator, AggregateUDFImpl, Signature}; +use datafusion::physical_expr::expressions::format_state_name; +use std::sync::Arc; + +/// Native implementation of Spark's `approx_percentile` / `percentile_approx`, +/// backed by a bit-for-bit `QuantileSummaries` (Greenwald-Khanna) port. The +/// child value is cast to Float64 by the serde; the original `input_type` is +/// carried so results can be cast back to Spark's output type. +#[derive(Debug)] +pub struct ApproxPercentile { + name: String, + signature: Signature, + percentiles: Vec<f64>, + accuracy: i64, + input_type: DataType, + return_array: bool, +} + +impl PartialEq for ApproxPercentile { + fn eq(&self, other: &Self) -> bool { + self.name == other.name + && self.percentiles == other.percentiles + && self.accuracy == other.accuracy + && self.input_type == other.input_type + && self.return_array == other.return_array + } +} +impl Eq for ApproxPercentile {} + +impl std::hash::Hash for ApproxPercentile { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + self.name.hash(state); + self.percentiles + .iter() + .for_each(|p| p.to_bits().hash(state)); + self.accuracy.hash(state); + self.input_type.hash(state); + self.return_array.hash(state); + } +} + +impl ApproxPercentile { + pub fn new( + percentiles: Vec<f64>, + accuracy: i64, + input_type: DataType, + return_array: bool, + ) -> Self { + Self { + name: "approx_percentile".to_string(), + signature: Signature::numeric(1, Immutable), + percentiles, + accuracy, + input_type, + return_array, + } + } +} + +impl AggregateUDFImpl for ApproxPercentile { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + if self.return_array { + Ok(DataType::List(Arc::new(Field::new( + "item", + self.input_type.clone(), + false, + )))) + } else { + Ok(self.input_type.clone()) + } + } + + fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + Ok(Box::new(ApproxPercentileAccumulator::new( + self.percentiles.clone(), + self.accuracy, + self.input_type.clone(), + self.return_array, + ))) + } + + fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + Ok(vec![Arc::new(Field::new( + format_state_name(&self.name, "digest"), + DataType::Binary, + true, + ))]) + } + + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + false + } +} + +#[derive(Debug)] +struct ApproxPercentileAccumulator { + summary: QuantileSummaries, + percentiles: Vec<f64>, + input_type: DataType, + return_array: bool, +} + +impl ApproxPercentileAccumulator { + fn new(percentiles: Vec<f64>, accuracy: i64, input_type: DataType, return_array: bool) -> Self { + let relative_error = 1.0 / accuracy as f64; + Self { + summary: QuantileSummaries::new( + QuantileSummaries::DEFAULT_COMPRESS_THRESHOLD, + relative_error, + ), + percentiles, + input_type, + return_array, + } + } + + /// Cast a double quantile back to Spark's output type. GK always returns an + /// actual inserted value (never an interpolation), so for the supported + /// numeric types this round-trips exactly and is always in range. + fn cast_back(&self, d: f64) -> ScalarValue { + match &self.input_type { + DataType::Int8 => ScalarValue::Int8(Some(d as i8)), + DataType::Int16 => ScalarValue::Int16(Some(d as i16)), + DataType::Int32 => ScalarValue::Int32(Some(d as i32)), + DataType::Int64 => ScalarValue::Int64(Some(d as i64)), + DataType::Float32 => ScalarValue::Float32(Some(d as f32)), + _ => ScalarValue::Float64(Some(d)), + } + } +} + +impl Accumulator for ApproxPercentileAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + let arr = downcast_value!(&values[0], Float64Array); + if arr.null_count() == 0 { + // Fast path: no validity checks needed, iterate the raw values. + for &v in arr.values() { + self.summary.insert(v); + } + } else { + for v in arr.iter().flatten() { + self.summary.insert(v); + } + } + Ok(()) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + let digests = downcast_value!(&states[0], BinaryArray); + self.summary.compress(); + for i in 0..digests.len() { + if digests.is_null(i) { + continue; + } + let peer = QuantileSummaries::from_bytes( + QuantileSummaries::DEFAULT_COMPRESS_THRESHOLD, + digests.value(i), + ); + if self.summary.count() == 0 { Review Comment: The `count() == 0` branch duplicates logic already inside `merge` (which handles `count == 0`). The only gain is avoiding one clone of the already-owned `peer`. If that is intentional for large digests a one-line note would help, otherwise `self.summary = self.summary.merge(&peer)` is simpler. ########## 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(); Review Comment: `new_samples` grows to `sampled.len() + sorted.len()`. `Vec::with_capacity(...)` here removes the reallocation on every flush. Same idea for reserving `head_sampled` in `update_batch` from the input array length. ########## 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: Not for this PR. `merge` takes `&self` / `&other` and allocates a fresh summary, so `merge_batch` reallocates on every incoming digest, and `with_head_buffer_inserted` rebuilds the whole `sampled` vector on every flush. A double-buffer for the flush and a consuming `merge(self, other: &Self) -> Self` (or `merge_into(&mut self, ...)`) would cut allocations if a future profile shows this path is hot. Noting it here as a possible follow-up, not a change to make now. -- 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]
