This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25511-21a3215b66a8620cb932899de419e6b43d1848e6 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit dcf9dd487a5cff6a53cbf33de5787956dd20c94c Author: Stefan Wang <[email protected]> AuthorDate: Thu Sep 24 13:35:01 2026 +0000 fix: respect volatility for all function expressions (#25511) ## Which issue does this PR close? Closes https://github.com/apache/datafusion/issues/25504. ## Rationale for this change Repeated calls to a volatile higher-order UDF can reuse the first call's result. Comparing two calls can also return true without invoking the function. Both queries now execute each call. ## What changes are included in this PR? The volatility check reads the declared signature for aggregate, window and higher-order expressions, as it already does for scalar expressions. This includes aggregates used as window functions. The join optimizer uses the shared repeatability check instead of checking aggregate volatility separately. ## What is the testing strategy for this PR? The regression registers a counter UDF and runs both queries through the SQL engine. It also checks every function kind, volatility level and aliased expression. Join-optimizer controls preserve DISTINCT handling and reject volatile expressions and subqueries. On macOS arm64 with Rust 1.98.1 and the pinned test data: ```sh cargo test --locked --profile ci -j6 -p datafusion --test user_defined_integration volatility -- --nocapture cargo test --locked --profile ci -j6 -p datafusion-expr -p datafusion-optimizer --lib cargo test --locked --profile ci -j6 -p datafusion-sqllogictest --test sqllogictests -- --test-threads 6 cse.slt expr.slt array_transform.slt eliminate_outer_join.slt aggregates_simplify.slt ``` <details><summary>Raw logs</summary> Before the fix: ```text +-------+--------+ | first | second | +-------+--------+ | 0 | 0 | +-------+--------+ calls=1 +-------+ | equal | +-------+ | true | +-------+ calls=0 ``` After: ```text +-------+--------+ | first | second | +-------+--------+ | 0 | 1 | +-------+--------+ calls=2 +-------+ | equal | +-------+ | false | +-------+ calls=2 Progress: 6/6 files completed (100%) slt_exit_code=0 ``` </details> ## Are there any user-facing changes? The affected queries return the results of separate function calls. There is no public API change. --------- Signed-off-by: 1fanwang <[email protected]> --- datafusion/core/tests/user_defined/mod.rs | 2 + datafusion/core/tests/user_defined/volatility.rs | 180 +++++++++++++++++++++++ datafusion/expr/src/expr.rs | 14 +- datafusion/optimizer/src/utils.rs | 10 +- 4 files changed, 198 insertions(+), 8 deletions(-) diff --git a/datafusion/core/tests/user_defined/mod.rs b/datafusion/core/tests/user_defined/mod.rs index 4dad3ec457..c2bc0ffa71 100644 --- a/datafusion/core/tests/user_defined/mod.rs +++ b/datafusion/core/tests/user_defined/mod.rs @@ -45,3 +45,5 @@ mod insert_operation; /// Tests for `StatisticsRequest`s flowing from a custom optimizer rule /// through the physical planner into a custom `TableProvider`. mod statistics_requests; + +mod volatility; diff --git a/datafusion/core/tests/user_defined/volatility.rs b/datafusion/core/tests/user_defined/volatility.rs new file mode 100644 index 0000000000..1e013ac1a5 --- /dev/null +++ b/datafusion/core/tests/user_defined/volatility.rs @@ -0,0 +1,180 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::{ + Arc, + atomic::{AtomicI64, Ordering}, +}; + +use datafusion::arrow::array::{ArrayRef, AsArray}; +use datafusion::arrow::datatypes::{DataType, Field, FieldRef, Int64Type}; +use datafusion::common::test_util::batches_to_string; +use datafusion::common::{Result, ScalarValue, assert_batches_eq}; +use datafusion::logical_expr::expr::{HigherOrderFunction, WindowFunction}; +use datafusion::logical_expr::{ + ColumnarValue, Expr, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, + HigherOrderSignature, HigherOrderUDF, HigherOrderUDFImpl, LambdaParametersProgress, + PartitionEvaluator, ValueOrLambda, Volatility, col, create_udaf, create_udf, + create_udwf, +}; +use datafusion::prelude::SessionContext; +use datafusion_functions_aggregate::average::AvgAccumulator; + +static NEXT_VALUE: AtomicI64 = AtomicI64::new(0); + +#[derive(Debug, PartialEq, Eq, Hash)] +struct NextValue { + signature: HigherOrderSignature, +} + +impl HigherOrderUDFImpl for NextValue { + fn name(&self) -> &str { + "next_value" + } + + fn signature(&self) -> &HigherOrderSignature { + &self.signature + } + + fn lambda_parameters( + &self, + _step: usize, + _fields: &[ValueOrLambda<FieldRef, Option<FieldRef>>], + ) -> Result<LambdaParametersProgress> { + Ok(LambdaParametersProgress::Complete(vec![])) + } + + fn return_field_from_args( + &self, + _args: HigherOrderReturnFieldArgs, + ) -> Result<FieldRef> { + Ok(Arc::new(Field::new("value", DataType::Int64, false))) + } + + fn invoke_with_args(&self, _args: HigherOrderFunctionArgs) -> Result<ColumnarValue> { + let value = if self.signature.volatility == Volatility::Volatile { + NEXT_VALUE.fetch_add(1, Ordering::Relaxed) + } else { + 0 + }; + Ok(ColumnarValue::Scalar(ScalarValue::Int64(Some(value)))) + } +} + +#[test] +fn function_volatility() { + #[derive(Debug)] + struct IdentityEvaluator; + + impl PartitionEvaluator for IdentityEvaluator { + fn evaluate_all( + &mut self, + values: &[ArrayRef], + _num_rows: usize, + ) -> Result<ArrayRef> { + Ok(Arc::clone(&values[0])) + } + } + + for volatility in [ + Volatility::Immutable, + Volatility::Stable, + Volatility::Volatile, + ] { + let scalar = create_udf( + "identity", + vec![DataType::Float64], + DataType::Float64, + volatility, + Arc::new(|args| Ok(args[0].clone())), + ); + let aggregate = Arc::new(create_udaf( + "average", + vec![DataType::Float64], + Arc::new(DataType::Float64), + volatility, + Arc::new(|_| Ok(Box::<AvgAccumulator>::default())), + Arc::new(vec![DataType::UInt64, DataType::Float64]), + )); + let window = create_udwf( + "identity_window", + DataType::Float64, + Arc::new(DataType::Float64), + volatility, + Arc::new(|| Ok(Box::new(IdentityEvaluator))), + ); + let higher_order = Arc::new(HigherOrderUDF::new_from_impl(NextValue { + signature: HigherOrderSignature::any(0, volatility), + })); + + for expr in [ + scalar.call(vec![col("value")]), + aggregate.call(vec![col("value")]), + WindowFunction::new(Arc::clone(&aggregate), vec![col("value")]).into(), + window.call(vec![col("value")]), + Expr::HigherOrderFunction(HigherOrderFunction::new(higher_order, vec![])), + ] { + let expected = volatility == Volatility::Volatile; + assert_eq!(expr.is_volatile_node(), expected, "{expr}"); + assert_eq!(expr.is_volatile(), expected, "{expr}"); + let aliased = expr.alias("result"); + assert!(!aliased.is_volatile_node()); + assert_eq!(aliased.is_volatile(), expected, "{aliased}"); + } + } +} + +#[tokio::test] +async fn volatile_higher_order_function_is_not_eliminated() -> Result<()> { + let ctx = SessionContext::new(); + ctx.register_higher_order_function(Arc::new(HigherOrderUDF::new_from_impl( + NextValue { + signature: HigherOrderSignature::any(0, Volatility::Volatile), + }, + ))); + let mut results = Vec::new(); + let mut calls = Vec::new(); + for sql in [ + "SELECT next_value() AS first, next_value() AS second", + "SELECT next_value() = next_value() AS equal", + ] { + NEXT_VALUE.store(0, Ordering::Relaxed); + let batches = ctx.sql(sql).await?.collect().await?; + let count = NEXT_VALUE.load(Ordering::Relaxed); + println!("{sql}\n{}\ncalls={count}", batches_to_string(&batches)); + results.push(batches); + calls.push(count); + } + + let batch = &results[0][0]; + assert_ne!( + batch.column(0).as_primitive::<Int64Type>().value(0), + batch.column(1).as_primitive::<Int64Type>().value(0) + ); + assert_batches_eq!( + [ + "+-------+", + "| equal |", + "+-------+", + "| false |", + "+-------+" + ], + &results[1] + ); + assert_eq!(calls, vec![2, 2]); + Ok(()) +} diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index e8e805c2b9..c197a6af6f 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -2157,7 +2157,19 @@ impl Expr { /// - `rand()` returns `true`, /// - `a + rand()` returns `false` 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(func) => func.func.signature().volatility, + Expr::AggregateFunction(func) => func.func.signature().volatility, + Expr::WindowFunction(func) => match &func.fun { + WindowFunctionDefinition::AggregateUDF(func) => { + func.signature().volatility + } + WindowFunctionDefinition::WindowUDF(func) => func.signature().volatility, + }, + Expr::HigherOrderFunction(func) => func.func.signature().volatility, + _ => return false, + }; + volatility == Volatility::Volatile } /// Returns true if the expression is volatile, i.e. whether it can return different diff --git a/datafusion/optimizer/src/utils.rs b/datafusion/optimizer/src/utils.rs index 716efc8197..ce4242bb9d 100644 --- a/datafusion/optimizer/src/utils.rs +++ b/datafusion/optimizer/src/utils.rs @@ -31,7 +31,7 @@ use datafusion_expr::expr::{Exists, InSubquery, SetComparison}; use datafusion_expr::expr_rewriter::replace_col; use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{ - ColumnarValue, DistinctHandling, Expr, Volatility, WriteOp, logical_plan::LogicalPlan, + ColumnarValue, DistinctHandling, Expr, WriteOp, logical_plan::LogicalPlan, }; use datafusion_physical_expr::create_physical_expr; use log::{debug, trace}; @@ -41,7 +41,7 @@ use std::sync::Arc; /// as it was initially placed here and then moved elsewhere. pub use datafusion_expr::expr_rewriter::NamePreserver; -/// Whether an expression is free of volatile scalar functions and subqueries. +/// Whether an expression is free of volatile functions and subqueries. /// Subqueries are conservative barriers because their plans may contain /// volatile expressions that [`Expr::is_volatile`] does not visit. pub(crate) fn is_repeatable(expr: &Expr) -> bool { @@ -79,11 +79,7 @@ pub(crate) fn is_duplicate_insensitive_aggregate(mut expr: &Expr) -> bool { // Variants added in the future are treated the same way. _ => false, }; - // Expr::is_volatile checks scalar functions only; check the aggregate - // function's own volatility separately. - ignores_duplicates - && aggregate.func.signature().volatility != Volatility::Volatile - && is_repeatable(expr) + ignores_duplicates && is_repeatable(expr) } /// Return the expression schema for a MERGE DML node. --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
