diff --git a/datafusion/functions-aggregate/src/median.rs b/datafusion/functions-aggregate/src/median.rs index fb74da87c7fc..f5aa8afa5e21 100644 --- a/datafusion/functions-aggregate/src/median.rs +++ b/datafusion/functions-aggregate/src/median.rs @@ -15,48 +15,25 @@ // specific language governing permissions and limitations // under the License. -use std::cmp::Ordering; -use std::fmt::{Debug, Formatter}; -use std::mem::{size_of, size_of_val}; +use std::fmt::Debug; use std::sync::Arc; -use arrow::array::{ - ArrowNumericType, BooleanArray, ListArray, PrimitiveArray, PrimitiveBuilder, - downcast_integer, -}; -use arrow::buffer::{OffsetBuffer, ScalarBuffer}; -use arrow::{ - array::{ArrayRef, AsArray}, - datatypes::{ - DataType, Decimal128Type, Decimal256Type, Field, Float16Type, Float32Type, - Float64Type, - }, -}; +use arrow::datatypes::DataType; -use arrow::array::Array; -use arrow::array::ArrowNativeTypeOp; -use arrow::datatypes::{ - ArrowNativeType, ArrowPrimitiveType, Decimal32Type, Decimal64Type, FieldRef, -}; +use arrow::datatypes::FieldRef; -use datafusion_common::hash_utils::RandomState; -use datafusion_common::types::{NativeType, logical_float64}; -use datafusion_common::{ - DataFusionError, Result, ScalarValue, assert_eq_or_internal_err, exec_datafusion_err, - internal_datafusion_err, internal_err, -}; +use crate::percentile_cont::PercentileCont; +use datafusion_common::types::logical_float64; +use datafusion_common::{Result, assert_eq_or_internal_err}; +use datafusion_expr::GroupsAccumulator; use datafusion_expr::function::StateFieldsArgs; use datafusion_expr::{ Accumulator, AggregateUDFImpl, Coercion, Documentation, Signature, TypeSignature, - TypeSignatureClass, Volatility, function::AccumulatorArgs, utils::format_state_name, + TypeSignatureClass, Volatility, function::AccumulatorArgs, }; -use datafusion_expr::{EmitTo, GroupsAccumulator}; -use datafusion_functions_aggregate_common::aggregate::groups_accumulator::accumulate::accumulate; -use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::filtered_null_mask; -use datafusion_functions_aggregate_common::noop_accumulator::NoopAccumulator; -use datafusion_functions_aggregate_common::utils::{GenericDistinctBuffer, Hashable}; use datafusion_macros::user_doc; -use std::collections::HashMap; +use datafusion_physical_expr::expressions::lit; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; make_udaf_expr_and_func!( Median, @@ -91,6 +68,7 @@ make_udaf_expr_and_func!( #[derive(PartialEq, Eq, Hash, Debug)] pub struct Median { signature: Signature, + percentile_cont: PercentileCont, } impl Default for Median { @@ -113,18 +91,30 @@ impl Median { TypeSignature::Coercible(vec![Coercion::new_exact( TypeSignatureClass::Float, )]), - TypeSignature::Coercible(vec![Coercion::new_implicit( - TypeSignatureClass::Native(logical_float64()), + TypeSignature::Coercible(vec![Coercion::new_implicit_native( + logical_float64(), vec![TypeSignatureClass::Integer], - NativeType::Float64, )]), ], Volatility::Immutable, ), + percentile_cont: PercentileCont::new(), } } } +type PercentileExprsArgs = ([Arc; 2], [FieldRef; 2]); + +/// Build arguments for `percentile_cont` UDF +fn percentile_exprs_args(args: &AccumulatorArgs) -> Result { + let percentile_expr = lit(0.5_f64); + let percentile_field = percentile_expr.return_field(args.schema)?; + Ok(( + [Arc::clone(&args.exprs[0]), percentile_expr], + [Arc::clone(&args.expr_fields[0]), percentile_field], + )) +} + impl AggregateUDFImpl for Median { fn name(&self) -> &str { "median" @@ -135,80 +125,38 @@ impl AggregateUDFImpl for Median { } fn return_type(&self, arg_types: &[DataType]) -> Result { - Ok(arg_types[0].clone()) + self.percentile_cont.return_type(arg_types) } fn state_fields(&self, args: StateFieldsArgs) -> Result> { - if args.input_fields[0].data_type().is_null() { - return Ok(vec![ - Field::new( - format_state_name(args.name, self.name()), - DataType::Null, - true, - ) - .into(), - ]); - } - - //Intermediate state is a list of the elements we have collected so far - let field = Field::new_list_field(args.input_fields[0].data_type().clone(), true); - let state_name = if args.is_distinct { - "distinct_median" - } else { - "median" - }; - - Ok(vec![ - Field::new( - format_state_name(args.name, state_name), - DataType::List(Arc::new(field)), - true, - ) - .into(), - ]) + self.percentile_cont.state_fields(args) } - fn accumulator(&self, acc_args: AccumulatorArgs) -> Result> { - macro_rules! helper { - ($t:ty, $dt:expr) => { - if acc_args.is_distinct { - Ok(Box::new(DistinctMedianAccumulator::<$t> { - data_type: $dt.clone(), - distinct_values: GenericDistinctBuffer::new($dt), - })) - } else { - Ok(Box::new(MedianAccumulator::<$t> { - data_type: $dt.clone(), - all_values: vec![], - })) - } - }; - } - - let dt = acc_args.expr_fields[0].data_type().clone(); - if dt.is_null() { - return Ok(Box::new(NoopAccumulator::default())); - } - - downcast_integer! { - dt => (helper, dt), - DataType::Float16 => helper!(Float16Type, dt), - DataType::Float32 => helper!(Float32Type, dt), - DataType::Float64 => helper!(Float64Type, dt), - DataType::Decimal32(_, _) => helper!(Decimal32Type, dt), - DataType::Decimal64(_, _) => helper!(Decimal64Type, dt), - DataType::Decimal128(_, _) => helper!(Decimal128Type, dt), - DataType::Decimal256(_, _) => helper!(Decimal256Type, dt), - _ => Err(DataFusionError::NotImplemented(format!( - "MedianAccumulator not supported for {} with {}", - acc_args.name, - dt, - ))), - } + fn accumulator(&self, args: AccumulatorArgs) -> Result> { + let num_args = args.exprs.len(); + assert_eq_or_internal_err!( + num_args, + 1, + "median should only have 1 arg, but found num args:{}", + num_args + ); + let (exprs, expr_fields) = percentile_exprs_args(&args)?; + let sub_args = AccumulatorArgs { + exprs: &exprs, + expr_fields: &expr_fields, + return_field: Arc::clone(&args.return_field), + schema: args.schema, + ignore_nulls: args.ignore_nulls, + order_bys: args.order_bys, + is_reversed: args.is_reversed, + name: args.name, + is_distinct: args.is_distinct, + }; + self.percentile_cont.accumulator(sub_args) } fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { - !args.is_distinct && !args.expr_fields[0].data_type().is_null() + self.percentile_cont.groups_accumulator_supported(args) } fn create_groups_accumulator( @@ -222,502 +170,22 @@ impl AggregateUDFImpl for Median { "median should only have 1 arg, but found num args:{}", num_args ); - - let dt = args.expr_fields[0].data_type().clone(); - - macro_rules! helper { - ($t:ty, $dt:expr) => { - Ok(Box::new(MedianGroupsAccumulator::<$t>::new($dt))) - }; - } - - downcast_integer! { - dt => (helper, dt), - DataType::Float16 => helper!(Float16Type, dt), - DataType::Float32 => helper!(Float32Type, dt), - DataType::Float64 => helper!(Float64Type, dt), - DataType::Decimal32(_, _) => helper!(Decimal32Type, dt), - DataType::Decimal64(_, _) => helper!(Decimal64Type, dt), - DataType::Decimal128(_, _) => helper!(Decimal128Type, dt), - DataType::Decimal256(_, _) => helper!(Decimal256Type, dt), - _ => Err(DataFusionError::NotImplemented(format!( - "MedianGroupsAccumulator not supported for {} with {}", - args.name, - dt, - ))), - } + let (exprs, expr_fields) = percentile_exprs_args(&args)?; + let sub_args = AccumulatorArgs { + exprs: &exprs, + expr_fields: &expr_fields, + return_field: Arc::clone(&args.return_field), + schema: args.schema, + ignore_nulls: args.ignore_nulls, + order_bys: args.order_bys, + is_reversed: args.is_reversed, + name: args.name, + is_distinct: args.is_distinct, + }; + self.percentile_cont.create_groups_accumulator(sub_args) } fn documentation(&self) -> Option<&Documentation> { self.doc() } } - -/// The median accumulator accumulates the raw input values -/// as `ScalarValue`s -/// -/// The intermediate state is represented as a List of scalar values updated by -/// `merge_batch` and a `Vec` of `ArrayRef` that are converted to scalar values -/// in the final evaluation step so that we avoid expensive conversions and -/// allocations during `update_batch`. -struct MedianAccumulator { - data_type: DataType, - all_values: Vec, -} - -impl Debug for MedianAccumulator { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "MedianAccumulator({})", self.data_type) - } -} - -impl Accumulator for MedianAccumulator { - fn state(&mut self) -> Result> { - // Convert `all_values` to `ListArray` and return a single List ScalarValue - - // Build offsets - let offsets = - OffsetBuffer::new(ScalarBuffer::from(vec![0, self.all_values.len() as i32])); - - // Build inner array - let values_array = PrimitiveArray::::new( - ScalarBuffer::from(std::mem::take(&mut self.all_values)), - None, - ) - .with_data_type(self.data_type.clone()); - - // Build the result list array - let list_array = ListArray::new( - Arc::new(Field::new_list_field(self.data_type.clone(), true)), - offsets, - Arc::new(values_array), - None, - ); - - Ok(vec![ScalarValue::List(Arc::new(list_array))]) - } - - fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - let values = values[0].as_primitive::(); - let additional = values.len() - values.null_count(); - self.all_values.try_reserve(additional).map_err(|e| { - exec_datafusion_err!( - "failed to reserve {additional} values for median accumulator: {e}" - ) - })?; - if values.null_count() > 0 { - self.all_values.extend(values.iter().flatten()); - } else { - // Fast path: no nulls, so the values buffer can be appended wholesale. - self.all_values.extend_from_slice(values.values()); - } - Ok(()) - } - - fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { - let array = states[0].as_list::(); - for v in array.iter().flatten() { - self.update_batch(&[v])? - } - Ok(()) - } - - fn evaluate(&mut self) -> Result { - let median = calculate_median::(&mut self.all_values); - ScalarValue::new_primitive::(median, &self.data_type) - } - - fn size(&self) -> usize { - size_of_val(self) + self.all_values.capacity() * size_of::() - } - - fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - let mut to_remove: HashMap, usize, RandomState> = - HashMap::default(); - - let arr = values[0].as_primitive::(); - if arr.null_count() > 0 { - for value in arr.iter().flatten() { - *to_remove.entry(Hashable(value)).or_default() += 1; - } - } else { - // Fast path: no nulls, so skip the per-element validity check. - for value in arr.values().iter() { - *to_remove.entry(Hashable(*value)).or_default() += 1; - } - } - - let mut i = 0; - while i < self.all_values.len() { - let k = Hashable(self.all_values[i]); - if let Some(count) = to_remove.get_mut(&k) - && *count > 0 - { - self.all_values.swap_remove(i); - *count -= 1; - if *count == 0 { - to_remove.remove(&k); - if to_remove.is_empty() { - break; - } - } - } else { - i += 1; - } - } - - // Retracting values that are not tracked means the accumulator state - // has diverged from the window frame; continuing would silently - // produce wrong results, so surface it as an error. - if !to_remove.is_empty() { - return internal_err!( - "median retract_batch: retracted value(s) not present in the window" - ); - } - Ok(()) - } - - fn supports_retract_batch(&self) -> bool { - true - } -} - -/// The median groups accumulator accumulates the raw input values -/// -/// For calculating the accurate medians of groups, we need to store all values -/// of groups before final evaluation. -/// So values in each group will be stored in a `Vec`, and the total group values -/// will be actually organized as a `Vec>`. -#[derive(Debug)] -struct MedianGroupsAccumulator { - data_type: DataType, - group_values: Vec>, -} - -impl MedianGroupsAccumulator { - pub fn new(data_type: DataType) -> Self { - Self { - data_type, - group_values: Vec::new(), - } - } -} - -impl GroupsAccumulator for MedianGroupsAccumulator { - fn update_batch( - &mut self, - values: &[ArrayRef], - group_indices: &[usize], - opt_filter: Option<&BooleanArray>, - total_num_groups: usize, - ) -> Result<()> { - assert_eq!(values.len(), 1, "single argument to update_batch"); - let values = values[0].as_primitive::(); - - // Push the `not nulls + not filtered` row into its group - self.group_values.resize(total_num_groups, Vec::new()); - accumulate( - group_indices, - values, - opt_filter, - |group_index, new_value| { - self.group_values[group_index].push(new_value); - }, - ); - - Ok(()) - } - - fn merge_batch( - &mut self, - values: &[ArrayRef], - group_indices: &[usize], - total_num_groups: usize, - ) -> Result<()> { - assert_eq!(values.len(), 1, "one argument to merge_batch"); - - // The merged values should be organized like as a `ListArray` which is nullable - // (input with nulls usually generated from `convert_to_state`), but `inner array` of - // `ListArray` is `non-nullable`. - // - // Following is the possible and impossible input `values`: - // - // # Possible values - // ```text - // group 0: [1, 2, 3] - // group 1: null (list array is nullable) - // group 2: [6, 7, 8] - // ... - // group n: [...] - // ``` - // - // # Impossible values - // ```text - // group x: [1, 2, null] (values in list array is non-nullable) - // ``` - // - let input_group_values = values[0].as_list::(); - - // Ensure group values big enough - self.group_values.resize(total_num_groups, Vec::new()); - - // Extend values to related groups - // TODO: avoid using iterator of the `ListArray`, this will lead to - // many calls of `slice` of its ``inner array`, and `slice` is not - // so efficient(due to the calculation of `null_count` for each `slice`). - group_indices - .iter() - .zip(input_group_values.iter()) - .for_each(|(&group_index, values_opt)| { - if let Some(values) = values_opt { - let values = values.as_primitive::(); - self.group_values[group_index].extend(values.values().iter()); - } - }); - - Ok(()) - } - - fn state(&mut self, emit_to: EmitTo) -> Result> { - // Emit values - let emit_group_values = emit_to.take_needed(&mut self.group_values); - - // Build offsets - let mut offsets = Vec::with_capacity(self.group_values.len() + 1); - offsets.push(0); - let mut cur_len = 0_i32; - for group_value in &emit_group_values { - cur_len += group_value.len() as i32; - offsets.push(cur_len); - } - // TODO: maybe we can use `OffsetBuffer::new_unchecked` like what in `convert_to_state`, - // but safety should be considered more carefully here(and I am not sure if it can get - // performance improvement when we introduce checks to keep the safety...). - // - // Can see more details in: - // https://github.com/apache/datafusion/pull/13681#discussion_r1931209791 - // - let offsets = OffsetBuffer::new(ScalarBuffer::from(offsets)); - - // Build inner array - let flatten_group_values = - emit_group_values.into_iter().flatten().collect::>(); - let group_values_array = - PrimitiveArray::::new(ScalarBuffer::from(flatten_group_values), None) - .with_data_type(self.data_type.clone()); - - // Build the result list array - let result_list_array = ListArray::new( - Arc::new(Field::new_list_field(self.data_type.clone(), true)), - offsets, - Arc::new(group_values_array), - None, - ); - - Ok(vec![Arc::new(result_list_array)]) - } - - fn evaluate(&mut self, emit_to: EmitTo) -> Result { - // Emit values - let emit_group_values = emit_to.take_needed(&mut self.group_values); - - // Calculate median for each group - let mut evaluate_result_builder = - PrimitiveBuilder::::new().with_data_type(self.data_type.clone()); - for mut values in emit_group_values { - let median = calculate_median::(&mut values); - evaluate_result_builder.append_option(median); - } - - Ok(Arc::new(evaluate_result_builder.finish())) - } - - fn convert_to_state( - &self, - values: &[ArrayRef], - opt_filter: Option<&BooleanArray>, - ) -> Result> { - assert_eq!(values.len(), 1, "one argument to merge_batch"); - - let input_array = values[0].as_primitive::(); - - // Directly convert the input array to states, each row will be - // seen as a respective group. - // For detail, the `input_array` will be converted to a `ListArray`. - // And if row is `not null + not filtered`, it will be converted to a list - // with only one element; otherwise, this row in `ListArray` will be set - // to null. - - // Reuse values buffer in `input_array` to build `values` in `ListArray` - let values = PrimitiveArray::::new(input_array.values().clone(), None) - .with_data_type(self.data_type.clone()); - - // `offsets` in `ListArray`, each row as a list element - let offset_end = i32::try_from(input_array.len()).map_err(|e| { - internal_datafusion_err!( - "cast array_len to i32 failed in convert_to_state of group median, err:{e:?}" - ) - })?; - let offsets = (0..=offset_end).collect::>(); - // Safety: all checks in `OffsetBuffer::new` are ensured to pass - let offsets = unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(offsets)) }; - - // `nulls` for converted `ListArray` - let nulls = filtered_null_mask(opt_filter, input_array); - - let converted_list_array = ListArray::new( - Arc::new(Field::new_list_field(self.data_type.clone(), true)), - offsets, - Arc::new(values), - nulls, - ); - - Ok(vec![Arc::new(converted_list_array)]) - } - fn size(&self) -> usize { - self.group_values - .iter() - .map(|values| values.capacity() * size_of::()) - .sum::() - // account for size of self.group_values too - + self.group_values.capacity() * size_of::>() - + size_of::>>() - } -} - -#[derive(Debug)] -struct DistinctMedianAccumulator { - distinct_values: GenericDistinctBuffer, - data_type: DataType, -} - -impl Accumulator for DistinctMedianAccumulator { - fn state(&mut self) -> Result> { - self.distinct_values.state() - } - - fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - self.distinct_values.update_batch(values) - } - - fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { - self.distinct_values.merge_batch(states) - } - - fn evaluate(&mut self) -> Result { - let mut d: Vec = - self.distinct_values.values.iter().map(|v| v.0).collect(); - let median = calculate_median::(&mut d); - ScalarValue::new_primitive::(median, &self.data_type) - } - - fn size(&self) -> usize { - size_of_val(self) + self.distinct_values.size() - } -} - -/// Get maximum entry in the slice, -fn slice_max(array: &[T::Native]) -> T::Native -where - T: ArrowPrimitiveType, - T::Native: PartialOrd, // Ensure the type supports PartialOrd for comparison -{ - // Make sure that, array is not empty. - debug_assert!(!array.is_empty()); - // `.unwrap()` is safe here as the array is supposed to be non-empty - *array - .iter() - .max_by(|x, y| x.partial_cmp(y).unwrap_or(Ordering::Less)) - .unwrap() -} - -fn calculate_median(values: &mut [T::Native]) -> Option { - let cmp = |x: &T::Native, y: &T::Native| x.compare(*y); - - let len = values.len(); - if len == 0 { - None - } else if len % 2 == 0 { - let (low, high, _) = values.select_nth_unstable_by(len / 2, cmp); - // Get the maximum of the low (left side after bi-partitioning) - let left_max = slice_max::(low); - // Calculate median as the average of the two middle values. - // Use checked arithmetic to detect overflow and fall back to safe formula. - let two = T::Native::usize_as(2); - let median = match left_max.add_checked(*high) { - Ok(sum) => sum.div_wrapping(two), - Err(_) => { - // Overflow detected - use safe midpoint formula: - // a/2 + b/2 + ((a%2 + b%2) / 2) - // This avoids overflow by dividing before adding. - let half_left = left_max.div_wrapping(two); - let half_right = (*high).div_wrapping(two); - let rem_left = left_max.mod_wrapping(two); - let rem_right = (*high).mod_wrapping(two); - // The sum of remainders (0, 1, or 2 for unsigned; -2 to 2 for signed) - // divided by 2 gives the correction factor (0 or 1 for unsigned; -1, 0, or 1 for signed) - let correction = rem_left.add_wrapping(rem_right).div_wrapping(two); - half_left.add_wrapping(half_right).add_wrapping(correction) - } - }; - Some(median) - } else { - let (_, median, _) = values.select_nth_unstable_by(len / 2, cmp); - Some(*median) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow::array::Float64Array; - - fn median_accumulator() -> MedianAccumulator { - MedianAccumulator { - data_type: DataType::Float64, - all_values: vec![], - } - } - - #[test] - fn retract_batch_errors_on_untracked_value() { - let mut acc = median_accumulator(); - let values: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0])); - acc.update_batch(std::slice::from_ref(&values)).unwrap(); - - let retract: ArrayRef = Arc::new(Float64Array::from(vec![3.0])); - let err = acc - .retract_batch(std::slice::from_ref(&retract)) - .unwrap_err() - .to_string(); - assert!( - err.contains("not present in the window"), - "unexpected error: {err}" - ); - } - - #[test] - fn update_batch_with_and_without_nulls_agree() { - // The null-free fast path must accumulate the same values as the - // general path. - let dense: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])); - let sparse: ArrayRef = Arc::new(Float64Array::from(vec![ - Some(1.0), - None, - Some(2.0), - None, - Some(3.0), - ])); - - let mut dense_acc = median_accumulator(); - dense_acc - .update_batch(std::slice::from_ref(&dense)) - .unwrap(); - let mut sparse_acc = median_accumulator(); - sparse_acc - .update_batch(std::slice::from_ref(&sparse)) - .unwrap(); - - assert_eq!(dense_acc.all_values, sparse_acc.all_values); - } -} diff --git a/datafusion/functions-aggregate/src/percentile_cont.rs b/datafusion/functions-aggregate/src/percentile_cont.rs index 3a98900bbb44..9703bb8f1e72 100644 --- a/datafusion/functions-aggregate/src/percentile_cont.rs +++ b/datafusion/functions-aggregate/src/percentile_cont.rs @@ -17,11 +17,13 @@ use std::collections::HashMap; use std::fmt::Debug; +use std::marker::PhantomData; use std::mem::{size_of, size_of_val}; use std::sync::Arc; use arrow::array::{ - ArrowNumericType, BooleanArray, ListArray, PrimitiveArray, PrimitiveBuilder, + ArrowNumericType, ArrowPrimitiveType, BooleanArray, ListArray, PrimitiveArray, + PrimitiveBuilder, downcast_integer, }; use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::{ @@ -32,6 +34,10 @@ use arrow::{ use num_traits::AsPrimitive; use arrow::array::ArrowNativeTypeOp; +use arrow::datatypes::{ + ArrowNativeType, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, + DecimalType, +}; use datafusion_common::hash_utils::RandomState; use datafusion_common::internal_err; use datafusion_common::types::{NativeType, logical_float64}; @@ -46,7 +52,7 @@ use datafusion_common::{ use datafusion_expr::utils::format_state_name; use datafusion_expr::{ Accumulator, AggregateUDFImpl, Coercion, Documentation, Expr, Signature, - TypeSignatureClass, Volatility, + TypeSignature, TypeSignatureClass, Volatility, }; use datafusion_expr::{EmitTo, GroupsAccumulator}; use datafusion_expr::{ @@ -61,21 +67,6 @@ use datafusion_macros::user_doc; use crate::utils::validate_percentile_expr; -/// Precision multiplier for linear interpolation calculations. -/// -/// This value of 1,000,000 was chosen to balance precision with overflow safety: -/// - Provides 6 decimal places of precision for the fractional component -/// - Small enough to avoid overflow when multiplied with typical numeric values -/// - Sufficient precision for most statistical applications -/// -/// The interpolation formula: `lower + (upper - lower) * fraction` -/// is computed as: `lower + ((upper - lower) * (fraction * PRECISION)) / PRECISION` -/// to avoid floating-point operations on integer types while maintaining precision. -/// -/// The interpolation arithmetic is performed in f64 and then cast back to the -/// native type to avoid overflowing Float16 intermediates. -const INTERPOLATION_PRECISION: f64 = 1_000_000.0; - create_func!(PercentileCont, percentile_cont_udaf); /// Computes the exact percentile continuous of a set of numbers @@ -143,18 +134,30 @@ impl Default for PercentileCont { impl PercentileCont { pub fn new() -> Self { Self { - signature: Signature::coercible( + signature: Signature::one_of( vec![ - Coercion::new_implicit( - TypeSignatureClass::Float, - vec![TypeSignatureClass::Numeric], - NativeType::Float64, - ), - Coercion::new_implicit( - TypeSignatureClass::Native(logical_float64()), - vec![TypeSignatureClass::Numeric], - NativeType::Float64, - ), + // Decimal signature: decimals, percentile + TypeSignature::Coercible(vec![ + // value + Coercion::new_exact(TypeSignatureClass::Decimal), + // percentile + Coercion::new_implicit_native( + logical_float64(), + vec![TypeSignatureClass::Numeric], + ), + ]), + // Float signature: float, percentile + TypeSignature::Coercible(vec![ + Coercion::new_implicit( + TypeSignatureClass::Float, + vec![TypeSignatureClass::Numeric], + NativeType::Float64, + ), + Coercion::new_implicit_native( + logical_float64(), + vec![TypeSignatureClass::Numeric], + ), + ]), ], Volatility::Immutable, ) @@ -179,10 +182,7 @@ impl AggregateUDFImpl for PercentileCont { } fn return_type(&self, arg_types: &[DataType]) -> Result { - match &arg_types[0] { - DataType::Null => Ok(DataType::Float64), - dt => Ok(dt.clone()), - } + Ok(arg_types[0].clone()) } fn state_fields(&self, args: StateFieldsArgs) -> Result> { @@ -216,39 +216,46 @@ impl AggregateUDFImpl for PercentileCont { } fn accumulator(&self, args: AccumulatorArgs) -> Result> { + // Always verify percentiles let percentile = get_percentile(&args)?; let input_dt = args.expr_fields[0].data_type(); + // Null input evaluates to null if input_dt.is_null() { - return Ok(Box::new(NoopAccumulator::new(ScalarValue::Float64(None)))); + return Ok(Box::new(NoopAccumulator::default())); } - if args.is_distinct { - match input_dt { - DataType::Float16 => Ok(Box::new(DistinctPercentileContAccumulator::< - Float16Type, - >::new(percentile))), - DataType::Float32 => Ok(Box::new(DistinctPercentileContAccumulator::< - Float32Type, - >::new(percentile))), - DataType::Float64 => Ok(Box::new(DistinctPercentileContAccumulator::< - Float64Type, - >::new(percentile))), - dt => internal_err!("Unsupported datatype for percentile cont: {dt}"), - } - } else { - match input_dt { - DataType::Float16 => Ok(Box::new( - PercentileContAccumulator::::new(percentile), - )), - DataType::Float32 => Ok(Box::new( - PercentileContAccumulator::::new(percentile), - )), - DataType::Float64 => Ok(Box::new( - PercentileContAccumulator::::new(percentile), - )), - dt => internal_err!("Unsupported datatype for percentile cont: {dt}"), - } + macro_rules! helper { + ($t:ty, $i:ty, $dt:expr) => { + if args.is_distinct { + Ok(Box::new(DistinctPercentileContAccumulator::<$t, $i>::new( + percentile, + $dt.clone(), + ))) + } else { + Ok(Box::new(PercentileContAccumulator::<$t, $i>::new( + percentile, + $dt.clone(), + ))) + } + }; + } + macro_rules! integer_helper { + ($t:ty, $dt:expr) => { + helper!($t, IntegerInterpolator, $dt) + }; + } + + downcast_integer! { + input_dt => (integer_helper, input_dt), + DataType::Float16 => helper!(Float16Type, FloatInterpolator, input_dt), + DataType::Float32 => helper!(Float32Type, FloatInterpolator, input_dt), + DataType::Float64 => helper!(Float64Type, FloatInterpolator, input_dt), + DataType::Decimal32(_, _) => helper!(Decimal32Type, DecimalInterpolator, input_dt), + DataType::Decimal64(_, _) => helper!(Decimal64Type, DecimalInterpolator, input_dt), + DataType::Decimal128(_, _) => helper!(Decimal128Type, DecimalInterpolator, input_dt), + DataType::Decimal256(_, _) => helper!(Decimal256Type, DecimalInterpolator, input_dt), + dt => internal_err!("Unsupported datatype for {} with {}", args.name, dt), } } @@ -260,20 +267,35 @@ impl AggregateUDFImpl for PercentileCont { &self, args: AccumulatorArgs, ) -> Result> { + // Always verify percentiles let percentile = get_percentile(&args)?; let input_dt = args.expr_fields[0].data_type(); - match input_dt { - DataType::Float16 => Ok(Box::new(PercentileContGroupsAccumulator::< - Float16Type, - >::new(percentile))), - DataType::Float32 => Ok(Box::new(PercentileContGroupsAccumulator::< - Float32Type, - >::new(percentile))), - DataType::Float64 => Ok(Box::new(PercentileContGroupsAccumulator::< - Float64Type, - >::new(percentile))), - dt => internal_err!("Unsupported datatype for percentile cont: {dt}"), + + macro_rules! helper { + ($t:ty, $i:ty, $dt:expr) => { + Ok(Box::new(PercentileContGroupsAccumulator::<$t, $i>::new( + percentile, + $dt.clone(), + ))) + }; + } + macro_rules! integer_helper { + ($t:ty, $dt:expr) => { + helper!($t, IntegerInterpolator, $dt) + }; + } + + downcast_integer! { + input_dt => (integer_helper, input_dt), + DataType::Float16 => helper!(Float16Type, FloatInterpolator, input_dt), + DataType::Float32 => helper!(Float32Type, FloatInterpolator, input_dt), + DataType::Float64 => helper!(Float64Type, FloatInterpolator, input_dt), + DataType::Decimal32(_, _) => helper!(Decimal32Type, DecimalInterpolator, input_dt), + DataType::Decimal64(_, _) => helper!(Decimal64Type, DecimalInterpolator, input_dt), + DataType::Decimal128(_, _) => helper!(Decimal128Type, DecimalInterpolator, input_dt), + DataType::Decimal256(_, _) => helper!(Decimal256Type, DecimalInterpolator, input_dt), + dt => internal_err!("Unsupported datatype for {} with {}", args.name, dt), } } @@ -377,25 +399,33 @@ fn simplify_percentile_cont_aggregate( /// in the final evaluation step so that we avoid expensive conversions and /// allocations during `update_batch`. #[derive(Debug)] -struct PercentileContAccumulator { +struct PercentileContAccumulator< + T: ArrowNumericType + Debug, + I: PercentileInterpolator, +> { all_values: Vec, percentile: f64, + data_type: DataType, + _interpolator: PhantomData, } -impl PercentileContAccumulator { - fn new(percentile: f64) -> Self { +impl> + PercentileContAccumulator +{ + fn new(percentile: f64, data_type: DataType) -> Self { Self { all_values: vec![], percentile, + data_type, + _interpolator: PhantomData, } } } -impl Accumulator for PercentileContAccumulator +impl Accumulator for PercentileContAccumulator where T: ArrowNumericType + Debug, - T::Native: Copy + AsPrimitive, - f64: AsPrimitive, + I: PercentileInterpolator + 'static, { fn state(&mut self) -> Result> { // Convert `all_values` to `ListArray` and return a single List ScalarValue @@ -408,11 +438,12 @@ where let values_array = PrimitiveArray::::new( ScalarBuffer::from(std::mem::take(&mut self.all_values)), None, - ); + ) + .with_data_type(self.data_type.clone()); // Build the result list array let list_array = ListArray::new( - Arc::new(Field::new_list_field(T::DATA_TYPE, true)), + Arc::new(Field::new_list_field(self.data_type.clone(), true)), offsets, Arc::new(values_array), None, @@ -440,13 +471,16 @@ where fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { let array = states[0].as_list::(); - self.update_batch(&[array.value(0)])?; + // Feed all list elements from a batch + for values in array.iter().flatten() { + self.update_batch(&[values])?; + } Ok(()) } fn evaluate(&mut self) -> Result { - let value = calculate_percentile::(&mut self.all_values, self.percentile); - ScalarValue::new_primitive::(value, &T::DATA_TYPE) + let value = calculate_percentile::(&mut self.all_values, self.percentile)?; + ScalarValue::new_primitive::(value, &self.data_type) } fn size(&self) -> usize { @@ -511,25 +545,33 @@ where /// So values in each group will be stored in a `Vec`, and the total group values /// will be actually organized as a `Vec>`. #[derive(Debug)] -struct PercentileContGroupsAccumulator { +struct PercentileContGroupsAccumulator< + T: ArrowNumericType + Debug, + I: PercentileInterpolator, +> { group_values: Vec>, percentile: f64, + data_type: DataType, + _interpolator: PhantomData, } -impl PercentileContGroupsAccumulator { - fn new(percentile: f64) -> Self { +impl> + PercentileContGroupsAccumulator +{ + fn new(percentile: f64, data_type: DataType) -> Self { Self { group_values: vec![], percentile, + data_type, + _interpolator: PhantomData, } } } -impl GroupsAccumulator for PercentileContGroupsAccumulator +impl GroupsAccumulator for PercentileContGroupsAccumulator where - T: ArrowNumericType + Send, - T::Native: Copy + AsPrimitive, - f64: AsPrimitive, + T: ArrowNumericType + Debug + Send, + I: PercentileInterpolator + 'static, { fn update_batch( &mut self, @@ -602,11 +644,12 @@ where let flatten_group_values = emit_group_values.into_iter().flatten().collect::>(); let group_values_array = - PrimitiveArray::::new(ScalarBuffer::from(flatten_group_values), None); + PrimitiveArray::::new(ScalarBuffer::from(flatten_group_values), None) + .with_data_type(self.data_type.clone()); // Build the result list array let result_list_array = ListArray::new( - Arc::new(Field::new_list_field(T::DATA_TYPE, true)), + Arc::new(Field::new_list_field(self.data_type.clone(), true)), offsets, Arc::new(group_values_array), None, @@ -621,9 +664,11 @@ where // Calculate percentile for each group let mut evaluate_result_builder = - PrimitiveBuilder::::with_capacity(emit_group_values.len()); + PrimitiveBuilder::::with_capacity(emit_group_values.len()) + .with_data_type(self.data_type.clone()); for values in &mut emit_group_values { - let value = calculate_percentile::(values.as_mut_slice(), self.percentile); + let value = + calculate_percentile::(values.as_mut_slice(), self.percentile)?; evaluate_result_builder.append_option(value); } @@ -647,7 +692,8 @@ where // to null. // Reuse values buffer in `input_array` to build `values` in `ListArray` - let values = PrimitiveArray::::new(input_array.values().clone(), None); + let values = PrimitiveArray::::new(input_array.values().clone(), None) + .with_data_type(self.data_type.clone()); // `offsets` in `ListArray`, each row as a list element let offset_end = i32::try_from(input_array.len()).map_err(|e| { @@ -668,7 +714,7 @@ where let nulls = filtered_null_mask(opt_filter, input_array); let converted_list_array = ListArray::new( - Arc::new(Field::new_list_field(T::DATA_TYPE, true)), + Arc::new(Field::new_list_field(self.data_type.clone(), true)), offsets, Arc::new(values), nulls, @@ -694,7 +740,10 @@ where /// left the window frame. The percentile is then computed over the set of keys /// with a positive count. #[derive(Debug)] -struct DistinctPercentileContAccumulator { +struct DistinctPercentileContAccumulator< + T: ArrowNumericType, + I: PercentileInterpolator, +> { /// Distinct value -> number of in-window rows carrying it. /// /// Uses the same fast (foldhash) `RandomState` as the shared @@ -702,22 +751,27 @@ struct DistinctPercentileContAccumulator { /// SipHash, which is considerably slower for this hot path. counts: HashMap, usize, RandomState>, percentile: f64, + data_type: DataType, + _interpolator: PhantomData, } -impl DistinctPercentileContAccumulator { - fn new(percentile: f64) -> Self { +impl> + DistinctPercentileContAccumulator +{ + fn new(percentile: f64, data_type: DataType) -> Self { Self { counts: HashMap::default(), percentile, + data_type, + _interpolator: PhantomData, } } } -impl Accumulator for DistinctPercentileContAccumulator +impl Accumulator for DistinctPercentileContAccumulator where T: ArrowNumericType + Debug, - T::Native: Copy + AsPrimitive, - f64: AsPrimitive, + I: PercentileInterpolator + 'static, { fn state(&mut self) -> Result> { // Emit the distinct keys as a single List scalar, matching the state @@ -726,7 +780,7 @@ where // cross-partition merges only need the distinct key set. let arr = Arc::new( PrimitiveArray::::from_iter_values(self.counts.keys().map(|v| v.0)) - .with_data_type(T::DATA_TYPE), + .with_data_type(self.data_type.clone()), ); Ok(vec![ SingleRowListArrayBuilder::new(arr).build_list_scalar(), @@ -763,8 +817,8 @@ where fn evaluate(&mut self) -> Result { let mut values: Vec = self.counts.keys().map(|v| v.0).collect(); - let value = calculate_percentile::(&mut values, self.percentile); - ScalarValue::new_primitive::(value, &T::DATA_TYPE) + let value = calculate_percentile::(&mut values, self.percentile)?; + ScalarValue::new_primitive::(value, &self.data_type) } fn size(&self) -> usize { @@ -816,6 +870,176 @@ where } } +/// A trait to abstract interpolation logic for percentile calculation +/// for floats and decimals. +trait PercentileInterpolator: Debug + Sync + Send { + fn interpolate( + lower: T::Native, + upper: T::Native, + fraction: f64, + ) -> Result; +} + +#[derive(Debug)] +struct FloatInterpolator; + +/// Precision multiplier for floating-point linear interpolation calculations. +/// +/// This value of 1,000,000 was chosen to balance precision with overflow safety: +/// - Provides 6 decimal places of precision for the fractional component +/// - Small enough to avoid overflow when multiplied with typical numeric values +/// - Sufficient precision for most statistical applications +/// +/// The interpolation formula: `lower + (upper - lower) * fraction` +/// is computed as: `lower + ((upper - lower) * (fraction * PRECISION)) / PRECISION` +/// to avoid floating-point operations on integer types while maintaining precision. +/// +/// The interpolation arithmetic for floats is performed in f64 and then cast back to the +/// native type to avoid overflowing Float16 intermediates. +const FLOAT_INTERPOLATION_PRECISION: f64 = 1_000_000.0; + +impl PercentileInterpolator for FloatInterpolator +where + T: ArrowNumericType, + T::Native: AsPrimitive, + f64: AsPrimitive, +{ + fn interpolate( + lower: T::Native, + upper: T::Native, + fraction: f64, + ) -> Result { + // Linear interpolation. + // We compute a quantized interpolation weight using `FLOAT_INTERPOLATION_PRECISION` because: + // 1. Both values come from the input data, so (upper - lower) is bounded by the value range + // 2. fraction is between 0 and 1; quantizing it provides stable, predictable results + // 3. The result is guaranteed to be between lower_value and upper_value (modulo cast rounding) + // 4. Arithmetic is performed in f64 and cast back to avoid overflowing Float16 intermediates + let scaled = (fraction * FLOAT_INTERPOLATION_PRECISION) as usize; + let weight = scaled as f64 / FLOAT_INTERPOLATION_PRECISION; + + let lower_f: f64 = lower.as_(); + let upper_f: f64 = upper.as_(); + let interpolated_f = lower_f + (upper_f - lower_f) * weight; + Ok(interpolated_f.as_()) + } +} + +#[derive(Debug)] +struct IntegerInterpolator; + +/// Precision multiplier for integer linear interpolation. +/// +/// The arithmetic is performed in `i128` so that it cannot overflow for the +/// widest supported integer: `(upper - lower)` is at most `u64::MAX` and +/// multiplying that by the precision stays well inside `i128`. +const INTEGER_INTERPOLATION_PRECISION: i128 = 1_000_000; + +impl PercentileInterpolator for IntegerInterpolator +where + T: ArrowNumericType, + T::Native: AsPrimitive, + i128: AsPrimitive, +{ + fn interpolate( + lower: T::Native, + upper: T::Native, + fraction: f64, + ) -> Result { + debug_assert!((0.0..=1.0).contains(&fraction)); + + let lower_i: i128 = lower.as_(); + let upper_i: i128 = upper.as_(); + debug_assert!(lower_i <= upper_i); + + // `lower + (upper - lower) * fraction`, rounding down. + let num = (fraction * INTEGER_INTERPOLATION_PRECISION as f64) as i128; + let interpolated = + lower_i + (upper_i - lower_i) * num / INTEGER_INTERPOLATION_PRECISION; + + // The result lies between `lower` and `upper`, so it fits in `T::Native`. + Ok(interpolated.as_()) + } +} + +#[derive(Debug)] +struct DecimalInterpolator; + +/// Precision multiplier for decimal linear interpolation calculations. +fn deduce_interpolation_precision() -> usize { + if T::BYTE_LENGTH == 4 { + // Decimal32 + 10_000 + } else { + 1_000_000 + } +} + +/// Compute a scaled value for interpolation using a formula `trunc(x * num / precision)` +/// The numerical method is separating `q * num + r * num / den` +/// where `q = x / precision` and `r = x % precision` +fn scale_by_num(x: T::Native, num: T::Native, den: T::Native) -> Result +where + T: DecimalType, +{ + debug_assert!(num >= T::Native::ZERO); + debug_assert!(num <= den); + + let q = x.div_wrapping(den); + let r = x.mod_wrapping(den); + // `q * num` cannot exceed `x`; `r * num` cannot exceed `(den - 1)^2` + + let a = q.mul_checked(num)?; + let b = r.mul_checked(num)?.div_wrapping(den); + a.add_checked(b) + .map_err(|e| exec_datafusion_err!("Arithmetic overflow in percentile_cont: {e}")) +} + +impl PercentileInterpolator for DecimalInterpolator +where + T: DecimalType, +{ + fn interpolate( + lower: T::Native, + upper: T::Native, + fraction: f64, + ) -> Result { + debug_assert!((0.0..=1.0).contains(&fraction)); + debug_assert!(lower <= upper); + + let interpolation_precision = deduce_interpolation_precision::(); + + let num = + T::Native::usize_as((fraction * interpolation_precision as f64) as usize); + let den = T::Native::usize_as(interpolation_precision); + + // Happy path: `upper - lower` does not overflow + // (could be a case for Decimal128 with max precision) + if let Ok(delta) = upper.sub_checked(lower) { + // Calculate the interpolation weight with the formula, where den is the precision: + // `lower + (upper - lower) * num / den` + let scaled: T::Native = scale_by_num::(delta, num, den)?; + lower.add_checked(scaled).map_err(|e| { + exec_datafusion_err!("Arithmetic overflow in percentile_cont: {e}") + }) + } else { + // Avoid overflow with the subtraction - split to two additive parts + // `a = lower * (precision-num) / precision` + // `b = upper * num / precision` + // The weights sum to 1, so the result is bounded by max(|lower|, |upper|) + // and never overflows, at the cost of a second truncation (2 ULP not 1). + let num_a = den.sub_wrapping(num); + let a: T::Native = scale_by_num::(lower, num_a, den)?; + + let b: T::Native = scale_by_num::(upper, num, den)?; + + a.add_checked(b).map_err(|e| { + exec_datafusion_err!("Arithmetic overflow in percentile_cont: {e}") + }) + } + } +} + /// Calculate the percentile value for a given set of values. /// This function performs an exact calculation by sorting all values. /// @@ -827,37 +1051,33 @@ where /// Note: This function takes a mutable slice and sorts it in place, but does not /// consume the data. This is important for window frame queries where evaluate() /// may be called multiple times on the same accumulator state. -fn calculate_percentile( +fn calculate_percentile>( values: &mut [T::Native], percentile: f64, -) -> Option -where - T::Native: Copy + AsPrimitive, - f64: AsPrimitive, -{ +) -> Result> { let cmp = |x: &T::Native, y: &T::Native| x.compare(*y); let len = values.len(); if len == 0 { - None + Ok(None) } else if len == 1 { - Some(values[0]) + Ok(Some(values[0])) } else if percentile == 0.0 { // Get minimum value - Some( + Ok(Some( *values .iter() .min_by(|a, b| cmp(a, b)) .expect("we checked for len > 0 a few lines above"), - ) + )) } else if percentile == 1.0 { // Get maximum value - Some( + Ok(Some( *values .iter() .max_by(|a, b| cmp(a, b)) .expect("we checked for len > 0 a few lines above"), - ) + )) } else { // Calculate the index using the formula: p * (n - 1) let index = percentile * ((len - 1) as f64); @@ -867,7 +1087,7 @@ where if lower_index == upper_index { // Exact index, return the value at that position let (_, value, _) = values.select_nth_unstable_by(lower_index, cmp); - Some(*value) + Ok(Some(*value)) } else { // Need to interpolate between two values // First, partition at lower_index to get the lower value @@ -878,20 +1098,11 @@ where let (_, upper_value, _) = values.select_nth_unstable_by(upper_index, cmp); let upper_value = *upper_value; - // Linear interpolation. - // We compute a quantized interpolation weight using `INTERPOLATION_PRECISION` because: - // 1. Both values come from the input data, so (upper - lower) is bounded by the value range - // 2. fraction is between 0 and 1; quantizing it provides stable, predictable results - // 3. The result is guaranteed to be between lower_value and upper_value (modulo cast rounding) - // 4. Arithmetic is performed in f64 and cast back to avoid overflowing Float16 intermediates let fraction = index - (lower_index as f64); - let scaled = (fraction * INTERPOLATION_PRECISION) as usize; - let weight = scaled as f64 / INTERPOLATION_PRECISION; - let lower_f: f64 = lower_value.as_(); - let upper_f: f64 = upper_value.as_(); - let interpolated_f = lower_f + (upper_f - lower_f) * weight; - Some(interpolated_f.as_()) + let interpolated = I::interpolate(lower_value, upper_value, fraction)?; + + Ok(Some(interpolated)) } } } @@ -900,11 +1111,17 @@ where mod tests { use super::*; use arrow::array::Float64Array; + use arrow::datatypes::{ + Decimal64Type, Float16Type, Float64Type, Int64Type, UInt32Type, + }; use half::f16; #[test] fn retract_batch_errors_on_untracked_value() { - let mut acc = PercentileContAccumulator::::new(0.5); + let mut acc = PercentileContAccumulator::::new( + 0.5, + DataType::Float64, + ); let values: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0])); acc.update_batch(std::slice::from_ref(&values)).unwrap(); @@ -932,11 +1149,19 @@ mod tests { Some(3.0), ])); - let mut dense_acc = PercentileContAccumulator::::new(0.5); + let mut dense_acc = + PercentileContAccumulator::::new( + 0.5, + DataType::Float64, + ); dense_acc .update_batch(std::slice::from_ref(&dense)) .unwrap(); - let mut sparse_acc = PercentileContAccumulator::::new(0.5); + let mut sparse_acc = + PercentileContAccumulator::::new( + 0.5, + DataType::Float64, + ); sparse_acc .update_batch(std::slice::from_ref(&sparse)) .unwrap(); @@ -950,8 +1175,10 @@ mod tests { // Interpolating between 0 and the max finite f16 value previously overflowed // intermediate f16 computations and produced NaN. let mut values = vec![f16::from_f32(0.0), f16::from_f32(65504.0)]; - let result = calculate_percentile::(&mut values, 0.5) - .expect("non-empty input"); + let result = + calculate_percentile::(&mut values, 0.5) + .expect("non-empty input") + .expect("non-empty result"); let result_f = result.to_f32(); assert!( !result_f.is_nan(), @@ -963,4 +1190,87 @@ mod tests { "unexpected result {result_f}" ); } + + #[test] + fn percentile_cont_decimal64() { + // Test values: [100.00, 200.00, 300.00, 400.00, 500.00] + // These are stored as i64 values scaled by 10^2 + let mut values = vec![ + 10000i64, // 100.00 + 20000i64, // 200.00 + 30000i64, // 300.00 + 40000i64, // 400.00 + 50000i64, // 500.00 + ]; + + // Test 50th percentile (median) + // Should return 300.00 (30000) + let result = + calculate_percentile::(&mut values, 0.5) + .expect("evaluate failed") + .expect("expected Some value"); + + assert_eq!(result, 30000i64, "50th percentile should be 300.00"); + + // Test 15th percentile + // Should return 160.00 (16000) + let result = + calculate_percentile::(&mut values, 0.15) + .expect("evaluate failed") + .expect("expected Some value"); + + assert_eq!(result, 16000i64, "15th percentile should be 160.00"); + + // Test 0th percentile (minimum) + let mut values = vec![10000i64, 20000i64, 30000i64]; + let result = + calculate_percentile::(&mut values, 0.0) + .expect("evaluate failed") + .expect("expected Some value"); + + assert_eq!( + result, 10000i64, + "0th percentile should be minimum value 100.00" + ); + + // Test 100th percentile (maximum) + let mut values = vec![10000i64, 20000i64, 30000i64]; + let result = + calculate_percentile::(&mut values, 1.0) + .expect("evaluate failed") + .expect("expected Some value"); + + assert_eq!( + result, 30000i64, + "100th percentile should be maximum value 300.00" + ); + } + + #[test] + fn percentile_cont_integer() { + // Cannot be tested with SLT since the integer coercion to float, + // but can affect dataframe use (e.g. `test_oom`) + let mut values = vec![1u32, 2, 3, 4]; + let result = + calculate_percentile::(&mut values, 0.5) + .expect("non-empty input") + .expect("non-empty result"); + assert_eq!(result, 2); + + // Interpolation rounds down rather than towards zero + let mut values = vec![-4i64, -3, -2, -1]; + let result = + calculate_percentile::(&mut values, 0.5) + .expect("non-empty input") + .expect("non-empty result"); + assert_eq!(result, -3); + + // No overflows + let mut values = vec![i64::MIN, i64::MAX]; + let result = + calculate_percentile::(&mut values, 0.5) + .expect("non-empty input") + .expect("non-empty result"); + assert_eq!(result, -1); + } } diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 460d4cd2ffda..e7eb43d41ace 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -1172,6 +1172,39 @@ FROM distinct_pct_nulls ORDER BY id; statement ok DROP TABLE distinct_pct_nulls; +# percentile_cont for shorter floats + +query RT +select percentile_cont(arrow_cast(col_f32, 'Float16'), 0.5), arrow_typeof(percentile_cont(arrow_cast(col_f32, 'Float16'), 0.5)) from median_table; +---- +2.75 Float16 + +query RT +select percentile_cont(col_f32, 0.5), arrow_typeof(percentile_cont(col_f32, 0.5)) from median_table; +---- +2.75 Float32 + +query RT +select percentile_cont(DISTINCT arrow_cast(col_f32, 'Float16'), 0.5), arrow_typeof(percentile_cont(DISTINCT arrow_cast(col_f32, 'Float16'), 0.5)) from median_table; +---- +2.75 Float16 + +query RT +select percentile_cont(DISTINCT col_f32, 0.5), arrow_typeof(percentile_cont(DISTINCT col_f32, 0.5)) from median_table; +---- +2.75 Float32 + +query RT +select median(arrow_cast(col_f32, 'Float16')), arrow_typeof(median(arrow_cast(col_f32, 'Float16'))) from median_table; +---- +2.75 Float16 + +query RT +select median(col_f32), arrow_typeof(median(col_f32)) from median_table; +---- +2.75 Float32 + + query RT select approx_median(arrow_cast(col_f32, 'Float16')), arrow_typeof(approx_median(arrow_cast(col_f32, 'Float16'))) from median_table; ---- @@ -4499,6 +4532,279 @@ SELECT percentile_cont(0.75) WITHIN GROUP (ORDER BY v DESC) FROM (VALUES (1), (2 2.75 +##################### +## percentile_cont tests for decimal types +##################### + +# Decimal128: typical percentiles (exact index, n=5) +statement ok +CREATE TABLE t (c DECIMAL(10, 2)) AS VALUES (10.10), (20.20), (30.30), (40.40), (50.50); + +query RT +SELECT percentile_cont(0.5) WITHIN GROUP (ORDER BY c), arrow_typeof(percentile_cont(0.5) WITHIN GROUP (ORDER BY c)) FROM t +---- +30.3 Decimal128(10, 2) + +query R +SELECT percentile_cont(0.25) WITHIN GROUP (ORDER BY c) FROM t +---- +20.2 + +query R +SELECT percentile_cont(0.75) WITHIN GROUP (ORDER BY c) FROM t +---- +40.4 + +query R +SELECT percentile_cont(0.0) WITHIN GROUP (ORDER BY c) FROM t +---- +10.1 + +query R +SELECT percentile_cont(1.0) WITHIN GROUP (ORDER BY c) FROM t +---- +50.5 + +# Alternate syntax should agree with WITHIN GROUP syntax +query B +SELECT percentile_cont(c, 0.5) = percentile_cont(0.5) WITHIN GROUP (ORDER BY c) FROM t +---- +true + +statement ok +DROP TABLE t; + +# Decimal128: interpolation between two values +statement ok +CREATE TABLE t (c DECIMAL(10, 2)) AS VALUES (10.00), (20.00); + +query RT +SELECT percentile_cont(0.5) WITHIN GROUP (ORDER BY c), arrow_typeof(percentile_cont(0.5) WITHIN GROUP (ORDER BY c)) FROM t +---- +15 Decimal128(10, 2) + +query R +SELECT percentile_cont(0.25) WITHIN GROUP (ORDER BY c) FROM t +---- +12.5 + +query R +SELECT percentile_cont(0.75) WITHIN GROUP (ORDER BY c) FROM t +---- +17.5 + +statement ok +DROP TABLE t; + +# Decimal128: negative numbers and NULL handling +statement ok +CREATE TABLE t (c DECIMAL(10, 2)) AS VALUES (-10.00), (-5.00), (0.00), (5.00), (10.00), (NULL); + +query R +SELECT percentile_cont(0.5) WITHIN GROUP (ORDER BY c) FROM t +---- +0 + +query R +SELECT percentile_cont(0.25) WITHIN GROUP (ORDER BY c) FROM t +---- +-5 + +statement ok +DROP TABLE t; + +# Decimal128: all NULLs +statement ok +CREATE TABLE t (c DECIMAL(10, 2)) AS VALUES (NULL); + +query R +SELECT percentile_cont(0.5) WITHIN GROUP (ORDER BY c) FROM t +---- +NULL + +statement ok +DROP TABLE t; + +# Decimal128: single value +statement ok +CREATE TABLE t (c DECIMAL(10, 2)) AS VALUES (42.00); + +query R +SELECT percentile_cont(0.5) WITHIN GROUP (ORDER BY c) FROM t +---- +42 + +statement ok +DROP TABLE t; + +# Decimal128: GROUP BY +statement ok +CREATE TABLE t (g INT, c DECIMAL(10, 2)) AS VALUES (1, 10.00), (1, 20.00), (2, 30.00), (2, 40.00), (2, 50.00); + +query IR +SELECT g, percentile_cont(0.5) WITHIN GROUP (ORDER BY c) FROM t GROUP BY g ORDER BY g +---- +1 15 +2 40 + +statement ok +DROP TABLE t; + +# Decimal128: DISTINCT +statement ok +CREATE TABLE t (c DECIMAL(10, 2)) AS VALUES (10.00), (10.00), (10.00), (10.00), (20.00); + +query R +SELECT percentile_cont(0.5) WITHIN GROUP (ORDER BY c) FROM t +---- +10 + +query RT +SELECT percentile_cont(DISTINCT c, 0.5), arrow_typeof(percentile_cont(DISTINCT c, 0.5)) FROM t +---- +15 Decimal128(10, 2) + +statement ok +DROP TABLE t; + +# Other decimals + +query RT +SELECT percentile_cont(0.5) WITHIN GROUP (ORDER BY v), arrow_typeof(percentile_cont(0.5) WITHIN GROUP (ORDER BY v)) +FROM (VALUES (arrow_cast(10.00, 'Decimal32(9, 2)')), (arrow_cast(20.00, 'Decimal32(9, 2)'))) as t (v) +---- +15 Decimal32(9, 2) + +query RT +SELECT percentile_cont(0.5) WITHIN GROUP (ORDER BY v), arrow_typeof(percentile_cont(0.5) WITHIN GROUP (ORDER BY v)) +FROM (VALUES (arrow_cast(10.00, 'Decimal64(18, 2)')), (arrow_cast(20.00, 'Decimal64(18, 2)'))) as t (v) +---- +15 Decimal64(18, 2) + +query RT +SELECT percentile_cont(0.5) WITHIN GROUP (ORDER BY v), arrow_typeof(percentile_cont(0.5) WITHIN GROUP (ORDER BY v)) +FROM (VALUES (arrow_cast(10.00, 'Decimal256(54, 6)')), (arrow_cast(20.00, 'Decimal256(54, 6)'))) as t (v) +---- +15 Decimal256(54, 6) + +# Other decimals: GROUP BY and DISTINCT + +statement ok +CREATE TABLE t (g INT, c DECIMAL(10, 2)) AS VALUES + (1, 10.00), (1, 10.00), (1, 10.00), (1, 20.00), (2, 30.00), (2, 40.00); + +query IRT +SELECT g, + percentile_cont(arrow_cast(c, 'Decimal32(9, 2)'), 0.5), + arrow_typeof(percentile_cont(arrow_cast(c, 'Decimal32(9, 2)'), 0.5)) +FROM t GROUP BY g ORDER BY g +---- +1 10 Decimal32(9, 2) +2 35 Decimal32(9, 2) + +query RRT +SELECT + percentile_cont(arrow_cast(c, 'Decimal32(9, 2)'), 0.5), + percentile_cont(DISTINCT arrow_cast(c, 'Decimal32(9, 2)'), 0.5), + arrow_typeof(percentile_cont(DISTINCT arrow_cast(c, 'Decimal32(9, 2)'), 0.5)) +FROM t +---- +15 25 Decimal32(9, 2) + +query IRT +SELECT g, + percentile_cont(arrow_cast(c, 'Decimal64(18, 2)'), 0.5), + arrow_typeof(percentile_cont(arrow_cast(c, 'Decimal64(18, 2)'), 0.5)) +FROM t GROUP BY g ORDER BY g +---- +1 10 Decimal64(18, 2) +2 35 Decimal64(18, 2) + +query RRT +SELECT + percentile_cont(arrow_cast(c, 'Decimal64(18, 2)'), 0.5), + percentile_cont(DISTINCT arrow_cast(c, 'Decimal64(18, 2)'), 0.5), + arrow_typeof(percentile_cont(DISTINCT arrow_cast(c, 'Decimal64(18, 2)'), 0.5)) +FROM t +---- +15 25 Decimal64(18, 2) + +query IRT +SELECT g, + percentile_cont(arrow_cast(c, 'Decimal256(54, 6)'), 0.5), + arrow_typeof(percentile_cont(arrow_cast(c, 'Decimal256(54, 6)'), 0.5)) +FROM t GROUP BY g ORDER BY g +---- +1 10 Decimal256(54, 6) +2 35 Decimal256(54, 6) + +query RRT +SELECT + percentile_cont(arrow_cast(c, 'Decimal256(54, 6)'), 0.5), + percentile_cont(DISTINCT arrow_cast(c, 'Decimal256(54, 6)'), 0.5), + arrow_typeof(percentile_cont(DISTINCT arrow_cast(c, 'Decimal256(54, 6)'), 0.5)) +FROM t +---- +15 25 Decimal256(54, 6) + +statement ok +DROP TABLE t; + +statement ok +CREATE TABLE t (c DECIMAL(50, 2)) AS VALUES (10.00), (20.00); + +query RT +SELECT percentile_cont(0.5) WITHIN GROUP (ORDER BY c), arrow_typeof(percentile_cont(0.5) WITHIN GROUP (ORDER BY c)) FROM t +---- +15 Decimal256(50, 2) + +statement ok +DROP TABLE t; + +# Decimal128 overflow cases +statement ok +CREATE TABLE t (c DECIMAL(38, 0)) AS VALUES + (arrow_cast('-90000000000000000000000000000000000000', 'Decimal128(38,0)')), + (arrow_cast('90000000000000000000000000000000000000', 'Decimal128(38,0)')); + +query RT +SELECT percentile_cont(0.25) WITHIN GROUP (ORDER BY c), arrow_typeof(percentile_cont(0.25) WITHIN GROUP (ORDER BY c)) FROM t +---- +-45000000000000000000000000000000000000 Decimal128(38, 0) + +statement ok +DROP TABLE t; + +statement ok +CREATE TABLE t (c DECIMAL(38, 0)) AS VALUES + (arrow_cast('-99999999999999999999999999999999999999', 'Decimal128(38,0)')), + (arrow_cast('99999999999999999999999999999999999999', 'Decimal128(38,0)')); + +query RT +SELECT percentile_cont(0.5) WITHIN GROUP (ORDER BY c), arrow_typeof(percentile_cont(0.5) WITHIN GROUP (ORDER BY c)) FROM t +---- +0 Decimal128(38, 0) + +statement ok +DROP TABLE t; + +# Decimal32 overflow cases +query RT +SELECT percentile_cont(0.5) WITHIN GROUP (ORDER BY v), arrow_typeof(percentile_cont(0.5) WITHIN GROUP (ORDER BY v)) +FROM (VALUES (arrow_cast(-999999999, 'Decimal32(9, 0)')), (arrow_cast(999999999, 'Decimal32(9, 0)'))) as t (v) +---- +0 Decimal32(9, 0) + +# A fraction lower than `1 / den` returns zero numerator, so the interpolation +# floors to the lower value +query RT +SELECT percentile_cont(0.0000001) WITHIN GROUP (ORDER BY v), arrow_typeof(percentile_cont(0.0000001) WITHIN GROUP (ORDER BY v)) +FROM (VALUES (arrow_cast(10.00, 'Decimal64(18, 2)')), (arrow_cast(20.00, 'Decimal64(18, 2)'))) as t (v) +---- +10 Decimal64(18, 2) + + + # variance_single_value query RRRR select var(sq.column1), var_pop(sq.column1), stddev(sq.column1), stddev_pop(sq.column1) from (values (1.0)) as sq; @@ -8817,10 +9123,11 @@ NULL NULL NULL NULL statement ok drop table distinct_avg; -query R -select percentile_cont(null, 0.5); +# Null evaluates to null +query ?T +select percentile_cont(null, 0.5), arrow_typeof(percentile_cont(null, 0.5)); ---- -NULL +NULL Null # Test string_agg window frame behavior (fix for issue #19612) statement ok