diff --git a/datafusion/core/tests/user_defined/mod.rs b/datafusion/core/tests/user_defined/mod.rs index 4dad3ec4577d9..c2bc0ffa713af 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 0000000000000..1e013ac1a5a6d --- /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>], + ) -> Result { + Ok(LambdaParametersProgress::Complete(vec![])) + } + + fn return_field_from_args( + &self, + _args: HigherOrderReturnFieldArgs, + ) -> Result { + Ok(Arc::new(Field::new("value", DataType::Int64, false))) + } + + fn invoke_with_args(&self, _args: HigherOrderFunctionArgs) -> Result { + 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 { + 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::::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::().value(0), + batch.column(1).as_primitive::().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 e8e805c2b9b75..c197a6af6fc40 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