diff --git a/datafusion/physical-expr/src/aggregate.rs b/datafusion/physical-expr/src/aggregate.rs index a660a272b959..6d95d8ea12bd 100644 --- a/datafusion/physical-expr/src/aggregate.rs +++ b/datafusion/physical-expr/src/aggregate.rs @@ -264,6 +264,15 @@ impl AggregateExprBuilder { } = self; assert_or_internal_err!(!args.is_empty(), "args should not be empty"); + // An order-insensitive aggregate ignores its ORDER BY, so drop it here. + // Everything derived from `order_bys` below, such as the ordering fields + // in the aggregate's state, then agrees that there is no ordering. + let order_bys = if fun.order_sensitivity().is_insensitive() { + vec![] + } else { + order_bys + }; + let ordering_types = order_bys .iter() .map(|e| e.expr.data_type(&schema)) @@ -772,11 +781,7 @@ impl AggregateFunctionExpr { /// Returns the ORDER BY expressions for the aggregate function. pub fn order_bys(&self) -> &[PhysicalSortExpr] { - if self.order_sensitivity().is_insensitive() { - &[] - } else { - &self.order_bys - } + &self.order_bys } /// Indicates whether aggregator can produce the correct result with any @@ -1052,10 +1057,7 @@ impl AggregateFunctionExpr { args: Vec>, order_by_exprs: Vec>, ) -> Option { - if args.len() != self.args.len() - || (self.order_sensitivity() != AggregateOrderSensitivity::Insensitive - && order_by_exprs.len() != self.order_bys.len()) - { + if args.len() != self.args.len() || order_by_exprs.len() != self.order_bys.len() { return None; } @@ -1192,7 +1194,9 @@ mod tests { use arrow::datatypes::Field; use datafusion_common::metadata::FieldMetadata; - use datafusion_expr::{col, test::function_stub::sum}; + use datafusion_expr::{ + AggregateUDFImpl, Signature, Volatility, col, test::function_stub::sum, + }; fn aggregate_test_schema() -> Result<(Schema, DFSchema)> { let schema = Schema::new(vec![Field::new("column1", DataType::Int64, true)]); @@ -1263,4 +1267,98 @@ mod tests { Ok(()) } + + /// An aggregate that uses the default `AggregateUDFImpl::state_fields`, + /// which appends the ordering fields to the state. + #[derive(Debug, PartialEq, Eq, Hash)] + struct DefaultStateUdaf { + signature: Signature, + order_insensitive: bool, + } + + impl DefaultStateUdaf { + fn new(order_insensitive: bool) -> Self { + Self { + signature: Signature::any(1, Volatility::Immutable), + order_insensitive, + } + } + } + + impl AggregateUDFImpl for DefaultStateUdaf { + fn name(&self) -> &str { + "default_state_udaf" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + Ok(arg_types[0].clone()) + } + + /// Always fails, reporting how many ORDER BY expressions it was given. + fn accumulator(&self, acc_args: AccumulatorArgs) -> Result> { + not_impl_err!("accumulator with {} order_bys", acc_args.order_bys.len()) + } + + fn order_sensitivity(&self) -> AggregateOrderSensitivity { + if self.order_insensitive { + AggregateOrderSensitivity::Insensitive + } else { + AggregateOrderSensitivity::HardRequirement + } + } + } + + /// Builds `default_state_udaf(v ORDER BY k)`. + fn build_with_order_by(order_insensitive: bool) -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("v", DataType::Int64, true), + Field::new("k", DataType::Int64, true), + ])); + let fun = AggregateUDF::from(DefaultStateUdaf::new(order_insensitive)); + AggregateExprBuilder::new(Arc::new(fun), vec![Arc::new(Column::new("v", 0))]) + .order_by(vec![PhysicalSortExpr { + expr: Arc::new(Column::new("k", 1)), + options: SortOptions::default(), + }]) + .schema(schema) + .alias("default_state_udaf(v) ORDER BY [k ASC NULLS LAST]") + .build() + } + + #[test] + fn order_insensitive_aggregate_discards_order_by() -> Result<()> { + let expr = build_with_order_by(true)?; + assert!(expr.order_bys().is_empty()); + assert!(expr.all_expressions().order_by_exprs.is_empty()); + // Only the value: the default `state_fields` has no ordering fields to append + assert_eq!(expr.state_fields()?.len(), 1); + let err = expr.create_accumulator().unwrap_err(); + assert!(err.message().contains("accumulator with 0 order_bys")); + + // Rewriting the expressions does not bring the ORDER BY back + let rewritten = expr + .with_new_expressions(expr.expressions(), vec![]) + .expect("rewrite is supported"); + assert!(rewritten.order_bys().is_empty()); + assert_eq!(rewritten.state_fields()?.len(), 1); + + Ok(()) + } + + #[test] + fn order_sensitive_aggregate_keeps_order_by() -> Result<()> { + let expr = build_with_order_by(false)?; + assert_eq!(expr.order_bys().len(), 1); + assert_eq!(expr.all_expressions().order_by_exprs.len(), 1); + // The value, followed by the ordering field + assert_eq!(expr.state_fields()?.len(), 2); + let err = expr.create_accumulator().unwrap_err(); + assert!(err.message().contains("accumulator with 1 order_bys")); + + Ok(()) + } } diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index 942a6f3cc898..97539ea479d4 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -3075,6 +3075,33 @@ FRA 50 200 GRC 30 80 TUR 75 100 +# An ORDER BY on an order-insensitive aggregator is ignored, so its partial +# state must not include ordering fields, with or without GROUP BY. MIN and MAX +# use the default state fields; SUM is a control that already defines its own. +query TRRR +SELECT country, MIN(amount ORDER BY ts DESC) AS min1, + MAX(amount ORDER BY ts DESC) AS max1, + SUM(amount ORDER BY ts DESC) AS sum1 + FROM sales_global + GROUP BY country + ORDER BY country +---- +FRA 50 200 250 +GRC 30 80 110 +TUR 75 100 175 + +# Without GROUP BY too. The filter makes the input multi-partition, so the +# aggregation runs in Partial and Final modes; over the single partition of +# `sales_global` alone it would run in Single mode and never build a state. +query RRR +SELECT MIN(amount ORDER BY ts DESC) AS min1, + MAX(amount ORDER BY ts DESC) AS max1, + SUM(amount ORDER BY ts DESC) AS sum1 + FROM sales_global + WHERE amount > 0 +---- +30 200 535 + # Conversion in between FIRST_VALUE and LAST_VALUE to resolve # contradictory requirements should work in multi partitions. query TT