diff --git a/datafusion/functions-aggregate/benches/percentile_cont.rs b/datafusion/functions-aggregate/benches/percentile_cont.rs index 17723ac1c2cc..a6f29c9174db 100644 --- a/datafusion/functions-aggregate/benches/percentile_cont.rs +++ b/datafusion/functions-aggregate/benches/percentile_cont.rs @@ -24,7 +24,7 @@ use criterion::{BatchSize, Criterion, criterion_group, criterion_main}; use datafusion_expr::function::AccumulatorArgs; use datafusion_expr::groups_accumulator::GroupsAccumulator; use datafusion_expr::{Accumulator, AggregateUDFImpl}; -use datafusion_functions_aggregate::approx_percentile_cont::ApproxPercentileAccumulator; +use datafusion_functions_aggregate::approx_percentile_cont::ApproxPercentileCont; use datafusion_functions_aggregate::percentile_cont::PercentileCont; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::GroupsAccumulatorAdapter; use datafusion_physical_expr::expressions::{col, lit}; @@ -54,6 +54,28 @@ fn prepare_accumulator() -> Box { PercentileCont::new().accumulator(accumulator_args).unwrap() } +fn prepare_approx_accumulator() -> Box { + let schema = Arc::new(Schema::new(vec![Field::new("f", DataType::Float64, true)])); + let value_expr = col("f", &schema).unwrap(); + let percentile_expr = lit(0.5_f64); + let value_field = value_expr.return_field(&schema).unwrap(); + let percentile_field = percentile_expr.return_field(&schema).unwrap(); + let accumulator_args = AccumulatorArgs { + return_field: Field::new("f", DataType::Float64, true).into(), + schema: &schema, + expr_fields: &[value_field, percentile_field], + ignore_nulls: false, + order_bys: &[], + is_reversed: false, + name: "approx_percentile_cont(f, 0.5)", + is_distinct: false, + exprs: &[value_expr, percentile_expr], + }; + ApproxPercentileCont::new() + .accumulator(accumulator_args) + .unwrap() +} + fn stream_array(len: usize, null_stride: Option) -> ArrayRef { let values = (0..len) .map(|idx| { @@ -111,14 +133,7 @@ fn grouped_adapter_benchmark(c: &mut Criterion) { "approx_percentile_cont/groups_accumulator_adapter/8192_groups", |b| { b.iter_batched( - || { - GroupsAccumulatorAdapter::new(|| { - Ok(Box::new(ApproxPercentileAccumulator::new( - 0.5, - DataType::Float64, - )) as Box) - }) - }, + || GroupsAccumulatorAdapter::new(|| Ok(prepare_approx_accumulator())), |mut accumulator| { accumulator .update_batch( diff --git a/datafusion/functions-aggregate/src/approx_median.rs b/datafusion/functions-aggregate/src/approx_median.rs index ddaa883a3fd6..3ea0a0408e56 100644 --- a/datafusion/functions-aggregate/src/approx_median.rs +++ b/datafusion/functions-aggregate/src/approx_median.rs @@ -34,6 +34,7 @@ use datafusion_expr::{ use datafusion_macros::user_doc; use crate::approx_percentile_cont::ApproxPercentileAccumulator; +use crate::utils::{PercentileParam, PercentileParamState}; make_udaf_expr_and_func!( ApproxMedian, @@ -108,6 +109,7 @@ impl AggregateUDFImpl for ApproxMedian { Field::new_list_field(Float64, true), false, ), + Field::new(format_state_name(args.name, "percentile"), Float64, true), ] .into_iter() .map(Arc::new) @@ -138,7 +140,11 @@ impl AggregateUDFImpl for ApproxMedian { Ok(Box::new(NoopAccumulator::default())) } else { Ok(Box::new(ApproxPercentileAccumulator::new( - 0.5_f64, + PercentileParam { + aggregate_fn_name: "APPROX_MEDIAN".to_string(), + state: PercentileParamState::Resolved(0.5_f64), + is_desc: false, + }, acc_args.expr_fields[0].data_type().clone(), ))) } diff --git a/datafusion/functions-aggregate/src/approx_percentile_cont.rs b/datafusion/functions-aggregate/src/approx_percentile_cont.rs index 5c8441aec092..3e227f71c972 100644 --- a/datafusion/functions-aggregate/src/approx_percentile_cont.rs +++ b/datafusion/functions-aggregate/src/approx_percentile_cont.rs @@ -28,7 +28,7 @@ use arrow::{ }; use datafusion_common::types::{NativeType, logical_float64}; use datafusion_common::{ - DataFusionError, Result, ScalarValue, downcast_value, internal_err, not_impl_err, + Result, ScalarValue, downcast_value, internal_err, not_impl_err, plan_datafusion_err, plan_err, }; use datafusion_expr::DistinctHandling; @@ -43,10 +43,13 @@ use datafusion_functions_aggregate_common::tdigest::{DEFAULT_MAX_SIZE, TDigest}; use datafusion_macros::user_doc; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; -use crate::utils::{get_scalar_value, validate_percentile_expr}; +use crate::utils::{PercentileParam, get_scalar_value}; create_func!(ApproxPercentileCont, approx_percentile_cont_udaf); +// state field index for the percentile argument +pub const STATE_PERCENTILE_IDX: usize = 6; + /// Computes the approximate percentile continuous of a set of numbers pub fn approx_percentile_cont( order_by: Sort, @@ -172,20 +175,17 @@ impl ApproxPercentileCont { &self, args: &AccumulatorArgs, ) -> Result { - let percentile = - validate_percentile_expr(&args.exprs[1], "APPROX_PERCENTILE_CONT")?; - let is_descending = args .order_bys .first() .map(|sort_expr| sort_expr.options.descending) .unwrap_or(false); - let percentile = if is_descending { - 1.0 - percentile - } else { - percentile - }; + let percentile = PercentileParam::try_new( + &args.exprs[1], + "APPROX_PERCENTILE_CONT", + is_descending, + )?; let tdigest_max_size = if args.exprs.len() == 3 { Some(validate_input_max_size_expr(&args.exprs[2])?) @@ -219,9 +219,8 @@ impl ApproxPercentileCont { fn validate_input_max_size_expr(expr: &Arc) -> Result { let scalar_value = get_scalar_value(expr).map_err(|_e| { - DataFusionError::Plan( + plan_datafusion_err!( "Tdigest max_size value for 'APPROX_PERCENTILE_CONT' must be a literal" - .to_string(), ) })?; @@ -280,6 +279,11 @@ impl AggregateUDFImpl for ApproxPercentileCont { Field::new_list_field(DataType::Float64, true), false, ), + Field::new( + format_state_name(args.name, "percentile"), + DataType::Float64, + true, + ), ] .into_iter() .map(Arc::new) @@ -337,12 +341,12 @@ impl AggregateUDFImpl for ApproxPercentileCont { #[derive(Debug)] pub struct ApproxPercentileAccumulator { digest: TDigest, - percentile: f64, + percentile: PercentileParam, return_type: DataType, } impl ApproxPercentileAccumulator { - pub fn new(percentile: f64, return_type: DataType) -> Self { + pub fn new(percentile: PercentileParam, return_type: DataType) -> Self { Self { digest: TDigest::new(DEFAULT_MAX_SIZE), percentile, @@ -350,8 +354,8 @@ impl ApproxPercentileAccumulator { } } - pub fn new_with_max_size( - percentile: f64, + pub(crate) fn new_with_max_size( + percentile: PercentileParam, return_type: DataType, max_size: usize, ) -> Self { @@ -362,6 +366,15 @@ impl ApproxPercentileAccumulator { } } + /// Callers must only invoke this when a percentile argument actually + /// exists in the current batch. + pub(crate) fn resolve_percentile( + &mut self, + percentile_array: &ArrayRef, + ) -> Result<()> { + self.percentile.resolve(percentile_array) + } + // pub(crate) for approx_percentile_cont_with_weight pub(crate) fn max_size(&self) -> usize { self.digest.max_size() @@ -405,10 +418,17 @@ impl ApproxPercentileAccumulator { impl Accumulator for ApproxPercentileAccumulator { fn state(&mut self) -> Result> { - Ok(self.digest.to_scalar_state().into_iter().collect()) + let mut state: Vec = + self.digest.to_scalar_state().into_iter().collect(); + state.push(ScalarValue::Float64(self.percentile.get().ok())); + Ok(state) } fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + if values.len() > 1 { + self.resolve_percentile(&values[1])?; + } + // Remove any nulls before computing the percentile let mut values = Arc::clone(&values[0]); if values.null_count() > 0 { @@ -424,7 +444,9 @@ impl Accumulator for ApproxPercentileAccumulator { if self.digest.count() == 0.0 { return ScalarValue::try_from(self.return_type.clone()); } - let q = self.digest.estimate_quantile(self.percentile); + let q = self + .digest + .estimate_quantile(self.percentile.effective_percentile()?); // These acceptable return types MUST match the validation in // ApproxPercentile::create_accumulator. @@ -441,9 +463,14 @@ impl Accumulator for ApproxPercentileAccumulator { return Ok(()); } - let states = (0..states[0].len()) + if let Some(percentile_array) = states.get(STATE_PERCENTILE_IDX) { + self.percentile.resolve(percentile_array)?; + } + + let tdigest_states = &states[..STATE_PERCENTILE_IDX]; + let states = (0..tdigest_states[0].len()) .map(|index| { - states + tdigest_states .iter() .map(|array| ScalarValue::try_from_array(array, index)) .collect::>>() @@ -470,6 +497,19 @@ mod tests { use datafusion_functions_aggregate_common::tdigest::TDigest; use crate::approx_percentile_cont::ApproxPercentileAccumulator; + use crate::utils::{PercentileParam, PercentileParamState}; + + fn make_accumulator() -> ApproxPercentileAccumulator { + ApproxPercentileAccumulator::new_with_max_size( + PercentileParam { + aggregate_fn_name: "APPROX_PERCENTILE_CONT".to_string(), + state: PercentileParamState::Resolved(0.5), + is_desc: false, + }, + DataType::Float64, + 100, + ) + } #[test] fn test_combine_approx_percentile_accumulator() { @@ -486,8 +526,7 @@ mod tests { let t1 = TDigest::merge_digests(&digests); let t2 = TDigest::merge_digests(&digests); - let mut accumulator = - ApproxPercentileAccumulator::new_with_max_size(0.5, DataType::Float64, 100); + let mut accumulator = make_accumulator(); accumulator.merge_digests(&[t1]); assert_eq!(accumulator.digest.count(), 50_000.0); diff --git a/datafusion/functions-aggregate/src/approx_percentile_cont_with_weight.rs b/datafusion/functions-aggregate/src/approx_percentile_cont_with_weight.rs index 1a810a3c3668..7fd68d2e380d 100644 --- a/datafusion/functions-aggregate/src/approx_percentile_cont_with_weight.rs +++ b/datafusion/functions-aggregate/src/approx_percentile_cont_with_weight.rs @@ -308,6 +308,9 @@ impl Accumulator for ApproxPercentileWithWeightAccumulator { } fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + self.approx_percentile_cont_accumulator + .resolve_percentile(&values[2])?; + let mut means = Arc::clone(&values[0]); let mut weights = Arc::clone(&values[1]); // If nulls are present in either array, need to filter those rows out in both arrays diff --git a/datafusion/functions-aggregate/src/median.rs b/datafusion/functions-aggregate/src/median.rs index c96baddfee81..97d663a6c206 100644 --- a/datafusion/functions-aggregate/src/median.rs +++ b/datafusion/functions-aggregate/src/median.rs @@ -24,6 +24,7 @@ use arrow::datatypes::FieldRef; use crate::percentile_cont::{ PercentileCont, create_percentile_accumulator, create_percentile_groups_accumulator, }; +use crate::utils::{PercentileParam, PercentileParamState}; use datafusion_common::Result; use datafusion_common::types::logical_float64; use datafusion_expr::GroupsAccumulator; @@ -123,7 +124,11 @@ impl AggregateUDFImpl for Median { fn accumulator(&self, args: AccumulatorArgs) -> Result> { create_percentile_accumulator( self.name(), - 0.5, + PercentileParam { + aggregate_fn_name: "MEDIAN".to_string(), + state: PercentileParamState::Resolved(0.5), + is_desc: false, + }, args.expr_fields[0].data_type(), args.is_distinct, ) @@ -139,7 +144,11 @@ impl AggregateUDFImpl for Median { ) -> Result> { create_percentile_groups_accumulator( self.name(), - 0.5, + PercentileParam { + aggregate_fn_name: "MEDIAN".to_string(), + state: PercentileParamState::Resolved(0.5), + is_desc: false, + }, args.expr_fields[0].data_type(), ) } diff --git a/datafusion/functions-aggregate/src/percentile_cont.rs b/datafusion/functions-aggregate/src/percentile_cont.rs index 5b1f6bdc1f6f..ac3a925b5596 100644 --- a/datafusion/functions-aggregate/src/percentile_cont.rs +++ b/datafusion/functions-aggregate/src/percentile_cont.rs @@ -46,7 +46,6 @@ use datafusion_common::utils::memory::estimate_memory_size; use datafusion_functions_aggregate_common::noop_accumulator::NoopAccumulator; use crate::min_max::{max_udaf, min_udaf}; -use crate::utils::validate_percentile_expr; use datafusion_common::{ Result, ScalarValue, exec_datafusion_err, internal_datafusion_err, utils::{SingleRowListArrayBuilder, take_function_args}, @@ -67,6 +66,8 @@ use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls: use datafusion_functions_aggregate_common::utils::Hashable; use datafusion_macros::user_doc; +use crate::utils::PercentileParam; + /// Precision multiplier for linear interpolation calculations. /// /// This value of 1,000,000 was chosen to balance precision with overflow safety: @@ -227,12 +228,18 @@ impl AggregateUDFImpl for PercentileCont { true, ) .into(), + Field::new( + format_state_name(args.name, "percentile"), + DataType::Float64, + true, + ) + .into(), ]) } fn accumulator(&self, args: AccumulatorArgs) -> Result> { // Always verify percentiles - let percentile = get_percentile(&args)?; + let percentile = get_percentile_param(&args)?; create_percentile_accumulator( self.name(), percentile, @@ -250,7 +257,7 @@ impl AggregateUDFImpl for PercentileCont { args: AccumulatorArgs, ) -> Result> { // Always verify percentiles - let percentile = get_percentile(&args)?; + let percentile = get_percentile_param(&args)?; create_percentile_groups_accumulator( self.name(), percentile, @@ -273,27 +280,21 @@ impl AggregateUDFImpl for PercentileCont { } } -fn get_percentile(args: &AccumulatorArgs) -> Result { - let percentile = validate_percentile_expr(&args.exprs[1], "PERCENTILE_CONT")?; - +fn get_percentile_param(args: &AccumulatorArgs) -> Result { let is_descending = args .order_bys .first() .map(|sort_expr| sort_expr.options.descending) .unwrap_or(false); - let percentile = if is_descending { - 1.0 - percentile - } else { - percentile - }; - + let percentile = + PercentileParam::try_new(&args.exprs[1], "PERCENTILE_CONT", is_descending)?; Ok(percentile) } pub fn create_percentile_accumulator( name: &str, - percentile: f64, + percentile: PercentileParam, input_dt: &DataType, is_distinct: bool, ) -> Result> { @@ -339,7 +340,7 @@ pub fn create_percentile_accumulator( pub fn create_percentile_groups_accumulator( name: &str, - percentile: f64, + percentile: PercentileParam, input_dt: &DataType, ) -> Result> { macro_rules! helper { @@ -442,7 +443,7 @@ struct PercentileContAccumulator< I: PercentileInterpolator, > { all_values: Vec, - percentile: f64, + percentile: PercentileParam, data_type: DataType, _interpolator: PhantomData, } @@ -450,7 +451,7 @@ struct PercentileContAccumulator< impl> PercentileContAccumulator { - fn new(percentile: f64, data_type: DataType) -> Self { + fn new(percentile: PercentileParam, data_type: DataType) -> Self { Self { all_values: vec![], percentile, @@ -487,10 +488,17 @@ where None, ); - Ok(vec![ScalarValue::List(Arc::new(list_array))]) + Ok(vec![ + ScalarValue::List(Arc::new(list_array)), + ScalarValue::Float64(self.percentile.get().ok()), + ]) } fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + if values.len() > 1 { + self.percentile.resolve(&values[1])?; + } + let values = values[0].as_primitive::(); let additional = values.len() - values.null_count(); self.all_values.try_reserve(additional).map_err(|e| { @@ -508,6 +516,9 @@ where } fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + if let Some(percentile_array) = states.get(1) { + self.percentile.resolve(percentile_array)?; + } let array = states[0].as_list::(); // Feed all list elements from a batch for values in array.iter().flatten() { @@ -517,7 +528,11 @@ where } fn evaluate(&mut self) -> Result { - let value = calculate_percentile::(&mut self.all_values, self.percentile)?; + if self.all_values.is_empty() { + return ScalarValue::new_primitive::(None, &self.data_type); + } + let percentile = self.percentile.effective_percentile()?; + let value = calculate_percentile::(&mut self.all_values, percentile)?; ScalarValue::new_primitive::(value, &self.data_type) } @@ -588,7 +603,11 @@ struct PercentileContGroupsAccumulator< I: PercentileInterpolator, > { group_values: Vec>, - percentile: f64, + percentile: PercentileParam, + /// Per-group percentiles, only used when the percentile is given as a + /// column reference: it then has to be constant within each group, but + /// may differ between groups. + group_percentiles: Option>>, data_type: DataType, _interpolator: PhantomData, } @@ -596,14 +615,46 @@ struct PercentileContGroupsAccumulator< impl> PercentileContGroupsAccumulator { - fn new(percentile: f64, data_type: DataType) -> Self { + fn new(percentile: PercentileParam, data_type: DataType) -> Self { + let group_percentiles = percentile.is_col_ref().then(Vec::new); Self { group_values: vec![], percentile, + group_percentiles, data_type, _interpolator: PhantomData, } } + + /// Resolves the percentile of each group from the `percentiles` argument, + /// either per group or, for a literal percentile, once for all groups. + fn resolve_percentiles( + &mut self, + percentiles: &ArrayRef, + group_indices: &[usize], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + let Some(group_percentiles) = self.group_percentiles.as_mut() else { + return self.percentile.resolve(percentiles); + }; + group_percentiles.resize(total_num_groups, None); + + let percentiles = self.percentile.to_float64_array(percentiles)?; + for (row, (&group_index, value)) in + group_indices.iter().zip(percentiles.iter()).enumerate() + { + let Some(value) = value else { continue }; + if opt_filter.is_some_and(|f| !f.is_valid(row) || !f.value(row)) { + continue; + } + group_percentiles[group_index] = Some( + self.percentile + .resolve_group_value(group_percentiles[group_index], value)?, + ); + } + Ok(()) + } } impl GroupsAccumulator for PercentileContGroupsAccumulator @@ -618,8 +669,14 @@ where opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { - // For ordered-set aggregates, we only care about the ORDER BY column (first element) - // The percentile parameter is already stored in self.percentile + if values.len() > 1 { + self.resolve_percentiles( + &values[1], + group_indices, + opt_filter, + total_num_groups, + )?; + } let values = values[0].as_primitive::(); @@ -643,7 +700,14 @@ where group_indices: &[usize], total_num_groups: usize, ) -> Result<()> { - assert_eq!(values.len(), 1, "one argument to merge_batch"); + if let Some(percentile_array) = values.get(1) { + self.resolve_percentiles( + percentile_array, + group_indices, + None, + total_num_groups, + )?; + } let input_group_values = values[0].as_list::(); @@ -667,6 +731,7 @@ where fn state(&mut self, emit_to: EmitTo) -> Result> { // Emit values let emit_group_values = emit_to.take_needed(&mut self.group_values); + let num_groups = emit_group_values.len(); // Build offsets let mut offsets = Vec::with_capacity(self.group_values.len() + 1); @@ -693,7 +758,19 @@ where None, ); - Ok(vec![Arc::new(result_list_array)]) + let percentile_array: ArrayRef = match self.group_percentiles.as_mut() { + Some(group_percentiles) => Arc::new(arrow::array::Float64Array::from( + emit_to.take_needed(group_percentiles), + )), + None => Arc::new(arrow::array::Float64Array::from(vec![ + self.percentile + .get() + .ok(); + num_groups + ])), + }; + + Ok(vec![Arc::new(result_list_array), percentile_array]) } fn evaluate(&mut self, emit_to: EmitTo) -> Result { @@ -704,10 +781,33 @@ where let mut evaluate_result_builder = 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)?; - evaluate_result_builder.append_option(value); + match self.group_percentiles.as_mut() { + Some(group_percentiles) => { + let mut emit_group_percentiles = emit_to.take_needed(group_percentiles); + emit_group_percentiles.resize(emit_group_values.len(), None); + for (values, percentile) in + emit_group_values.iter_mut().zip(emit_group_percentiles) + { + if values.is_empty() { + evaluate_result_builder.append_null(); + continue; + } + let percentile = self + .percentile + .apply_order(self.percentile.get_group_value(percentile)?); + let value = + calculate_percentile::(values.as_mut_slice(), percentile)?; + evaluate_result_builder.append_option(value); + } + } + None => { + let percentile = self.percentile.effective_percentile()?; + for values in &mut emit_group_values { + let value = + calculate_percentile::(values.as_mut_slice(), percentile)?; + evaluate_result_builder.append_option(value); + } + } } Ok(Arc::new(evaluate_result_builder.finish())) @@ -729,8 +829,8 @@ where // 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) + // Reuse values buffer in `input_array` to build `list_values` in `ListArray` + let list_values = PrimitiveArray::::new(input_array.values().clone(), None) .with_data_type(self.data_type.clone()); // `offsets` in `ListArray`, each row as a list element @@ -754,11 +854,32 @@ where let converted_list_array = ListArray::new( Arc::new(Field::new_list_field(self.data_type.clone(), true)), offsets, - Arc::new(values), + Arc::new(list_values), nulls, ); - Ok(vec![Arc::new(converted_list_array)]) + // `state_fields` always declares a second (percentile) state field, so + // `convert_to_state` must also produce one array per row to match it. + // When the percentile is given as a column reference, propagate its + // per-row values (nulled out for filtered rows); otherwise broadcast + // the resolved constant. + let percentile_array: ArrayRef = if let Some(percentile_values) = values.get(1) { + let percentile_values = + self.percentile.to_float64_array(percentile_values)?; + let nulls = filtered_null_mask(opt_filter, &percentile_values); + Arc::new(arrow::array::Float64Array::new( + percentile_values.values().clone(), + nulls, + )) + } else { + let percentile = self.percentile.get().ok(); + Arc::new(arrow::array::Float64Array::from(vec![ + percentile; + input_array.len() + ])) + }; + + Ok(vec![Arc::new(converted_list_array), percentile_array]) } fn size(&self) -> usize { self.group_values @@ -767,6 +888,10 @@ where .sum::() // account for size of self.group_values too + self.group_values.capacity() * size_of::>() + + self + .group_percentiles + .as_ref() + .map_or(0, |p| p.capacity() * size_of::>()) } } @@ -788,7 +913,7 @@ struct DistinctPercentileContAccumulator< /// `GenericDistinctBuffer` rather than the standard library's default /// SipHash, which is considerably slower for this hot path. counts: HashMap, usize, RandomState>, - percentile: f64, + percentile: PercentileParam, data_type: DataType, _interpolator: PhantomData, } @@ -796,7 +921,7 @@ struct DistinctPercentileContAccumulator< impl> DistinctPercentileContAccumulator { - fn new(percentile: f64, data_type: DataType) -> Self { + fn new(percentile: PercentileParam, data_type: DataType) -> Self { Self { counts: HashMap::default(), percentile, @@ -822,12 +947,15 @@ where ); Ok(vec![ SingleRowListArrayBuilder::new(arr).build_list_scalar(), + ScalarValue::Float64(self.percentile.get().ok()), ]) } fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - // `values` may carry extra argument columns (e.g. the percentile - // literal); only the first column holds the aggregated values. + if values.len() > 1 { + self.percentile.resolve(&values[1])?; + } + let arr = values[0].as_primitive::(); if arr.null_count() > 0 { for value in arr.iter().flatten() { @@ -843,6 +971,9 @@ where } fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + if let Some(percentile_array) = states.get(1) { + self.percentile.resolve(percentile_array)?; + } let list = states[0].as_list::(); for values in list.iter().flatten() { let arr = values.as_primitive::(); @@ -855,7 +986,11 @@ 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)?; + if values.is_empty() { + return ScalarValue::new_primitive::(None, &self.data_type); + } + let percentile = self.percentile.effective_percentile()?; + let value = calculate_percentile::(&mut values, percentile)?; ScalarValue::new_primitive::(value, &self.data_type) } @@ -1102,16 +1237,25 @@ fn calculate_percentile>( #[cfg(test)] mod tests { use super::*; + use crate::utils::PercentileParamState; use arrow::array::Float64Array; use arrow::datatypes::{Decimal64Type, Decimal128Type, Float16Type, Float64Type}; use half::f16; + fn create_accumulator() -> PercentileContAccumulator { + PercentileContAccumulator::::new( + PercentileParam { + aggregate_fn_name: "PERCENTILE_CONT".to_string(), + state: PercentileParamState::Resolved(0.5), + is_desc: false, + }, + DataType::Float64, + ) + } + #[test] fn retract_batch_errors_on_untracked_value() { - let mut acc = PercentileContAccumulator::::new( - 0.5, - DataType::Float64, - ); + let mut acc = create_accumulator(); let values: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0])); acc.update_batch(std::slice::from_ref(&values)).unwrap(); @@ -1139,19 +1283,11 @@ mod tests { Some(3.0), ])); - let mut dense_acc = - PercentileContAccumulator::::new( - 0.5, - DataType::Float64, - ); + let mut dense_acc = create_accumulator(); dense_acc .update_batch(std::slice::from_ref(&dense)) .unwrap(); - let mut sparse_acc = - PercentileContAccumulator::::new( - 0.5, - DataType::Float64, - ); + let mut sparse_acc = create_accumulator(); sparse_acc .update_batch(std::slice::from_ref(&sparse)) .unwrap(); diff --git a/datafusion/functions-aggregate/src/utils.rs b/datafusion/functions-aggregate/src/utils.rs index 6d816e54bdaf..71d98e6d3e77 100644 --- a/datafusion/functions-aggregate/src/utils.rs +++ b/datafusion/functions-aggregate/src/utils.rs @@ -17,9 +17,13 @@ use std::sync::Arc; -use arrow::array::RecordBatch; -use arrow::datatypes::Schema; -use datafusion_common::{DataFusionError, Result, ScalarValue, internal_err, plan_err}; +use arrow::array::{Array, ArrayRef, Float32Array, Float64Array, RecordBatch}; +use arrow::compute::{cast, max, min}; +use arrow::datatypes::{DataType, Schema}; +use datafusion_common::{ + Result, ScalarValue, downcast_value, exec_datafusion_err, exec_err, + internal_datafusion_err, internal_err, plan_err, +}; use datafusion_expr::ColumnarValue; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; @@ -37,37 +41,29 @@ pub(crate) fn get_scalar_value(expr: &Arc) -> Result, - fn_name: &str, -) -> Result { - let scalar_value = get_scalar_value(expr).map_err(|_e| { - DataFusionError::Plan(format!( - "Percentile value for '{fn_name}' must be a literal" - )) - })?; - +/// Validates that a percentile scalar is a Float32/Float64 value between 0.0 and 1.0. +fn scalar_to_percentile(scalar_value: ScalarValue, fn_name: &str) -> Result { let percentile = match scalar_value { ScalarValue::Float32(Some(value)) => value as f64, ScalarValue::Float64(Some(value)) => value, ScalarValue::Float32(None) | ScalarValue::Float64(None) => { return plan_err!( - "Percentile value for '{fn_name}' must be Float32 or Float64 literal (got null)" + "Percentile value for '{fn_name}' must be Float32 or Float64 (got null)" ); } sv => { return plan_err!( - "Percentile value for '{fn_name}' must be Float32 or Float64 literal (got data type {})", + "Percentile value for '{fn_name}' must be Float32 or Float64 (got data type {})", sv.data_type() ); } }; - // Ensure the percentile is between 0 and 1. + check_percentile_range(percentile) +} + +/// Ensures the percentile is between 0 and 1. +fn check_percentile_range(percentile: f64) -> Result { if !(0.0..=1.0).contains(&percentile) { return plan_err!( "Percentile value must be between 0.0 and 1.0 inclusive, {percentile} is invalid" @@ -75,3 +71,210 @@ pub(crate) fn validate_percentile_expr( } Ok(percentile) } + +/// State of the PercentileParam resolution. +/// Either already resolved or still requiring a non-empty record batch. +#[derive(Debug, Clone)] +pub(crate) enum PercentileParamState { + Resolved(f64), + Pending, +} + +/// Percentile argument for `aggregate_fn_name` and its state. +#[derive(Debug)] +pub struct PercentileParam { + pub aggregate_fn_name: String, + pub(crate) state: PercentileParamState, + pub(crate) is_desc: bool, +} + +impl PercentileParam { + /// Try to resolve the percentile eagerly. If the expression can't be + /// evaluated without row data (i.e. it references a column), defer + /// resolution to the first batch instead of erroring here. + pub fn try_new( + expr: &Arc, + fn_name: &str, + is_desc: bool, + ) -> Result { + match get_scalar_value(expr) { + Ok(scalar_value) => Ok(PercentileParam { + aggregate_fn_name: fn_name.to_string(), + state: PercentileParamState::Resolved(scalar_to_percentile( + scalar_value, + fn_name, + )?), + is_desc, + }), + Err(_) + if !datafusion_physical_expr::utils::collect_columns(expr).is_empty() => + { + Ok(PercentileParam { + aggregate_fn_name: fn_name.to_string(), + state: PercentileParamState::Pending, + is_desc, + }) + } + _ => plan_err!( + "PercentileParam must be a scalar value or reference a column, got {}", + expr + ), + } + } + + /// Resolve using the current batch if `Pending` + /// and validate that the argument is constant across all batches. + pub(crate) fn resolve(&mut self, array: &ArrayRef) -> Result<()> { + if array.null_count() >= array.len() { + return Ok(()); + } + + let (batch_min, batch_max) = match array.data_type() { + DataType::Float64 => { + let float_array = downcast_value!(array, Float64Array); + (min(float_array), max(float_array)) + } + DataType::Float32 => { + let float_array = downcast_value!(array, Float32Array); + ( + min(float_array).map(|v| v as f64), + max(float_array).map(|v| v as f64), + ) + } + data_type => { + return exec_err!( + "Percentile value for {} must be Float32 or Float64 (got {})", + self.aggregate_fn_name, + data_type + ); + } + }; + let batch_min = batch_min.ok_or_else(|| { + internal_datafusion_err!("expected a non-null percentile value") + })?; + let batch_max = batch_max.ok_or_else(|| { + internal_datafusion_err!("expected a non-null percentile value") + })?; + if batch_min.is_nan() || batch_max.is_nan() { + return exec_err!( + "Percentile value must be between 0.0 and 1.0 inclusive, NaN is invalid" + ); + } + if batch_min != batch_max { + return exec_err!( + "Percentile value for '{}' must be constant across the aggregation, found differing values {} and {}", + self.aggregate_fn_name, + batch_min, + batch_max + ); + } + + match self.state { + PercentileParamState::Resolved(resolved) => { + if batch_min != resolved { + return exec_err!( + "Percentile value for '{}' must be constant across the aggregation, found differing values {} and {}", + self.aggregate_fn_name, + batch_min, + resolved + ); + } + } + PercentileParamState::Pending => { + let resolved = scalar_to_percentile( + ScalarValue::Float64(Some(batch_min)), + &self.aggregate_fn_name, + ) + .map_err(|e| exec_datafusion_err!("{e}"))?; + self.state = PercentileParamState::Resolved(resolved); + } + } + + Ok(()) + } + + /// Returns the resolved percentile. + /// Errors if it has not been yet resolved. + pub(crate) fn get(&self) -> Result { + match self.state { + PercentileParamState::Resolved(value) => Ok(value), + PercentileParamState::Pending => { + let aggregate_fn_name = self.aggregate_fn_name.clone(); + exec_err!( + "Percentile value for '{aggregate_fn_name}' could not be determined: no non-null percentile value was seen" + ) + } + } + } + + /// The percentile to use, applying the `1.0 - p` flip for descending + /// `WITHIN GROUP (ORDER BY ... DESC)`. + pub(crate) fn effective_percentile(&self) -> Result { + Ok(self.apply_order(self.get()?)) + } + + /// Applies the `1.0 - p` flip for descending `WITHIN GROUP (ORDER BY ... DESC)`. + pub(crate) fn apply_order(&self, percentile: f64) -> f64 { + if self.is_desc { + 1.0 - percentile + } else { + percentile + } + } + + /// Returns `true` if the percentile must be read from the input rows, + /// i.e. it was given as a column reference rather than a literal. + pub(crate) fn is_col_ref(&self) -> bool { + matches!(self.state, PercentileParamState::Pending) + } + + /// Casts a percentile argument array to `Float64`. + pub(crate) fn to_float64_array(&self, array: &ArrayRef) -> Result { + match array.data_type() { + DataType::Float64 | DataType::Float32 => { + let array = cast(array, &DataType::Float64)?; + Ok(downcast_value!(array, Float64Array).clone()) + } + data_type => { + let agg_fn_name = &self.aggregate_fn_name; + exec_err!( + "Percentile value for {agg_fn_name} must be Float32 or Float64 (got {data_type})" + ) + } + } + } + + /// Validates `value` as the percentile of a single group: it must be in + /// range and agree with the group's `current` percentile, if any. + pub(crate) fn resolve_group_value( + &self, + current: Option, + value: f64, + ) -> Result { + match current { + Some(current) if current == value => Ok(current), + Some(current) => { + let agg_fn_name = &self.aggregate_fn_name; + exec_err!( + "Percentile value for '{agg_fn_name}' must be constant across the aggregation, found differing values {} and {}", + current, + value + ) + } + None => { + check_percentile_range(value).map_err(|e| exec_datafusion_err!("{e}")) + } + } + } + + /// Returns the resolved percentile of a group. + /// Errors if no percentile was seen for it. + pub(crate) fn get_group_value(&self, value: Option) -> Result { + value.ok_or_else(|| { + exec_datafusion_err!( + "Percentile value for '{}' could not be determined: no non-null percentile value was seen", + &self.aggregate_fn_name + ) + }) + } +} diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index ac0331975bd4..162414346242 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -190,7 +190,7 @@ SELECT c1, approx_percentile_cont(0.95, -1000) WITHIN GROUP (ORDER BY c3) AS c3_ statement error Function 'approx_percentile_cont' failed to match any signature SELECT approx_percentile_cont(0.95, c1) WITHIN GROUP (ORDER BY c3) FROM aggregate_test_100 -statement error DataFusion error: Error during planning: Percentile value for 'APPROX_PERCENTILE_CONT' must be a literal +statement error DataFusion error: Execution error: Percentile value for 'APPROX_PERCENTILE_CONT' must be constant across the aggregation, found differing values SELECT approx_percentile_cont(c12) WITHIN GROUP (ORDER BY c12) FROM aggregate_test_100 statement error DataFusion error: Error during planning: Tdigest max_size value for 'APPROX_PERCENTILE_CONT' must be a literal @@ -2483,7 +2483,7 @@ SELECT APPROX_PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY v) FROM (VALUES (CAST( ---- NULL -query error DataFusion error: Error during planning: Percentile value for 'APPROX_PERCENTILE_CONT' must be Float32 or Float64 literal \(got null\) +query error DataFusion error: Error during planning: Percentile value for 'APPROX_PERCENTILE_CONT' must be Float32 or Float64 \(got null\) SELECT APPROX_PERCENTILE_CONT(NULL) WITHIN GROUP (ORDER BY v) FROM (VALUES (CAST(NULL as INT))) as t (v); # percentile_cont_with_weight_with_nulls @@ -2806,6 +2806,278 @@ c 5 122.5 d 5 124.2 e 5 115.6 + +# percentile from a projected constant column +query R +SELECT approx_percentile_cont(y, m) AS median FROM (SELECT t.x + 1 as y, 0.5 as m FROM (VALUES (10), (20), (30)) AS t(x)) +---- +21 + +query R +SELECT percentile_cont(y, m) AS median FROM (SELECT t.x + 1 as y, 0.5 as m FROM (VALUES (10), (20), (30)) AS t(x)) +---- +21 + +query R +SELECT approx_percentile_cont(m) WITHIN GROUP (ORDER BY y) AS median FROM (SELECT t.x + 1 as y, 0.5 as m FROM (VALUES (10), (20), (30)) AS t(x)) +---- +21 + +query R +SELECT percentile_cont(m) WITHIN GROUP (ORDER BY y) AS median FROM (SELECT t.x + 1 as y, 0.5 as m FROM (VALUES (10), (20), (30)) AS t(x)) +---- +21 + +# percentile as a plain column straight off a VALUES table +query R +SELECT approx_percentile_cont(x, m) FROM (VALUES (10,0.5),(20,0.5),(30,0.5),(40,0.5)) AS t(x,m) +---- +25 + +query R +SELECT percentile_cont(x, m) FROM (VALUES (10,0.5),(20,0.5),(30,0.5),(40,0.5)) AS t(x,m) +---- +25 + +query R +SELECT percentile_cont(x, m) FROM (VALUES (10,0.5),(20,0.5),(30,0.5),(40,0.5)) AS t(x,m) WHERE x > 1000 +---- +NULL + +query R +SELECT approx_percentile_cont(x, m) FROM (VALUES (10,0.5),(20,0.5),(30,0.5),(40,0.5)) AS t(x,m) WHERE x > 1000 +---- +NULL + +query R +SELECT percentile_cont(DISTINCT x, m) FROM (VALUES (10,0.5),(20,0.5),(30,0.5),(40,0.5)) AS t(x,m) WHERE x > 1000 +---- +NULL + + +query R +SELECT approx_percentile_cont_with_weight(y, 1, m) AS median FROM (SELECT t.x + 1 as y, 0.5 as m FROM (VALUES (10), (20), (30)) AS t(x)) +---- +21 + +query IR +SELECT g, approx_percentile_cont(y, m) AS median FROM (SELECT t.x % 3 as g, t.x + 1 as y, 0.5 as m FROM (VALUES (10),(20),(30),(40),(50),(60),(70),(80),(90),(100),(110),(120)) AS t(x)) GROUP BY g ORDER BY g +---- +0 76 +1 56 +2 66 + +query IR +SELECT g, percentile_cont(y, m) AS median FROM (SELECT t.x % 3 as g, t.x + 1 as y, 0.5 as m FROM (VALUES (10),(20),(30),(40),(50),(60),(70),(80),(90),(100),(110),(120)) AS t(x)) GROUP BY g ORDER BY g +---- +0 76 +1 56 +2 66 + +query error DataFusion error: Execution error: Percentile value must be between 0\.0 and 1\.0 inclusive, NaN is invalid +SELECT approx_percentile_cont(x, m) FROM (VALUES (10,'NaN'::FLOAT),(20,'NaN'::FLOAT),(30,'NaN'::FLOAT),(40,'NaN'::FLOAT)) AS t(x,m) + +query error DataFusion error: Execution error: Percentile value must be between 0\.0 and 1\.0 inclusive, NaN is invalid +SELECT percentile_cont(x, m) FROM (VALUES (10,'NaN'::FLOAT),(20,'NaN'::FLOAT),(30,'NaN'::FLOAT),(40,'NaN'::FLOAT)) AS t(x,m) + +query error DataFusion error: Execution error: Error during planning: Percentile value must be between 0\.0 and 1\.0 inclusive, inf is invalid +SELECT approx_percentile_cont(x, m) FROM (VALUES (10,'+Inf'::FLOAT),(20,'+Inf'::FLOAT),(30,'+Inf'::FLOAT),(40,'+Inf'::FLOAT)) AS t(x,m) + +query error DataFusion error: Execution error: Error during planning: Percentile value must be between 0\.0 and 1\.0 inclusive, inf is invalid +SELECT percentile_cont(x, m) FROM (VALUES (10,'+Inf'::FLOAT),(20,'+Inf'::FLOAT),(30,'+Inf'::FLOAT),(40,'+Inf'::FLOAT)) AS t(x,m) + +query error DataFusion error: Execution error: Error during planning: Percentile value must be between 0\.0 and 1\.0 inclusive, \-inf is invalid +SELECT approx_percentile_cont(x, m) FROM (VALUES (10,'-Inf'::FLOAT),(20,'-Inf'::FLOAT),(30,'-Inf'::FLOAT),(40,'-Inf'::FLOAT)) AS t(x,m) + +query error DataFusion error: Execution error: Error during planning: Percentile value must be between 0\.0 and 1\.0 inclusive, \-inf is invalid +SELECT percentile_cont(x, m) FROM (VALUES (10,'-Inf'::FLOAT),(20,'-Inf'::FLOAT),(30,'-Inf'::FLOAT),(40,'-Inf'::FLOAT)) AS t(x,m) + +# the percentile only has to be constant within a group, not across groups +query IR +SELECT g, percentile_cont(x, p) FROM (VALUES (1, 10, 0.0), (1, 20, 0.0), (2, 30, 1.0), (2, 40, 1.0)) AS t(g, x, p) GROUP BY g ORDER BY g +---- +1 10 +2 40 + +query IR +SELECT g, approx_percentile_cont(x, p) FROM (VALUES (1, 10, 0.0), (1, 20, 0.0), (2, 30, 1.0), (2, 40, 1.0)) AS t(g, x, p) GROUP BY g ORDER BY g +---- +1 10 +2 40 + +query IR +SELECT g, percentile_cont(p) WITHIN GROUP (ORDER BY x DESC) FROM (VALUES (1, 10, 0.0), (1, 20, 0.0), (2, 30, 1.0), (2, 40, 1.0)) AS t(g, x, p) GROUP BY g ORDER BY g +---- +1 20 +2 30 + +# a group with no non-null values evaluates to NULL even without a percentile +query IR +SELECT g, percentile_cont(x, p) FROM (VALUES (1, 10, 0.0), (1, 20, 0.0), (2, NULL, NULL)) AS t(g, x, p) GROUP BY g ORDER BY g +---- +1 10 +2 NULL + +# the percentile must still be constant within each group +query error DataFusion error: Execution error: Percentile value for 'PERCENTILE_CONT' must be constant across the aggregation, found differing values 0 and 1 +SELECT g, percentile_cont(x, p) FROM (VALUES (1, 10, 0.0), (1, 20, 1.0), (2, 30, 1.0), (2, 40, 1.0)) AS t(g, x, p) GROUP BY g + +query error DataFusion error: Execution error: Percentile value for 'APPROX_PERCENTILE_CONT' must be constant across the aggregation, found differing values +SELECT g, approx_percentile_cont(x, p) FROM (VALUES (1, 10, 0.0), (1, 20, 1.0), (2, 30, 1.0), (2, 40, 1.0)) AS t(g, x, p) GROUP BY g + +# a genuinely non-constant percentile must error clearly +query error DataFusion error: Execution error: Percentile value for 'APPROX_PERCENTILE_CONT' must be constant across the aggregation, found differing values +SELECT approx_percentile_cont(x, m) FROM (VALUES (10,0.1),(20,0.5),(30,0.9),(40,0.5)) AS t(x,m) + +query error DataFusion error: Execution error: Percentile value for 'PERCENTILE_CONT' must be constant across the aggregation, found differing values 0.1 and 0.9 +SELECT percentile_cont(x, m) FROM (VALUES (10,0.1),(20,0.5),(30,0.9),(40,0.5)) AS t(x,m) + +# percentile column with DESC ordering is flipped to 1 - p +query RR +SELECT percentile_cont(m) WITHIN GROUP (ORDER BY x DESC), approx_percentile_cont(m) WITHIN GROUP (ORDER BY x DESC) FROM (VALUES (10, 0.25), (20, 0.25), (30, 0.25), (40, 0.25), (50, 0.25)) AS t(x, m) +---- +40 42.5 + +query RR +SELECT percentile_cont(m) WITHIN GROUP (ORDER BY x), approx_percentile_cont(m) WITHIN GROUP (ORDER BY x) FROM (VALUES (10, 0.25), (20, 0.25), (30, 0.25), (40, 0.25), (50, 0.25)) AS t(x, m) +---- +20 17.5 + +# Float32 percentile column +query RR +SELECT percentile_cont(x, arrow_cast(m, 'Float32')), approx_percentile_cont(x, arrow_cast(m, 'Float32')) FROM (VALUES (10, 0.5), (20, 0.5), (30, 0.5)) AS t(x, m) +---- +20 20 + +# null percentile rows are ignored when resolving the percentile +query RR +SELECT percentile_cont(x, m), approx_percentile_cont(x, m) FROM (VALUES (10, 0.5), (20, NULL), (30, 0.5), (40, NULL)) AS t(x, m) +---- +25 25 + +# an all-null percentile column cannot be resolved +query error DataFusion error: Execution error: Percentile value for 'PERCENTILE_CONT' could not be determined: no non\-null percentile value was seen +SELECT percentile_cont(x, m) FROM (VALUES (10, CAST(NULL AS DOUBLE)), (20, NULL)) AS t(x, m) + +query error DataFusion error: Execution error: Percentile value for 'APPROX_PERCENTILE_CONT' could not be determined: no non\-null percentile value was seen +SELECT approx_percentile_cont(x, m) FROM (VALUES (10, CAST(NULL AS DOUBLE)), (20, NULL)) AS t(x, m) + +# out of range percentile column +query error DataFusion error: Execution error: Error during planning: Percentile value must be between 0.0 and 1.0 inclusive, 1.5 is invalid +SELECT percentile_cont(x, m) FROM (VALUES (10, 1.5), (20, 1.5)) AS t(x, m) + +query error DataFusion error: Execution error: Error during planning: Percentile value must be between 0.0 and 1.0 inclusive, -0.5 is invalid +SELECT approx_percentile_cont(x, m) FROM (VALUES (10, -0.5), (20, -0.5)) AS t(x, m) + +# Percentile column spread across multiple batches and partitions: the +# percentile is resolved per batch and must survive partial -> final merges +statement ok +set datafusion.execution.batch_size = 10; + +query R +SELECT percentile_cont(x, m) FROM (SELECT value AS x, 0.5 AS m FROM generate_series(1, 100)) +---- +50.5 + +query R +SELECT approx_percentile_cont(x, m) FROM (SELECT value AS x, 0.5 AS m FROM generate_series(1, 100)) +---- +50.5 + +query R +SELECT percentile_cont(DISTINCT x, m) FROM (SELECT value % 10 AS x, 0.5 AS m FROM generate_series(1, 100)) +---- +4.5 + +query R +SELECT approx_percentile_cont_with_weight(x, 1, m) FROM (SELECT value AS x, 0.5 AS m FROM generate_series(1, 100)) +---- +50.5 + +query IR +SELECT g, percentile_cont(x, m) FROM (SELECT value % 4 AS g, value AS x, 0.25 AS m FROM generate_series(1, 100)) GROUP BY g ORDER BY g +---- +0 28 +1 25 +2 26 +3 27 + +query IR +SELECT g, approx_percentile_cont(x, m) FROM (SELECT value % 4 AS g, value AS x, 0.25 AS m FROM generate_series(1, 100)) GROUP BY g ORDER BY g +---- +0 27 +1 24 +2 25 +3 26 + +# per-group percentiles spread across batches and partitions +query IR +SELECT g, percentile_cont(x, p) FROM (SELECT value % 2 AS g, value AS x, CASE WHEN value % 2 = 0 THEN 0.0 ELSE 1.0 END AS p FROM generate_series(1, 100)) GROUP BY g ORDER BY g +---- +0 2 +1 99 + +query IR +SELECT g, approx_percentile_cont(x, p) FROM (SELECT value % 2 AS g, value AS x, CASE WHEN value % 2 = 0 THEN 0.0 ELSE 1.0 END AS p FROM generate_series(1, 100)) GROUP BY g ORDER BY g +---- +0 2 +1 99 + +# Each batch is constant on its own, but the percentile differs across +# batches, so the mismatch must be detected across batches / on merge +query error DataFusion error: Execution error: Percentile value for 'PERCENTILE_CONT' must be constant across the aggregation, found differing values 0.9 and 0.5 +SELECT percentile_cont(x, m) FROM (SELECT value AS x, CASE WHEN value <= 50 THEN 0.5 ELSE 0.9 END AS m FROM generate_series(1, 100)) + +query error DataFusion error: Execution error: Percentile value for 'APPROX_PERCENTILE_CONT' must be constant across the aggregation, found differing values +SELECT approx_percentile_cont(x, m) FROM (SELECT value AS x, CASE WHEN value <= 50 THEN 0.5 ELSE 0.9 END AS m FROM generate_series(1, 100)) + +query error DataFusion error: Execution error: Percentile value for 'APPROX_PERCENTILE_CONT' must be constant across the aggregation, found differing values +SELECT approx_percentile_cont_with_weight(x, 1, m) FROM (SELECT value AS x, CASE WHEN value <= 50 THEN 0.5 ELSE 0.9 END AS m FROM generate_series(1, 100)) + +query error DataFusion error: Execution error: Percentile value for 'PERCENTILE_CONT' must be constant across the aggregation, found differing values 0.5 and 0.9 +SELECT g, percentile_cont(x, m) FROM (SELECT value % 4 AS g, value AS x, CASE WHEN value <= 50 THEN 0.5 ELSE 0.9 END AS m FROM generate_series(1, 100)) GROUP BY g + +# Force skipping partial aggregation so the groups accumulator goes through +# `convert_to_state` with a percentile column +statement ok +set datafusion.execution.skip_partial_aggregation_probe_rows_threshold = 10; + +statement ok +set datafusion.execution.skip_partial_aggregation_probe_ratio_threshold = 0.0; + +query IR +SELECT g, percentile_cont(x, m) FROM (SELECT value % 4 AS g, value AS x, 0.25 AS m FROM generate_series(1, 100)) GROUP BY g ORDER BY g +---- +0 28 +1 25 +2 26 +3 27 + +query error DataFusion error: Execution error: Percentile value for 'PERCENTILE_CONT' must be constant across the aggregation, found differing values 0.5 and 0.9 +SELECT g, percentile_cont(x, m) FROM (SELECT value % 4 AS g, value AS x, CASE WHEN value <= 50 THEN 0.5 ELSE 0.9 END AS m FROM generate_series(1, 100)) GROUP BY g + +query IR +SELECT g, percentile_cont(x, p) FROM (SELECT value % 2 AS g, value AS x, CASE WHEN value % 2 = 0 THEN 0.0 ELSE 1.0 END AS p FROM generate_series(1, 100)) GROUP BY g ORDER BY g +---- +0 2 +1 99 + +# filtered-out rows must not contribute their percentile +query IR +SELECT g, percentile_cont(x, p) FILTER (WHERE x > 10) FROM (SELECT value % 2 AS g, value AS x, CASE WHEN value <= 10 THEN 0.5 WHEN value % 2 = 0 THEN 0.0 ELSE 1.0 END AS p FROM generate_series(1, 100)) GROUP BY g ORDER BY g +---- +0 12 +1 99 + +statement ok +reset datafusion.execution.skip_partial_aggregation_probe_rows_threshold; + +statement ok +reset datafusion.execution.skip_partial_aggregation_probe_ratio_threshold; + +statement ok +reset datafusion.execution.batch_size; + # error is unique to this UDAF query TRR SELECT c1, avg(c2) AS c2, avg(c3) AS c3 FROM aggregate_test_100 GROUP BY 1 ORDER BY 1 @@ -4301,7 +4573,7 @@ SELECT percentile_cont(1.5) WITHIN GROUP (ORDER BY c3) FROM aggregate_test_100 statement error DataFusion error: Error during planning: Percentile value must be between 0.0 and 1.0 inclusive SELECT percentile_cont(-0.1) WITHIN GROUP (ORDER BY c3) FROM aggregate_test_100 -statement error DataFusion error: Error during planning: Percentile value for 'PERCENTILE_CONT' must be a literal +statement error DataFusion error: Execution error: Percentile value for 'PERCENTILE_CONT' must be constant across the aggregation, found differing values 1 and 5 SELECT percentile_cont(c2) WITHIN GROUP (ORDER BY c3) FROM aggregate_test_100 statement error DataFusion error: Error during planning: \[IGNORE | RESPECT\] NULLS are not permitted for percentile_cont