From ec9f03c5ea2f36335a52f50f774361124beb0d94 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E6=9E=97=E4=BC=9F?= Date: Fri, 14 Feb 2025 11:27:40 +0800 Subject: [PATCH 1/4] Migrate math functions to implement invoke_with_args --- datafusion/functions/src/math/abs.rs | 11 ++++------- datafusion/functions/src/math/cot.rs | 10 +++------- datafusion/functions/src/math/factorial.rs | 11 ++++------- datafusion/functions/src/math/gcd.rs | 11 ++++------- datafusion/functions/src/math/iszero.rs | 11 ++++------- datafusion/functions/src/math/lcm.rs | 11 ++++------- datafusion/functions/src/math/log.rs | 11 ++++------- datafusion/functions/src/math/nans.rs | 10 +++------- datafusion/functions/src/math/nanvl.rs | 11 ++++------- datafusion/functions/src/math/pi.rs | 11 ++++------- datafusion/functions/src/math/power.rs | 12 +++++------- datafusion/functions/src/math/random.rs | 12 ++++-------- datafusion/functions/src/math/round.rs | 11 ++++------- datafusion/functions/src/math/signum.rs | 11 ++++------- datafusion/functions/src/math/trunc.rs | 11 ++++------- 15 files changed, 59 insertions(+), 106 deletions(-) diff --git a/datafusion/functions/src/math/abs.rs b/datafusion/functions/src/math/abs.rs index ff6a82113262..0c686a59016a 100644 --- a/datafusion/functions/src/math/abs.rs +++ b/datafusion/functions/src/math/abs.rs @@ -32,7 +32,8 @@ use datafusion_common::{ use datafusion_expr::interval_arithmetic::Interval; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, }; use datafusion_macros::user_doc; @@ -168,12 +169,8 @@ impl ScalarUDFImpl for AbsFunc { } } - fn invoke_batch( - &self, - args: &[ColumnarValue], - _number_rows: usize, - ) -> Result { - let args = ColumnarValue::values_to_arrays(args)?; + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let args = ColumnarValue::values_to_arrays(&args.args)?; let [input] = take_function_args(self.name(), args)?; let input_data_type = input.data_type(); diff --git a/datafusion/functions/src/math/cot.rs b/datafusion/functions/src/math/cot.rs index 8b4f9317fe5f..4e56212ddbee 100644 --- a/datafusion/functions/src/math/cot.rs +++ b/datafusion/functions/src/math/cot.rs @@ -24,7 +24,7 @@ use arrow::datatypes::{DataType, Float32Type, Float64Type}; use crate::utils::make_scalar_function; use datafusion_common::{exec_err, Result}; -use datafusion_expr::{ColumnarValue, Documentation}; +use datafusion_expr::{ColumnarValue, Documentation, ScalarFunctionArgs}; use datafusion_expr::{ScalarUDFImpl, Signature, Volatility}; use datafusion_macros::user_doc; @@ -87,12 +87,8 @@ impl ScalarUDFImpl for CotFunc { self.doc() } - fn invoke_batch( - &self, - args: &[ColumnarValue], - _number_rows: usize, - ) -> Result { - make_scalar_function(cot, vec![])(args) + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(cot, vec![])(&args.args) } } diff --git a/datafusion/functions/src/math/factorial.rs b/datafusion/functions/src/math/factorial.rs index 18f10863a01b..c2ac21b78f21 100644 --- a/datafusion/functions/src/math/factorial.rs +++ b/datafusion/functions/src/math/factorial.rs @@ -30,7 +30,8 @@ use datafusion_common::{ arrow_datafusion_err, exec_err, internal_datafusion_err, DataFusionError, Result, }; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, }; use datafusion_macros::user_doc; @@ -76,12 +77,8 @@ impl ScalarUDFImpl for FactorialFunc { Ok(Int64) } - fn invoke_batch( - &self, - args: &[ColumnarValue], - _number_rows: usize, - ) -> Result { - make_scalar_function(factorial, vec![])(args) + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(factorial, vec![])(&args.args) } fn documentation(&self) -> Option<&Documentation> { diff --git a/datafusion/functions/src/math/gcd.rs b/datafusion/functions/src/math/gcd.rs index 14503701f661..911e00308ab7 100644 --- a/datafusion/functions/src/math/gcd.rs +++ b/datafusion/functions/src/math/gcd.rs @@ -29,7 +29,8 @@ use datafusion_common::{ arrow_datafusion_err, exec_err, internal_datafusion_err, DataFusionError, Result, }; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, }; use datafusion_macros::user_doc; @@ -77,12 +78,8 @@ impl ScalarUDFImpl for GcdFunc { Ok(Int64) } - fn invoke_batch( - &self, - args: &[ColumnarValue], - _number_rows: usize, - ) -> Result { - make_scalar_function(gcd, vec![])(args) + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(gcd, vec![])(&args.args) } fn documentation(&self) -> Option<&Documentation> { diff --git a/datafusion/functions/src/math/iszero.rs b/datafusion/functions/src/math/iszero.rs index 8e72ee285518..bc12dfb7898e 100644 --- a/datafusion/functions/src/math/iszero.rs +++ b/datafusion/functions/src/math/iszero.rs @@ -25,7 +25,8 @@ use arrow::datatypes::{DataType, Float32Type, Float64Type}; use datafusion_common::{exec_err, Result}; use datafusion_expr::TypeSignature::Exact; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, }; use datafusion_macros::user_doc; @@ -77,12 +78,8 @@ impl ScalarUDFImpl for IsZeroFunc { Ok(Boolean) } - fn invoke_batch( - &self, - args: &[ColumnarValue], - _number_rows: usize, - ) -> Result { - make_scalar_function(iszero, vec![])(args) + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(iszero, vec![])(&args.args) } fn documentation(&self) -> Option<&Documentation> { diff --git a/datafusion/functions/src/math/lcm.rs b/datafusion/functions/src/math/lcm.rs index c2c72c89841d..fc6bf9461f28 100644 --- a/datafusion/functions/src/math/lcm.rs +++ b/datafusion/functions/src/math/lcm.rs @@ -27,7 +27,8 @@ use datafusion_common::{ arrow_datafusion_err, exec_err, internal_datafusion_err, DataFusionError, Result, }; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, }; use datafusion_macros::user_doc; @@ -78,12 +79,8 @@ impl ScalarUDFImpl for LcmFunc { Ok(Int64) } - fn invoke_batch( - &self, - args: &[ColumnarValue], - _number_rows: usize, - ) -> Result { - make_scalar_function(lcm, vec![])(args) + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(lcm, vec![])(&args.args) } fn documentation(&self) -> Option<&Documentation> { diff --git a/datafusion/functions/src/math/log.rs b/datafusion/functions/src/math/log.rs index 88a624806874..9334f7a5c6f1 100644 --- a/datafusion/functions/src/math/log.rs +++ b/datafusion/functions/src/math/log.rs @@ -31,7 +31,8 @@ use datafusion_expr::expr::ScalarFunction; use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo}; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::{ - lit, ColumnarValue, Documentation, Expr, ScalarUDF, TypeSignature::*, + lit, ColumnarValue, Documentation, Expr, ScalarFunctionArgs, ScalarUDF, + TypeSignature::*, }; use datafusion_expr::{ScalarUDFImpl, Signature, Volatility}; use datafusion_macros::user_doc; @@ -117,12 +118,8 @@ impl ScalarUDFImpl for LogFunc { } // Support overloaded log(base, x) and log(x) which defaults to log(10, x) - fn invoke_batch( - &self, - args: &[ColumnarValue], - _number_rows: usize, - ) -> Result { - let args = ColumnarValue::values_to_arrays(args)?; + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let args = ColumnarValue::values_to_arrays(&args.args)?; let mut base = ColumnarValue::Scalar(ScalarValue::Float32(Some(10.0))); diff --git a/datafusion/functions/src/math/nans.rs b/datafusion/functions/src/math/nans.rs index 30c920c29a21..34a5c2a1c16b 100644 --- a/datafusion/functions/src/math/nans.rs +++ b/datafusion/functions/src/math/nans.rs @@ -19,7 +19,7 @@ use arrow::datatypes::{DataType, Float32Type, Float64Type}; use datafusion_common::{exec_err, Result}; -use datafusion_expr::{ColumnarValue, TypeSignature}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, TypeSignature}; use arrow::array::{ArrayRef, AsArray, BooleanArray}; use datafusion_expr::{Documentation, ScalarUDFImpl, Signature, Volatility}; @@ -75,12 +75,8 @@ impl ScalarUDFImpl for IsNanFunc { Ok(DataType::Boolean) } - fn invoke_batch( - &self, - args: &[ColumnarValue], - _number_rows: usize, - ) -> Result { - let args = ColumnarValue::values_to_arrays(args)?; + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let args = ColumnarValue::values_to_arrays(&args.args)?; let arr: ArrayRef = match args[0].data_type() { DataType::Float64 => Arc::new(BooleanArray::from_unary( diff --git a/datafusion/functions/src/math/nanvl.rs b/datafusion/functions/src/math/nanvl.rs index 33823acce751..9effb82896ee 100644 --- a/datafusion/functions/src/math/nanvl.rs +++ b/datafusion/functions/src/math/nanvl.rs @@ -26,7 +26,8 @@ use arrow::datatypes::{DataType, Float32Type, Float64Type}; use datafusion_common::{exec_err, DataFusionError, Result}; use datafusion_expr::TypeSignature::Exact; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, }; use datafusion_macros::user_doc; @@ -87,12 +88,8 @@ impl ScalarUDFImpl for NanvlFunc { } } - fn invoke_batch( - &self, - args: &[ColumnarValue], - _number_rows: usize, - ) -> Result { - make_scalar_function(nanvl, vec![])(args) + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(nanvl, vec![])(&args.args) } fn documentation(&self) -> Option<&Documentation> { diff --git a/datafusion/functions/src/math/pi.rs b/datafusion/functions/src/math/pi.rs index 06f7a01544f8..5339a9b14a28 100644 --- a/datafusion/functions/src/math/pi.rs +++ b/datafusion/functions/src/math/pi.rs @@ -22,7 +22,8 @@ use arrow::datatypes::DataType::Float64; use datafusion_common::{internal_err, Result, ScalarValue}; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, }; use datafusion_macros::user_doc; @@ -67,12 +68,8 @@ impl ScalarUDFImpl for PiFunc { Ok(Float64) } - fn invoke_batch( - &self, - args: &[ColumnarValue], - _number_rows: usize, - ) -> Result { - if !args.is_empty() { + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + if !args.args.is_empty() { return internal_err!("{} function does not accept arguments", self.name()); } Ok(ColumnarValue::Scalar(ScalarValue::Float64(Some( diff --git a/datafusion/functions/src/math/power.rs b/datafusion/functions/src/math/power.rs index 7fab858d34a0..b667414e8a25 100644 --- a/datafusion/functions/src/math/power.rs +++ b/datafusion/functions/src/math/power.rs @@ -29,7 +29,9 @@ use datafusion_common::{ }; use datafusion_expr::expr::ScalarFunction; use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo}; -use datafusion_expr::{ColumnarValue, Documentation, Expr, ScalarUDF, TypeSignature}; +use datafusion_expr::{ + ColumnarValue, Documentation, Expr, ScalarFunctionArgs, ScalarUDF, TypeSignature, +}; use datafusion_expr::{ScalarUDFImpl, Signature, Volatility}; use datafusion_macros::user_doc; @@ -91,12 +93,8 @@ impl ScalarUDFImpl for PowerFunc { &self.aliases } - fn invoke_batch( - &self, - args: &[ColumnarValue], - _number_rows: usize, - ) -> Result { - let args = ColumnarValue::values_to_arrays(args)?; + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let args = ColumnarValue::values_to_arrays(&args.args)?; let arr: ArrayRef = match args[0].data_type() { DataType::Float64 => { diff --git a/datafusion/functions/src/math/random.rs b/datafusion/functions/src/math/random.rs index 197d065ea408..607f9fb09f2a 100644 --- a/datafusion/functions/src/math/random.rs +++ b/datafusion/functions/src/math/random.rs @@ -24,7 +24,7 @@ use arrow::datatypes::DataType::Float64; use rand::{thread_rng, Rng}; use datafusion_common::{internal_err, Result}; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_expr::{Documentation, ScalarUDFImpl, Signature, Volatility}; use datafusion_macros::user_doc; @@ -70,16 +70,12 @@ impl ScalarUDFImpl for RandomFunc { Ok(Float64) } - fn invoke_batch( - &self, - args: &[ColumnarValue], - num_rows: usize, - ) -> Result { - if !args.is_empty() { + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + if !args.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]; + let mut values = vec![0.0; args.number_rows]; // Equivalent to set each element with rng.gen_range(0.0..1.0), but more efficient rng.fill(&mut values[..]); let array = Float64Array::from(values); diff --git a/datafusion/functions/src/math/round.rs b/datafusion/functions/src/math/round.rs index b3442c321c99..fc87b7e63a62 100644 --- a/datafusion/functions/src/math/round.rs +++ b/datafusion/functions/src/math/round.rs @@ -28,7 +28,8 @@ use datafusion_common::{exec_datafusion_err, exec_err, Result, ScalarValue}; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::TypeSignature::Exact; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, }; use datafusion_macros::user_doc; @@ -90,12 +91,8 @@ impl ScalarUDFImpl for RoundFunc { } } - fn invoke_batch( - &self, - args: &[ColumnarValue], - _number_rows: usize, - ) -> Result { - make_scalar_function(round, vec![])(args) + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(round, vec![])(&args.args) } fn output_ordering(&self, input: &[ExprProperties]) -> Result { diff --git a/datafusion/functions/src/math/signum.rs b/datafusion/functions/src/math/signum.rs index f68834db375e..bf8c5e74975d 100644 --- a/datafusion/functions/src/math/signum.rs +++ b/datafusion/functions/src/math/signum.rs @@ -25,7 +25,8 @@ use arrow::datatypes::{DataType, Float32Type, Float64Type}; use datafusion_common::{exec_err, Result}; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, }; use datafusion_macros::user_doc; @@ -88,12 +89,8 @@ impl ScalarUDFImpl for SignumFunc { Ok(input[0].sort_properties) } - fn invoke_batch( - &self, - args: &[ColumnarValue], - _number_rows: usize, - ) -> Result { - make_scalar_function(signum, vec![])(args) + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(signum, vec![])(&args.args) } fn documentation(&self) -> Option<&Documentation> { diff --git a/datafusion/functions/src/math/trunc.rs b/datafusion/functions/src/math/trunc.rs index 8d791370d7f8..2ac291204a0b 100644 --- a/datafusion/functions/src/math/trunc.rs +++ b/datafusion/functions/src/math/trunc.rs @@ -28,7 +28,8 @@ use datafusion_common::{exec_err, Result}; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::TypeSignature::Exact; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, }; use datafusion_macros::user_doc; @@ -99,12 +100,8 @@ impl ScalarUDFImpl for TruncFunc { } } - fn invoke_batch( - &self, - args: &[ColumnarValue], - _number_rows: usize, - ) -> Result { - make_scalar_function(trunc, vec![])(args) + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(trunc, vec![])(&args.args) } fn output_ordering(&self, input: &[ExprProperties]) -> Result { From fb46b3bc8d8dde582b53e3f16412cccaabd41b57 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E6=9E=97=E4=BC=9F?= Date: Fri, 14 Feb 2025 13:57:23 +0800 Subject: [PATCH 2/4] Fix unit tests --- datafusion/functions/src/math/log.rs | 165 ++++++++++++++---------- datafusion/functions/src/math/power.rs | 34 +++-- datafusion/functions/src/math/signum.rs | 22 ++-- 3 files changed, 136 insertions(+), 85 deletions(-) diff --git a/datafusion/functions/src/math/log.rs b/datafusion/functions/src/math/log.rs index 9334f7a5c6f1..fd135f4c5ec0 100644 --- a/datafusion/functions/src/math/log.rs +++ b/datafusion/functions/src/math/log.rs @@ -264,34 +264,44 @@ mod tests { #[test] #[should_panic] fn test_log_invalid_base_type() { - let args = [ - ColumnarValue::Array(Arc::new(Float64Array::from(vec![ - 10.0, 100.0, 1000.0, 10000.0, - ]))), // num - ColumnarValue::Array(Arc::new(Int64Array::from(vec![5, 10, 15, 20]))), - ]; - #[allow(deprecated)] // TODO: migrate to invoke_with_args - let _ = LogFunc::new().invoke_batch(&args, 4); + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(Arc::new(Float64Array::from(vec![ + 10.0, 100.0, 1000.0, 10000.0, + ]))), // num + ColumnarValue::Array(Arc::new(Int64Array::from(vec![5, 10, 15, 20]))), + ], + number_rows: 4, + return_type: &DataType::Float64, + }; + let _ = LogFunc::new().invoke_with_args(args); } #[test] fn test_log_invalid_value() { - let args = [ - ColumnarValue::Array(Arc::new(Int64Array::from(vec![10]))), // num - ]; - #[allow(deprecated)] // TODO: migrate to invoke_with_args - let result = LogFunc::new().invoke_batch(&args, 1); + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(Arc::new(Int64Array::from(vec![10]))), // num + ], + number_rows: 1, + return_type: &DataType::Float64, + }; + + let result = LogFunc::new().invoke_with_args(args); result.expect_err("expected error"); } #[test] fn test_log_scalar_f32_unary() { - let args = [ - ColumnarValue::Scalar(ScalarValue::Float32(Some(10.0))), // num - ]; - #[allow(deprecated)] // TODO: migrate to invoke_with_args + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Scalar(ScalarValue::Float32(Some(10.0))), // num + ], + number_rows: 1, + return_type: &DataType::Float32, + }; let result = LogFunc::new() - .invoke_batch(&args, 1) + .invoke_with_args(args) .expect("failed to initialize function log"); match result { @@ -310,12 +320,15 @@ mod tests { #[test] fn test_log_scalar_f64_unary() { - let args = [ - ColumnarValue::Scalar(ScalarValue::Float64(Some(10.0))), // num - ]; - #[allow(deprecated)] // TODO: migrate to invoke_with_args + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Scalar(ScalarValue::Float64(Some(10.0))), // num + ], + number_rows: 1, + return_type: &DataType::Float64, + }; let result = LogFunc::new() - .invoke_batch(&args, 1) + .invoke_with_args(args) .expect("failed to initialize function log"); match result { @@ -334,13 +347,16 @@ mod tests { #[test] fn test_log_scalar_f32() { - let args = [ - ColumnarValue::Scalar(ScalarValue::Float32(Some(2.0))), // num - ColumnarValue::Scalar(ScalarValue::Float32(Some(32.0))), // num - ]; - #[allow(deprecated)] // TODO: migrate to invoke_with_args + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Scalar(ScalarValue::Float32(Some(2.0))), // num + ColumnarValue::Scalar(ScalarValue::Float32(Some(32.0))), // num + ], + number_rows: 1, + return_type: &DataType::Float32, + }; let result = LogFunc::new() - .invoke_batch(&args, 1) + .invoke_with_args(args) .expect("failed to initialize function log"); match result { @@ -359,13 +375,16 @@ mod tests { #[test] fn test_log_scalar_f64() { - let args = [ - ColumnarValue::Scalar(ScalarValue::Float64(Some(2.0))), // num - ColumnarValue::Scalar(ScalarValue::Float64(Some(64.0))), // num - ]; - #[allow(deprecated)] // TODO: migrate to invoke_with_args + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Scalar(ScalarValue::Float64(Some(2.0))), // num + ColumnarValue::Scalar(ScalarValue::Float64(Some(64.0))), // num + ], + number_rows: 1, + return_type: &DataType::Float64, + }; let result = LogFunc::new() - .invoke_batch(&args, 1) + .invoke_with_args(args) .expect("failed to initialize function log"); match result { @@ -384,14 +403,17 @@ mod tests { #[test] fn test_log_f64_unary() { - let args = [ - ColumnarValue::Array(Arc::new(Float64Array::from(vec![ - 10.0, 100.0, 1000.0, 10000.0, - ]))), // num - ]; - #[allow(deprecated)] // TODO: migrate to invoke_with_args + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(Arc::new(Float64Array::from(vec![ + 10.0, 100.0, 1000.0, 10000.0, + ]))), // num + ], + number_rows: 4, + return_type: &DataType::Float64, + }; let result = LogFunc::new() - .invoke_batch(&args, 4) + .invoke_with_args(args) .expect("failed to initialize function log"); match result { @@ -413,14 +435,17 @@ mod tests { #[test] fn test_log_f32_unary() { - let args = [ - ColumnarValue::Array(Arc::new(Float32Array::from(vec![ - 10.0, 100.0, 1000.0, 10000.0, - ]))), // num - ]; - #[allow(deprecated)] // TODO: migrate to invoke_with_args + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(Arc::new(Float32Array::from(vec![ + 10.0, 100.0, 1000.0, 10000.0, + ]))), // num + ], + number_rows: 4, + return_type: &DataType::Float32, + }; let result = LogFunc::new() - .invoke_batch(&args, 4) + .invoke_with_args(args) .expect("failed to initialize function log"); match result { @@ -442,15 +467,20 @@ mod tests { #[test] fn test_log_f64() { - let args = [ - ColumnarValue::Array(Arc::new(Float64Array::from(vec![2.0, 2.0, 3.0, 5.0]))), // base - ColumnarValue::Array(Arc::new(Float64Array::from(vec![ - 8.0, 4.0, 81.0, 625.0, - ]))), // num - ]; - #[allow(deprecated)] // TODO: migrate to invoke_with_args + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(Arc::new(Float64Array::from(vec![ + 2.0, 2.0, 3.0, 5.0, + ]))), // base + ColumnarValue::Array(Arc::new(Float64Array::from(vec![ + 8.0, 4.0, 81.0, 625.0, + ]))), // num + ], + number_rows: 4, + return_type: &DataType::Float64, + }; let result = LogFunc::new() - .invoke_batch(&args, 4) + .invoke_with_args(args) .expect("failed to initialize function log"); match result { @@ -472,15 +502,20 @@ mod tests { #[test] fn test_log_f32() { - let args = [ - ColumnarValue::Array(Arc::new(Float32Array::from(vec![2.0, 2.0, 3.0, 5.0]))), // base - ColumnarValue::Array(Arc::new(Float32Array::from(vec![ - 8.0, 4.0, 81.0, 625.0, - ]))), // num - ]; - #[allow(deprecated)] // TODO: migrate to invoke_with_args + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(Arc::new(Float32Array::from(vec![ + 2.0, 2.0, 3.0, 5.0, + ]))), // base + ColumnarValue::Array(Arc::new(Float32Array::from(vec![ + 8.0, 4.0, 81.0, 625.0, + ]))), // num + ], + number_rows: 4, + return_type: &DataType::Float32, + }; let result = LogFunc::new() - .invoke_batch(&args, 4) + .invoke_with_args(args) .expect("failed to initialize function log"); match result { diff --git a/datafusion/functions/src/math/power.rs b/datafusion/functions/src/math/power.rs index b667414e8a25..028ec2fef793 100644 --- a/datafusion/functions/src/math/power.rs +++ b/datafusion/functions/src/math/power.rs @@ -193,13 +193,20 @@ mod tests { #[test] fn test_power_f64() { - let args = [ - ColumnarValue::Array(Arc::new(Float64Array::from(vec![2.0, 2.0, 3.0, 5.0]))), // base - ColumnarValue::Array(Arc::new(Float64Array::from(vec![3.0, 2.0, 4.0, 4.0]))), // exponent - ]; - #[allow(deprecated)] // TODO: migrate to invoke_with_args + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(Arc::new(Float64Array::from(vec![ + 2.0, 2.0, 3.0, 5.0, + ]))), // base + ColumnarValue::Array(Arc::new(Float64Array::from(vec![ + 3.0, 2.0, 4.0, 4.0, + ]))), // exponent + ], + number_rows: 4, + return_type: &DataType::Float64, + }; let result = PowerFunc::new() - .invoke_batch(&args, 4) + .invoke_with_args(args) .expect("failed to initialize function power"); match result { @@ -220,13 +227,16 @@ mod tests { #[test] fn test_power_i64() { - let args = [ - ColumnarValue::Array(Arc::new(Int64Array::from(vec![2, 2, 3, 5]))), // base - ColumnarValue::Array(Arc::new(Int64Array::from(vec![3, 2, 4, 4]))), // exponent - ]; - #[allow(deprecated)] // TODO: migrate to invoke_with_args + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(Arc::new(Int64Array::from(vec![2, 2, 3, 5]))), // base + ColumnarValue::Array(Arc::new(Int64Array::from(vec![3, 2, 4, 4]))), // exponent + ], + number_rows: 4, + return_type: &DataType::Int64, + }; let result = PowerFunc::new() - .invoke_batch(&args, 4) + .invoke_with_args(args) .expect("failed to initialize function power"); match result { diff --git a/datafusion/functions/src/math/signum.rs b/datafusion/functions/src/math/signum.rs index bf8c5e74975d..4aa0a598fdb1 100644 --- a/datafusion/functions/src/math/signum.rs +++ b/datafusion/functions/src/math/signum.rs @@ -138,9 +138,9 @@ mod test { use std::sync::Arc; use arrow::array::{Float32Array, Float64Array}; - + use arrow::datatypes::DataType; use datafusion_common::cast::{as_float32_array, as_float64_array}; - use datafusion_expr::{ColumnarValue, ScalarUDFImpl}; + use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; use crate::math::signum::SignumFunc; @@ -157,10 +157,13 @@ mod test { f32::INFINITY, f32::NEG_INFINITY, ])); - let batch_size = array.len(); - #[allow(deprecated)] // TODO: migrate to invoke_with_args + let args = ScalarFunctionArgs { + args: vec![ColumnarValue::Array(array.clone())], + number_rows: array.len(), + return_type: &DataType::Float32, + }; let result = SignumFunc::new() - .invoke_batch(&[ColumnarValue::Array(array)], batch_size) + .invoke_with_args(args) .expect("failed to initialize function signum"); match result { @@ -198,10 +201,13 @@ mod test { f64::INFINITY, f64::NEG_INFINITY, ])); - let batch_size = array.len(); - #[allow(deprecated)] // TODO: migrate to invoke_with_args + let args = ScalarFunctionArgs { + args: vec![ColumnarValue::Array(array.clone())], + number_rows: array.len(), + return_type: &DataType::Float64, + }; let result = SignumFunc::new() - .invoke_batch(&[ColumnarValue::Array(array)], batch_size) + .invoke_with_args(args) .expect("failed to initialize function signum"); match result { From 0eed501eb315d4fbd01148d2a0e614e80145cd74 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E6=9E=97=E4=BC=9F?= Date: Fri, 14 Feb 2025 14:15:05 +0800 Subject: [PATCH 3/4] Fix clippy lints --- datafusion/functions/src/math/signum.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/functions/src/math/signum.rs b/datafusion/functions/src/math/signum.rs index 4aa0a598fdb1..eb2cb00161d8 100644 --- a/datafusion/functions/src/math/signum.rs +++ b/datafusion/functions/src/math/signum.rs @@ -158,7 +158,7 @@ mod test { f32::NEG_INFINITY, ])); let args = ScalarFunctionArgs { - args: vec![ColumnarValue::Array(array.clone())], + args: vec![ColumnarValue::Array(Arc::clone(&array))], number_rows: array.len(), return_type: &DataType::Float32, }; @@ -202,7 +202,7 @@ mod test { f64::NEG_INFINITY, ])); let args = ScalarFunctionArgs { - args: vec![ColumnarValue::Array(array.clone())], + args: vec![ColumnarValue::Array(Arc::clone(&array))], number_rows: array.len(), return_type: &DataType::Float64, }; From 50c76545bde20a6162aa9e706cd3eb7ec68f253e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E6=9E=97=E4=BC=9F?= Date: Fri, 14 Feb 2025 14:37:11 +0800 Subject: [PATCH 4/4] Fix compile --- datafusion/functions/src/math/signum.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/datafusion/functions/src/math/signum.rs b/datafusion/functions/src/math/signum.rs index eb2cb00161d8..ba5422afa768 100644 --- a/datafusion/functions/src/math/signum.rs +++ b/datafusion/functions/src/math/signum.rs @@ -137,7 +137,7 @@ pub fn signum(args: &[ArrayRef]) -> Result { mod test { use std::sync::Arc; - use arrow::array::{Float32Array, Float64Array}; + use arrow::array::{ArrayRef, Float32Array, Float64Array}; use arrow::datatypes::DataType; use datafusion_common::cast::{as_float32_array, as_float64_array}; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; @@ -158,7 +158,7 @@ mod test { f32::NEG_INFINITY, ])); let args = ScalarFunctionArgs { - args: vec![ColumnarValue::Array(Arc::clone(&array))], + args: vec![ColumnarValue::Array(Arc::clone(&array) as ArrayRef)], number_rows: array.len(), return_type: &DataType::Float32, }; @@ -202,7 +202,7 @@ mod test { f64::NEG_INFINITY, ])); let args = ScalarFunctionArgs { - args: vec![ColumnarValue::Array(Arc::clone(&array))], + args: vec![ColumnarValue::Array(Arc::clone(&array) as ArrayRef)], number_rows: array.len(), return_type: &DataType::Float64, };