From e43466a79cee7272e3279dbbda24c39c21631e3f Mon Sep 17 00:00:00 2001 From: Piotr Findeisen Date: Tue, 5 Nov 2024 17:03:21 +0100 Subject: [PATCH] Migrate invoke_no_args to invoke_batch (#13246) Migrate scalar UDFs from deprecated invoke_no_args to invoke_batch. --- .../user_defined/user_defined_scalar_functions.rs | 10 +++++----- datafusion/expr/src/udf.rs | 14 ++++++-------- datafusion/functions-nested/src/make_array.rs | 10 +++++----- datafusion/functions/src/core/version.rs | 15 +++++++++------ datafusion/functions/src/math/pi.rs | 11 +++++++++-- datafusion/functions/src/math/random.rs | 15 +++++++++------ datafusion/functions/src/string/uuid.rs | 15 +++++++++------ 7 files changed, 52 insertions(+), 38 deletions(-) diff --git a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs index 8453a360cd25..cf403e5d640f 100644 --- a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs @@ -209,11 +209,11 @@ impl ScalarUDFImpl for Simple0ArgsScalarUDF { Ok(self.return_type.clone()) } - fn invoke(&self, _args: &[ColumnarValue]) -> Result { - not_impl_err!("{} function does not accept arguments", self.name()) - } - - fn invoke_no_args(&self, _number_rows: usize) -> Result { + fn invoke_batch( + &self, + _args: &[ColumnarValue], + _number_rows: usize, + ) -> Result { Ok(ColumnarValue::Scalar(ScalarValue::Int32(Some(100)))) } } diff --git a/datafusion/expr/src/udf.rs b/datafusion/expr/src/udf.rs index 003a3ed36a60..9ee9f9253847 100644 --- a/datafusion/expr/src/udf.rs +++ b/datafusion/expr/src/udf.rs @@ -749,14 +749,12 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl { self.inner.return_type_from_exprs(args, schema, arg_types) } - fn invoke(&self, args: &[ColumnarValue]) -> Result { - #[allow(deprecated)] - self.inner.invoke(args) - } - - fn invoke_no_args(&self, number_rows: usize) -> Result { - #[allow(deprecated)] - self.inner.invoke_no_args(number_rows) + fn invoke_batch( + &self, + args: &[ColumnarValue], + number_rows: usize, + ) -> Result { + self.inner.invoke_batch(args, number_rows) } fn simplify( diff --git a/datafusion/functions-nested/src/make_array.rs b/datafusion/functions-nested/src/make_array.rs index c2c6f24948b8..7aa3445f6858 100644 --- a/datafusion/functions-nested/src/make_array.rs +++ b/datafusion/functions-nested/src/make_array.rs @@ -98,14 +98,14 @@ impl ScalarUDFImpl for MakeArray { } } - fn invoke(&self, args: &[ColumnarValue]) -> Result { + fn invoke_batch( + &self, + args: &[ColumnarValue], + _number_rows: usize, + ) -> Result { make_scalar_function(make_array_inner)(args) } - fn invoke_no_args(&self, _number_rows: usize) -> Result { - make_scalar_function(make_array_inner)(&[]) - } - fn aliases(&self) -> &[String] { &self.aliases } diff --git a/datafusion/functions/src/core/version.rs b/datafusion/functions/src/core/version.rs index e7ac749ddddc..36cf07e9e5da 100644 --- a/datafusion/functions/src/core/version.rs +++ b/datafusion/functions/src/core/version.rs @@ -18,7 +18,7 @@ //! [`VersionFunc`]: Implementation of the `version` function. use arrow::datatypes::DataType; -use datafusion_common::{not_impl_err, plan_err, Result, ScalarValue}; +use datafusion_common::{internal_err, plan_err, Result, ScalarValue}; use datafusion_expr::scalar_doc_sections::DOC_SECTION_OTHER; use datafusion_expr::{ ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, @@ -66,11 +66,14 @@ impl ScalarUDFImpl for VersionFunc { } } - fn invoke(&self, _: &[ColumnarValue]) -> Result { - not_impl_err!("version does not take any arguments") - } - - fn invoke_no_args(&self, _: usize) -> Result { + fn invoke_batch( + &self, + args: &[ColumnarValue], + _number_rows: usize, + ) -> Result { + if !args.is_empty() { + return internal_err!("{} function does not accept arguments", self.name()); + } // TODO it would be great to add rust version and arrow version, // but that requires a `build.rs` script and/or adding a version const to arrow-rs let version = format!( diff --git a/datafusion/functions/src/math/pi.rs b/datafusion/functions/src/math/pi.rs index ea0f33161772..502429d0ca5d 100644 --- a/datafusion/functions/src/math/pi.rs +++ b/datafusion/functions/src/math/pi.rs @@ -20,7 +20,7 @@ use std::sync::OnceLock; use arrow::datatypes::DataType; use arrow::datatypes::DataType::Float64; -use datafusion_common::{not_impl_err, Result, ScalarValue}; +use datafusion_common::{internal_err, not_impl_err, Result, ScalarValue}; use datafusion_expr::scalar_doc_sections::DOC_SECTION_MATH; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::{ @@ -67,7 +67,14 @@ impl ScalarUDFImpl for PiFunc { not_impl_err!("{} function does not accept arguments", self.name()) } - fn invoke_no_args(&self, _number_rows: usize) -> Result { + fn invoke_batch( + &self, + args: &[ColumnarValue], + _number_rows: usize, + ) -> Result { + if !args.is_empty() { + return internal_err!("{} function does not accept arguments", self.name()); + } Ok(ColumnarValue::Scalar(ScalarValue::Float64(Some( std::f64::consts::PI, )))) diff --git a/datafusion/functions/src/math/random.rs b/datafusion/functions/src/math/random.rs index cf564e5328a5..cd92798d67dd 100644 --- a/datafusion/functions/src/math/random.rs +++ b/datafusion/functions/src/math/random.rs @@ -23,7 +23,7 @@ use arrow::datatypes::DataType; use arrow::datatypes::DataType::Float64; use rand::{thread_rng, Rng}; -use datafusion_common::{not_impl_err, Result}; +use datafusion_common::{internal_err, Result}; use datafusion_expr::scalar_doc_sections::DOC_SECTION_MATH; use datafusion_expr::ColumnarValue; use datafusion_expr::{Documentation, ScalarUDFImpl, Signature, Volatility}; @@ -64,11 +64,14 @@ impl ScalarUDFImpl for RandomFunc { Ok(Float64) } - fn invoke(&self, _args: &[ColumnarValue]) -> Result { - not_impl_err!("{} function does not accept arguments", self.name()) - } - - fn invoke_no_args(&self, num_rows: usize) -> Result { + fn invoke_batch( + &self, + args: &[ColumnarValue], + num_rows: usize, + ) -> Result { + if !args.is_empty() { + return internal_err!("{} function does not accept arguments", self.name()); + } let mut rng = thread_rng(); let mut values = vec![0.0; num_rows]; // Equivalent to set each element with rng.gen_range(0.0..1.0), but more efficient diff --git a/datafusion/functions/src/string/uuid.rs b/datafusion/functions/src/string/uuid.rs index 0fbdce16ccd1..2282670dac79 100644 --- a/datafusion/functions/src/string/uuid.rs +++ b/datafusion/functions/src/string/uuid.rs @@ -23,7 +23,7 @@ use arrow::datatypes::DataType; use arrow::datatypes::DataType::Utf8; use uuid::Uuid; -use datafusion_common::{not_impl_err, Result}; +use datafusion_common::{internal_err, Result}; use datafusion_expr::scalar_doc_sections::DOC_SECTION_STRING; use datafusion_expr::{ColumnarValue, Documentation, Volatility}; use datafusion_expr::{ScalarUDFImpl, Signature}; @@ -64,13 +64,16 @@ impl ScalarUDFImpl for UuidFunc { Ok(Utf8) } - fn invoke(&self, _args: &[ColumnarValue]) -> Result { - not_impl_err!("{} function does not accept arguments", self.name()) - } - /// Prints random (v4) uuid values per row /// uuid() = 'a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a11' - fn invoke_no_args(&self, num_rows: usize) -> Result { + fn invoke_batch( + &self, + args: &[ColumnarValue], + num_rows: usize, + ) -> Result { + if !args.is_empty() { + return internal_err!("{} function does not accept arguments", self.name()); + } let values = std::iter::repeat_with(|| Uuid::new_v4().to_string()).take(num_rows); let array = GenericStringArray::::from_iter_values(values); Ok(ColumnarValue::Array(Arc::new(array)))