From c225aae8ece9654a70f23ab93d7ac8b9ab5e81aa Mon Sep 17 00:00:00 2001 From: Anish Mehta Date: Sun, 20 Sep 2026 01:52:55 +0530 Subject: [PATCH] fix: honor the volatility of aggregate, window and higher-order functions in `Expr::is_volatile` `Expr::is_volatile_node` only looked at `Expr::ScalarFunction`, so a volatile `AggregateUDF`, `WindowUDF` or `HigherOrderUDF` looked deterministic to `Expr::is_volatile` and to every optimizer rule built on it (common subexpression elimination, filter and TopK pushdown, the simplifier). The physical `HigherOrderFunctionExpr::is_volatile_node` already checked the signature; the logical `Expr` was the outlier. Closes #25504 --- datafusion/expr/src/expr.rs | 183 +++++++++++++++++++++++++++++++++++- 1 file changed, 182 insertions(+), 1 deletion(-) diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index e8e805c2b9b75..9930721f498fb 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -2156,8 +2156,24 @@ impl Expr { /// Note: unlike [`Self::is_volatile`], this function does not consider inputs: /// - `rand()` returns `true`, /// - `a + rand()` returns `false` + /// + /// Scalar, aggregate, window and higher-order function calls are volatile + /// when their function declares [`Volatility::Volatile`]. pub fn is_volatile_node(&self) -> bool { - matches!(self, Expr::ScalarFunction(func) if func.func.signature().volatility == Volatility::Volatile) + let volatility = match self { + Expr::ScalarFunction(ScalarFunction { func, .. }) => { + func.signature().volatility + } + Expr::AggregateFunction(AggregateFunction { func, .. }) => { + func.signature().volatility + } + Expr::WindowFunction(window) => window.fun.signature().volatility, + Expr::HigherOrderFunction(HigherOrderFunction { func, .. }) => { + func.signature().volatility + } + _ => return false, + }; + volatility == Volatility::Volatile } /// Returns true if the expression is volatile, i.e. whether it can return different @@ -4326,6 +4342,171 @@ mod test { assert_eq!(udf.signature().volatility, Volatility::Volatile); } + #[test] + fn test_is_volatile_aggregate_window_and_higher_order_funcs() { + use crate::function::{ + AccumulatorArgs, PartitionEvaluatorArgs, WindowUDFFieldArgs, + }; + use crate::higher_order_function::{ + HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, HigherOrderSignature, + HigherOrderTypeSignature, HigherOrderUDFImpl, + }; + use crate::{ + Accumulator, AggregateUDFImpl, PartitionEvaluator, WindowUDF, WindowUDFImpl, + }; + + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestAggregateUDF { + signature: Signature, + } + impl AggregateUDFImpl for TestAggregateUDF { + fn name(&self) -> &str { + "TestAggregateUDF" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Float32) + } + fn accumulator( + &self, + _acc_args: AccumulatorArgs, + ) -> Result> { + unimplemented!() + } + } + + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestWindowUDF { + signature: Signature, + } + impl WindowUDFImpl for TestWindowUDF { + fn name(&self) -> &str { + "TestWindowUDF" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn partition_evaluator( + &self, + _partition_evaluator_args: PartitionEvaluatorArgs, + ) -> Result> { + unimplemented!() + } + fn field(&self, _field_args: WindowUDFFieldArgs) -> Result { + Ok(Arc::new(Field::new( + "TestWindowUDF", + DataType::Float32, + true, + ))) + } + } + + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestHigherOrderUDF { + signature: HigherOrderSignature, + } + impl HigherOrderUDFImpl for TestHigherOrderUDF { + fn name(&self) -> &str { + "TestHigherOrderUDF" + } + fn signature(&self) -> &HigherOrderSignature { + &self.signature + } + fn lambda_parameters( + &self, + _step: usize, + _fields: &[ValueOrLambda>], + ) -> Result { + unimplemented!() + } + fn return_field_from_args( + &self, + _args: HigherOrderReturnFieldArgs, + ) -> Result { + unimplemented!() + } + fn invoke_with_args( + &self, + _args: HigherOrderFunctionArgs, + ) -> Result { + unimplemented!() + } + } + + for (volatility, expected) in [ + (Volatility::Immutable, false), + (Volatility::Stable, false), + (Volatility::Volatile, true), + ] { + let signature = Signature::uniform(1, vec![DataType::Float32], volatility); + + let udaf = Arc::new(AggregateUDF::from(TestAggregateUDF { + signature: signature.clone(), + })); + let aggregate = Expr::AggregateFunction(AggregateFunction::new_udf( + Arc::clone(&udaf), + vec![col("a")], + false, + None, + vec![], + None, + )); + assert_eq!( + aggregate.is_volatile_node(), + expected, + "{volatility:?} aggregate" + ); + assert_eq!( + aggregate.is_volatile(), + expected, + "{volatility:?} aggregate" + ); + + // An aggregate function used as a window function keeps its volatility + let window_aggregate = Expr::from(WindowFunction::new(udaf, vec![col("a")])); + assert_eq!( + window_aggregate.is_volatile_node(), + expected, + "{volatility:?} window aggregate" + ); + + let udwf = Arc::new(WindowUDF::from(TestWindowUDF { + signature: signature.clone(), + })); + let window = Expr::from(WindowFunction::new(udwf, vec![col("a")])); + assert_eq!(window.is_volatile_node(), expected, "{volatility:?} window"); + assert_eq!(window.is_volatile(), expected, "{volatility:?} window"); + + let udhof = Arc::new(HigherOrderUDF::from(TestHigherOrderUDF { + signature: HigherOrderSignature::new( + HigherOrderTypeSignature::Any(1), + volatility, + ), + })); + let higher_order = Expr::HigherOrderFunction(HigherOrderFunction::new( + udhof, + vec![col("a")], + )); + assert_eq!( + higher_order.is_volatile_node(), + expected, + "{volatility:?} higher-order" + ); + assert_eq!( + higher_order.is_volatile(), + expected, + "{volatility:?} higher-order" + ); + + // The volatility is found through the argument of an ordinary expression too + let nested = lit(1.0f32) + aggregate; + assert!(!nested.is_volatile_node()); + assert_eq!(nested.is_volatile(), expected, "{volatility:?} nested"); + } + } + use super::*; use crate::logical_plan::{EmptyRelation, LogicalPlan};