diff --git a/datafusion/functions-aggregate/src/percentile_cont.rs b/datafusion/functions-aggregate/src/percentile_cont.rs index 4b5e892fb5cdf..baf1978ea24af 100644 --- a/datafusion/functions-aggregate/src/percentile_cont.rs +++ b/datafusion/functions-aggregate/src/percentile_cont.rs @@ -32,8 +32,10 @@ use arrow::{ use num_traits::AsPrimitive; use arrow::array::ArrowNativeTypeOp; +use datafusion_common::hash_utils::RandomState; use datafusion_common::internal_err; use datafusion_common::types::{NativeType, logical_float64}; +use datafusion_common::utils::memory::estimate_memory_size; use datafusion_functions_aggregate_common::noop_accumulator::NoopAccumulator; use crate::min_max::{max_udaf, min_udaf}; @@ -672,7 +674,11 @@ where #[derive(Debug)] struct DistinctPercentileContAccumulator { /// Distinct value -> number of in-window rows carrying it. - counts: HashMap, usize>, + /// + /// Uses the same fast (foldhash) `RandomState` as the shared + /// `GenericDistinctBuffer` rather than the standard library's default + /// SipHash, which is considerably slower for this hot path. + counts: HashMap, usize, RandomState>, percentile: f64, } @@ -709,8 +715,15 @@ where // `values` may carry extra argument columns (e.g. the percentile // literal); only the first column holds the aggregated values. let arr = values[0].as_primitive::(); - for value in arr.iter().flatten() { - *self.counts.entry(Hashable(value)).or_default() += 1; + if arr.null_count() > 0 { + for value in arr.iter().flatten() { + *self.counts.entry(Hashable(value)).or_default() += 1; + } + } else { + // Fast path: no nulls, so skip the per-element validity check. + for value in arr.values().iter() { + *self.counts.entry(Hashable(*value)).or_default() += 1; + } } Ok(()) } @@ -733,9 +746,11 @@ where } fn size(&self) -> usize { - size_of_val(self) - + self.counts.capacity() - * (size_of::>() + size_of::()) + estimate_memory_size::<(Hashable, usize)>( + self.counts.capacity(), + size_of_val(self), + ) + .unwrap() } fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { @@ -744,12 +759,31 @@ where } let arr = values[0].as_primitive::(); - for value in arr.iter().flatten() { - if let Some(count) = self.counts.get_mut(&Hashable(value)) { - *count -= 1; - if *count == 0 { - self.counts.remove(&Hashable(value)); + let mut decrement = |value: T::Native| { + match self.counts.get_mut(&Hashable(value)) { + Some(count) => { + *count -= 1; + if *count == 0 { + self.counts.remove(&Hashable(value)); + } + Ok(()) } + // Retracting a value that isn't tracked means the accumulator + // state has diverged from the window frame; continuing would + // silently produce wrong results, so surface it as an error. + None => internal_err!( + "percentile_cont(DISTINCT) retract_batch: retracted a value not present in the window" + ), + } + }; + if arr.null_count() > 0 { + for value in arr.iter().flatten() { + decrement(value)?; + } + } else { + // Fast path: no nulls, so skip the per-element validity check. + for value in arr.values().iter() { + decrement(*value)?; } } Ok(()) diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index f6380c9a5f06d..26b8a78f3921a 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -1122,7 +1122,8 @@ SELECT percentile_cont(DISTINCT x, 0.5) FROM distinct_pct; query IR SELECT id, percentile_cont(DISTINCT x, 0.5) OVER (ORDER BY id ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) -FROM distinct_pct; +FROM distinct_pct +ORDER BY id; ---- 1 5 2 5 @@ -1131,6 +1132,46 @@ FROM distinct_pct; statement ok DROP TABLE distinct_pct; +# Regression: grouped percentile_cont(DISTINCT ...) forces two-phase +# (Partial + FinalPartitioned) aggregation, exercising the distinct +# accumulator's state()/merge_batch() paths. Duplicate values within a +# group must be de-duplicated across the per-partition merge. +# Group 1 distinct {1,5,9} -> median 5; group 2 distinct {3,7} -> median 5. +statement ok +CREATE TABLE grp_distinct_pct(g INT, x DOUBLE) AS VALUES + (1, 5), (1, 5), (1, 9), (1, 1), + (2, 7), (2, 7), (2, 3); + +query IR +SELECT g, percentile_cont(DISTINCT x, 0.5) FROM grp_distinct_pct GROUP BY g ORDER BY g; +---- +1 5 +2 5 + +statement ok +DROP TABLE grp_distinct_pct; + +# Regression: sliding-window percentile_cont(DISTINCT ...) over data with +# NULLs exercises the null_count() > 0 slow path in BOTH update_batch (a NULL +# row enters the frame) and retract_batch (a NULL row leaves the frame as the +# window slides). NULLs are ignored; distinct dedups the non-null values. +statement ok +CREATE TABLE distinct_pct_nulls(id INT, x DOUBLE) AS VALUES + (1, 5), (2, NULL), (3, 9), (4, 5); + +query IR +SELECT id, percentile_cont(DISTINCT x, 0.5) + OVER (ORDER BY id ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) +FROM distinct_pct_nulls ORDER BY id; +---- +1 5 +2 5 +3 9 +4 7 + +statement ok +DROP TABLE distinct_pct_nulls; + query RT select approx_median(arrow_cast(col_f32, 'Float16')), arrow_typeof(approx_median(arrow_cast(col_f32, 'Float16'))) from median_table; ----