From ac999a0295abd4d94ba0e09a462399568694e640 Mon Sep 17 00:00:00 2001 From: irenjj Date: Sun, 10 Nov 2024 23:35:50 +0800 Subject: [PATCH 1/2] migrate UDF invoke to invoke_batch --- datafusion/functions/src/datetime/date_bin.rs | 275 +++++++++++------- .../functions/src/datetime/date_trunc.rs | 24 +- .../functions/src/datetime/make_date.rs | 101 ++++--- datafusion/functions/src/datetime/to_char.rs | 57 ++-- datafusion/functions/src/datetime/to_date.rs | 58 ++-- datafusion/functions/src/string/concat.rs | 3 +- datafusion/functions/src/string/concat_ws.rs | 6 +- datafusion/functions/src/string/contains.rs | 3 +- datafusion/functions/src/string/lower.rs | 3 +- datafusion/functions/src/string/upper.rs | 3 +- 10 files changed, 306 insertions(+), 227 deletions(-) diff --git a/datafusion/functions/src/datetime/date_bin.rs b/datafusion/functions/src/datetime/date_bin.rs index 065201e1caa39..4385b1fedf0b1 100644 --- a/datafusion/functions/src/datetime/date_bin.rs +++ b/datafusion/functions/src/datetime/date_bin.rs @@ -491,50 +491,67 @@ mod tests { use chrono::TimeDelta; #[test] - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch fn test_date_bin() { - let res = DateBinFunc::new().invoke(&[ - ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(IntervalDayTime { - days: 0, - milliseconds: 1, - }))), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ]); + let res = DateBinFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some( + IntervalDayTime { + days: 0, + milliseconds: 1, + }, + ))), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ], + 0, + ); assert!(res.is_ok()); let timestamps = Arc::new((1..6).map(Some).collect::()); - let res = DateBinFunc::new().invoke(&[ - ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(IntervalDayTime { - days: 0, - milliseconds: 1, - }))), - ColumnarValue::Array(timestamps), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ]); + let res = DateBinFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some( + IntervalDayTime { + days: 0, + milliseconds: 1, + }, + ))), + ColumnarValue::Array(timestamps), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ], + 0, + ); assert!(res.is_ok()); - let res = DateBinFunc::new().invoke(&[ - ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(IntervalDayTime { - days: 0, - milliseconds: 1, - }))), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ]); + let res = DateBinFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some( + IntervalDayTime { + days: 0, + milliseconds: 1, + }, + ))), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ], + 0, + ); assert!(res.is_ok()); // stride supports month-day-nano - let res = DateBinFunc::new().invoke(&[ - ColumnarValue::Scalar(ScalarValue::IntervalMonthDayNano(Some( - IntervalMonthDayNano { - months: 0, - days: 0, - nanoseconds: 1, - }, - ))), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ]); + let res = DateBinFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::IntervalMonthDayNano(Some( + IntervalMonthDayNano { + months: 0, + days: 0, + nanoseconds: 1, + }, + ))), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ], + 0, + ); assert!(res.is_ok()); // @@ -542,99 +559,129 @@ mod tests { // // invalid number of arguments - let res = DateBinFunc::new().invoke(&[ColumnarValue::Scalar( - ScalarValue::IntervalDayTime(Some(IntervalDayTime { - days: 0, - milliseconds: 1, - })), - )]); + let res = DateBinFunc::new().invoke_batch( + &[ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some( + IntervalDayTime { + days: 0, + milliseconds: 1, + }, + )))], + 0, + ); assert_eq!( res.err().unwrap().strip_backtrace(), "Execution error: DATE_BIN expected two or three arguments" ); // stride: invalid type - let res = DateBinFunc::new().invoke(&[ - ColumnarValue::Scalar(ScalarValue::IntervalYearMonth(Some(1))), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ]); + let res = DateBinFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::IntervalYearMonth(Some(1))), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ], + 0, + ); assert_eq!( res.err().unwrap().strip_backtrace(), "Execution error: DATE_BIN expects stride argument to be an INTERVAL but got Interval(YearMonth)" ); // stride: invalid value - let res = DateBinFunc::new().invoke(&[ - ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(IntervalDayTime { - days: 0, - milliseconds: 0, - }))), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ]); + let res = DateBinFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some( + IntervalDayTime { + days: 0, + milliseconds: 0, + }, + ))), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ], + 0, + ); assert_eq!( res.err().unwrap().strip_backtrace(), "Execution error: DATE_BIN stride must be non-zero" ); // stride: overflow of day-time interval - let res = DateBinFunc::new().invoke(&[ - ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some( - IntervalDayTime::MAX, - ))), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ]); + let res = DateBinFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some( + IntervalDayTime::MAX, + ))), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ], + 0, + ); assert_eq!( res.err().unwrap().strip_backtrace(), "Execution error: DATE_BIN stride argument is too large" ); // stride: overflow of month-day-nano interval - let res = DateBinFunc::new().invoke(&[ - ColumnarValue::Scalar(ScalarValue::new_interval_mdn(0, i32::MAX, 1)), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ]); + let res = DateBinFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::new_interval_mdn(0, i32::MAX, 1)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ], + 0, + ); assert_eq!( res.err().unwrap().strip_backtrace(), "Execution error: DATE_BIN stride argument is too large" ); // stride: month intervals - let res = DateBinFunc::new().invoke(&[ - ColumnarValue::Scalar(ScalarValue::new_interval_mdn(1, 1, 1)), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ]); + let res = DateBinFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::new_interval_mdn(1, 1, 1)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ], + 0, + ); assert_eq!( res.err().unwrap().strip_backtrace(), "This feature is not implemented: DATE_BIN stride does not support combination of month, day and nanosecond intervals" ); // origin: invalid type - let res = DateBinFunc::new().invoke(&[ - ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(IntervalDayTime { - days: 0, - milliseconds: 1, - }))), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(Some(1), None)), - ]); + let res = DateBinFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some( + IntervalDayTime { + days: 0, + milliseconds: 1, + }, + ))), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(Some(1), None)), + ], + 0, + ); assert_eq!( res.err().unwrap().strip_backtrace(), "Execution error: DATE_BIN expects origin argument to be a TIMESTAMP with nanosecond precision but got Timestamp(Microsecond, None)" ); - let res = DateBinFunc::new().invoke(&[ - ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(IntervalDayTime { - days: 0, - milliseconds: 1, - }))), - ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(Some(1), None)), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ]); + let res = DateBinFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some( + IntervalDayTime { + days: 0, + milliseconds: 1, + }, + ))), + ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(Some(1), None)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ], + 0, + ); assert!(res.is_ok()); // unsupported array type for stride @@ -648,11 +695,14 @@ mod tests { }) .collect::(), ); - let res = DateBinFunc::new().invoke(&[ - ColumnarValue::Array(intervals), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ]); + let res = DateBinFunc::new().invoke_batch( + &[ + ColumnarValue::Array(intervals), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ], + 0, + ); assert_eq!( res.err().unwrap().strip_backtrace(), "This feature is not implemented: DATE_BIN only supports literal values for the stride argument, not arrays" @@ -660,14 +710,19 @@ mod tests { // unsupported array type for origin let timestamps = Arc::new((1..6).map(Some).collect::()); - let res = DateBinFunc::new().invoke(&[ - ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(IntervalDayTime { - days: 0, - milliseconds: 1, - }))), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ColumnarValue::Array(timestamps), - ]); + let res = DateBinFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some( + IntervalDayTime { + days: 0, + milliseconds: 1, + }, + ))), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ColumnarValue::Array(timestamps), + ], + 0, + ); assert_eq!( res.err().unwrap().strip_backtrace(), "This feature is not implemented: DATE_BIN only supports literal values for the origin argument, not arrays" @@ -782,16 +837,18 @@ mod tests { .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) .collect::() .with_timezone_opt(tz_opt.clone()); - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch let result = DateBinFunc::new() - .invoke(&[ - ColumnarValue::Scalar(ScalarValue::new_interval_dt(1, 0)), - ColumnarValue::Array(Arc::new(input)), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond( - Some(string_to_timestamp_nanos(origin).unwrap()), - tz_opt.clone(), - )), - ]) + .invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::new_interval_dt(1, 0)), + ColumnarValue::Array(Arc::new(input)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond( + Some(string_to_timestamp_nanos(origin).unwrap()), + tz_opt.clone(), + )), + ], + 0, + ) .unwrap(); if let ColumnarValue::Array(result) = result { assert_eq!( diff --git a/datafusion/functions/src/datetime/date_trunc.rs b/datafusion/functions/src/datetime/date_trunc.rs index f8abef601f70c..a0d472d81d8f2 100644 --- a/datafusion/functions/src/datetime/date_trunc.rs +++ b/datafusion/functions/src/datetime/date_trunc.rs @@ -724,12 +724,14 @@ mod tests { .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) .collect::() .with_timezone_opt(tz_opt.clone()); - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch let result = DateTruncFunc::new() - .invoke(&[ - ColumnarValue::Scalar(ScalarValue::from("day")), - ColumnarValue::Array(Arc::new(input)), - ]) + .invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::from("day")), + ColumnarValue::Array(Arc::new(input)), + ], + 0, + ) .unwrap(); if let ColumnarValue::Array(result) = result { assert_eq!( @@ -883,12 +885,14 @@ mod tests { .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) .collect::() .with_timezone_opt(tz_opt.clone()); - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch let result = DateTruncFunc::new() - .invoke(&[ - ColumnarValue::Scalar(ScalarValue::from("hour")), - ColumnarValue::Array(Arc::new(input)), - ]) + .invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::from("hour")), + ColumnarValue::Array(Arc::new(input)), + ], + 0, + ) .unwrap(); if let ColumnarValue::Array(result) = result { assert_eq!( diff --git a/datafusion/functions/src/datetime/make_date.rs b/datafusion/functions/src/datetime/make_date.rs index 6b246cb088a2b..4960b8c25254e 100644 --- a/datafusion/functions/src/datetime/make_date.rs +++ b/datafusion/functions/src/datetime/make_date.rs @@ -234,13 +234,15 @@ mod tests { #[test] fn test_make_date() { - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch let res = MakeDateFunc::new() - .invoke(&[ - ColumnarValue::Scalar(ScalarValue::Int32(Some(2024))), - ColumnarValue::Scalar(ScalarValue::Int64(Some(1))), - ColumnarValue::Scalar(ScalarValue::UInt32(Some(14))), - ]) + .invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::Int32(Some(2024))), + ColumnarValue::Scalar(ScalarValue::Int64(Some(1))), + ColumnarValue::Scalar(ScalarValue::UInt32(Some(14))), + ], + 0, + ) .expect("that make_date parsed values without error"); if let ColumnarValue::Scalar(ScalarValue::Date32(date)) = res { @@ -249,13 +251,15 @@ mod tests { panic!("Expected a scalar value") } - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch let res = MakeDateFunc::new() - .invoke(&[ - ColumnarValue::Scalar(ScalarValue::Int64(Some(2024))), - ColumnarValue::Scalar(ScalarValue::UInt64(Some(1))), - ColumnarValue::Scalar(ScalarValue::UInt32(Some(14))), - ]) + .invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::Int64(Some(2024))), + ColumnarValue::Scalar(ScalarValue::UInt64(Some(1))), + ColumnarValue::Scalar(ScalarValue::UInt32(Some(14))), + ], + 0, + ) .expect("that make_date parsed values without error"); if let ColumnarValue::Scalar(ScalarValue::Date32(date)) = res { @@ -264,13 +268,15 @@ mod tests { panic!("Expected a scalar value") } - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch let res = MakeDateFunc::new() - .invoke(&[ - ColumnarValue::Scalar(ScalarValue::Utf8(Some("2024".to_string()))), - ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some("1".to_string()))), - ColumnarValue::Scalar(ScalarValue::Utf8(Some("14".to_string()))), - ]) + .invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::Utf8(Some("2024".to_string()))), + ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some("1".to_string()))), + ColumnarValue::Scalar(ScalarValue::Utf8(Some("14".to_string()))), + ], + 0, + ) .expect("that make_date parsed values without error"); if let ColumnarValue::Scalar(ScalarValue::Date32(date)) = res { @@ -282,13 +288,15 @@ mod tests { let years = Arc::new((2021..2025).map(Some).collect::()); let months = Arc::new((1..5).map(Some).collect::()); let days = Arc::new((11..15).map(Some).collect::()); - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch let res = MakeDateFunc::new() - .invoke(&[ - ColumnarValue::Array(years), - ColumnarValue::Array(months), - ColumnarValue::Array(days), - ]) + .invoke_batch( + &[ + ColumnarValue::Array(years), + ColumnarValue::Array(months), + ColumnarValue::Array(days), + ], + 0, + ) .expect("that make_date parsed values without error"); if let ColumnarValue::Array(array) = res { @@ -308,45 +316,50 @@ mod tests { // // invalid number of arguments - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch let res = MakeDateFunc::new() - .invoke(&[ColumnarValue::Scalar(ScalarValue::Int32(Some(1)))]); + .invoke_batch(&[ColumnarValue::Scalar(ScalarValue::Int32(Some(1)))], 0); assert_eq!( res.err().unwrap().strip_backtrace(), "Execution error: make_date function requires 3 arguments, got 1" ); // invalid type - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch - let res = MakeDateFunc::new().invoke(&[ - ColumnarValue::Scalar(ScalarValue::IntervalYearMonth(Some(1))), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ]); + let res = MakeDateFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::IntervalYearMonth(Some(1))), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ], + 0, + ); assert_eq!( res.err().unwrap().strip_backtrace(), "Arrow error: Cast error: Casting from Interval(YearMonth) to Int32 not supported" ); // overflow of month - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch - let res = MakeDateFunc::new().invoke(&[ - ColumnarValue::Scalar(ScalarValue::Int32(Some(2023))), - ColumnarValue::Scalar(ScalarValue::UInt64(Some(u64::MAX))), - ColumnarValue::Scalar(ScalarValue::Int32(Some(22))), - ]); + let res = MakeDateFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::Int32(Some(2023))), + ColumnarValue::Scalar(ScalarValue::UInt64(Some(u64::MAX))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(22))), + ], + 0, + ); assert_eq!( res.err().unwrap().strip_backtrace(), "Arrow error: Cast error: Can't cast value 18446744073709551615 to type Int32" ); // overflow of day - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch - let res = MakeDateFunc::new().invoke(&[ - ColumnarValue::Scalar(ScalarValue::Int32(Some(2023))), - ColumnarValue::Scalar(ScalarValue::Int32(Some(22))), - ColumnarValue::Scalar(ScalarValue::UInt32(Some(u32::MAX))), - ]); + let res = MakeDateFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::Int32(Some(2023))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(22))), + ColumnarValue::Scalar(ScalarValue::UInt32(Some(u32::MAX))), + ], + 0, + ); assert_eq!( res.err().unwrap().strip_backtrace(), "Arrow error: Cast error: Can't cast value 4294967295 to type Int32" diff --git a/datafusion/functions/src/datetime/to_char.rs b/datafusion/functions/src/datetime/to_char.rs index ef5d6a4f69907..c31575afa8570 100644 --- a/datafusion/functions/src/datetime/to_char.rs +++ b/datafusion/functions/src/datetime/to_char.rs @@ -384,9 +384,11 @@ mod tests { ]; for (value, format, expected) in scalar_data { - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch let result = ToCharFunc::new() - .invoke(&[ColumnarValue::Scalar(value), ColumnarValue::Scalar(format)]) + .invoke_batch( + &[ColumnarValue::Scalar(value), ColumnarValue::Scalar(format)], + 0, + ) .expect("that to_char parsed values without error"); if let ColumnarValue::Scalar(ScalarValue::Utf8(date)) = result { @@ -459,12 +461,14 @@ mod tests { ]; for (value, format, expected) in scalar_array_data { - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch let result = ToCharFunc::new() - .invoke(&[ - ColumnarValue::Scalar(value), - ColumnarValue::Array(Arc::new(format) as ArrayRef), - ]) + .invoke_batch( + &[ + ColumnarValue::Scalar(value), + ColumnarValue::Array(Arc::new(format) as ArrayRef), + ], + 0, + ) .expect("that to_char parsed values without error"); if let ColumnarValue::Scalar(ScalarValue::Utf8(date)) = result { @@ -585,12 +589,14 @@ mod tests { ]; for (value, format, expected) in array_scalar_data { - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch let result = ToCharFunc::new() - .invoke(&[ - ColumnarValue::Array(value as ArrayRef), - ColumnarValue::Scalar(format), - ]) + .invoke_batch( + &[ + ColumnarValue::Array(value as ArrayRef), + ColumnarValue::Scalar(format), + ], + 0, + ) .expect("that to_char parsed values without error"); if let ColumnarValue::Array(result) = result { @@ -602,12 +608,14 @@ mod tests { } for (value, format, expected) in array_array_data { - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch let result = ToCharFunc::new() - .invoke(&[ - ColumnarValue::Array(value), - ColumnarValue::Array(Arc::new(format) as ArrayRef), - ]) + .invoke_batch( + &[ + ColumnarValue::Array(value), + ColumnarValue::Array(Arc::new(format) as ArrayRef), + ], + 0, + ) .expect("that to_char parsed values without error"); if let ColumnarValue::Array(result) = result { @@ -623,20 +631,21 @@ mod tests { // // invalid number of arguments - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch let result = ToCharFunc::new() - .invoke(&[ColumnarValue::Scalar(ScalarValue::Int32(Some(1)))]); + .invoke_batch(&[ColumnarValue::Scalar(ScalarValue::Int32(Some(1)))], 0); assert_eq!( result.err().unwrap().strip_backtrace(), "Execution error: to_char function requires 2 arguments, got 1" ); // invalid type - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch - let result = ToCharFunc::new().invoke(&[ - ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), - ]); + let result = ToCharFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), + ], + 0, + ); assert_eq!( result.err().unwrap().strip_backtrace(), "Execution error: Format for `to_char` must be non-null Utf8, received Timestamp(Nanosecond, None)" diff --git a/datafusion/functions/src/datetime/to_date.rs b/datafusion/functions/src/datetime/to_date.rs index 8f72100416e84..a654735a3d81b 100644 --- a/datafusion/functions/src/datetime/to_date.rs +++ b/datafusion/functions/src/datetime/to_date.rs @@ -213,8 +213,8 @@ mod tests { } fn test_scalar(sv: ScalarValue, tc: &TestCase) { - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch - let to_date_result = ToDateFunc::new().invoke(&[ColumnarValue::Scalar(sv)]); + let to_date_result = + ToDateFunc::new().invoke_batch(&[ColumnarValue::Scalar(sv)], 0); match to_date_result { Ok(ColumnarValue::Scalar(ScalarValue::Date32(date_val))) => { @@ -234,9 +234,8 @@ mod tests { A: From> + Array + 'static, { let date_array = A::from(vec![tc.date_str]); - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch - let to_date_result = - ToDateFunc::new().invoke(&[ColumnarValue::Array(Arc::new(date_array))]); + let to_date_result = ToDateFunc::new() + .invoke_batch(&[ColumnarValue::Array(Arc::new(date_array))], 0); match to_date_result { Ok(ColumnarValue::Array(a)) => { @@ -325,11 +324,13 @@ mod tests { fn test_scalar(sv: ScalarValue, tc: &TestCase) { let format_scalar = ScalarValue::Utf8(Some(tc.format_str.to_string())); - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch - let to_date_result = ToDateFunc::new().invoke(&[ - ColumnarValue::Scalar(sv), - ColumnarValue::Scalar(format_scalar), - ]); + let to_date_result = ToDateFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(sv), + ColumnarValue::Scalar(format_scalar), + ], + 0, + ); match to_date_result { Ok(ColumnarValue::Scalar(ScalarValue::Date32(date_val))) => { @@ -350,11 +351,13 @@ mod tests { let date_array = A::from(vec![tc.formatted_date]); let format_array = A::from(vec![tc.format_str]); - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch - let to_date_result = ToDateFunc::new().invoke(&[ - ColumnarValue::Array(Arc::new(date_array)), - ColumnarValue::Array(Arc::new(format_array)), - ]); + let to_date_result = ToDateFunc::new().invoke_batch( + &[ + ColumnarValue::Array(Arc::new(date_array)), + ColumnarValue::Array(Arc::new(format_array)), + ], + 0, + ); match to_date_result { Ok(ColumnarValue::Array(a)) => { @@ -386,12 +389,14 @@ mod tests { let format1_scalar = ScalarValue::Utf8(Some("%Y-%m-%d".into())); let format2_scalar = ScalarValue::Utf8(Some("%Y/%m/%d".into())); - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch - let to_date_result = ToDateFunc::new().invoke(&[ - ColumnarValue::Scalar(formatted_date_scalar), - ColumnarValue::Scalar(format1_scalar), - ColumnarValue::Scalar(format2_scalar), - ]); + let to_date_result = ToDateFunc::new().invoke_batch( + &[ + ColumnarValue::Scalar(formatted_date_scalar), + ColumnarValue::Scalar(format1_scalar), + ColumnarValue::Scalar(format2_scalar), + ], + 0, + ); match to_date_result { Ok(ColumnarValue::Scalar(ScalarValue::Date32(date_val))) => { @@ -415,9 +420,8 @@ mod tests { for date_str in test_cases { let formatted_date_scalar = ScalarValue::Utf8(Some(date_str.into())); - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch - let to_date_result = - ToDateFunc::new().invoke(&[ColumnarValue::Scalar(formatted_date_scalar)]); + let to_date_result = ToDateFunc::new() + .invoke_batch(&[ColumnarValue::Scalar(formatted_date_scalar)], 0); match to_date_result { Ok(ColumnarValue::Scalar(ScalarValue::Date32(date_val))) => { @@ -434,9 +438,8 @@ mod tests { let date_str = "20241231"; let date_scalar = ScalarValue::Utf8(Some(date_str.into())); - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch let to_date_result = - ToDateFunc::new().invoke(&[ColumnarValue::Scalar(date_scalar)]); + ToDateFunc::new().invoke_batch(&[ColumnarValue::Scalar(date_scalar)], 0); match to_date_result { Ok(ColumnarValue::Scalar(ScalarValue::Date32(date_val))) => { @@ -456,9 +459,8 @@ mod tests { let date_str = "202412311"; let date_scalar = ScalarValue::Utf8(Some(date_str.into())); - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch let to_date_result = - ToDateFunc::new().invoke(&[ColumnarValue::Scalar(date_scalar)]); + ToDateFunc::new().invoke_batch(&[ColumnarValue::Scalar(date_scalar)], 0); if let Ok(ColumnarValue::Scalar(ScalarValue::Date32(_))) = to_date_result { panic!( diff --git a/datafusion/functions/src/string/concat.rs b/datafusion/functions/src/string/concat.rs index e429a938b27da..1c53c448a889f 100644 --- a/datafusion/functions/src/string/concat.rs +++ b/datafusion/functions/src/string/concat.rs @@ -408,8 +408,7 @@ mod tests { ]))); let args = &[c0, c1, c2]; - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch - let result = ConcatFunc::new().invoke(args)?; + let result = ConcatFunc::new().invoke_batch(args, 0)?; let expected = Arc::new(StringArray::from(vec!["foo,x", "bar,", "baz,z"])) as ArrayRef; match &result { diff --git a/datafusion/functions/src/string/concat_ws.rs b/datafusion/functions/src/string/concat_ws.rs index 611c48a9634a3..f770f5ce41381 100644 --- a/datafusion/functions/src/string/concat_ws.rs +++ b/datafusion/functions/src/string/concat_ws.rs @@ -467,8 +467,7 @@ mod tests { ]))); let args = &[c0, c1, c2]; - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch - let result = ConcatWsFunc::new().invoke(args)?; + let result = ConcatWsFunc::new().invoke_batch(args, 0)?; let expected = Arc::new(StringArray::from(vec!["foo,x", "bar", "baz,z"])) as ArrayRef; match &result { @@ -493,8 +492,7 @@ mod tests { ]))); let args = &[c0, c1, c2]; - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch - let result = ConcatWsFunc::new().invoke(args)?; + let result = ConcatWsFunc::new().invoke_batch(args, 0)?; let expected = Arc::new(StringArray::from(vec![Some("foo,x"), None, Some("baz+z")])) as ArrayRef; diff --git a/datafusion/functions/src/string/contains.rs b/datafusion/functions/src/string/contains.rs index 0c665a139152a..6700ddec3c966 100644 --- a/datafusion/functions/src/string/contains.rs +++ b/datafusion/functions/src/string/contains.rs @@ -145,8 +145,7 @@ mod test { Some("yyy?()"), ]))); let scalar = ColumnarValue::Scalar(ScalarValue::Utf8(Some("x?(".to_string()))); - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch - let actual = udf.invoke(&[array, scalar]).unwrap(); + let actual = udf.invoke_batch(&[array, scalar], 0).unwrap(); let expect = ColumnarValue::Array(Arc::new(BooleanArray::from(vec![ Some(true), Some(false), diff --git a/datafusion/functions/src/string/lower.rs b/datafusion/functions/src/string/lower.rs index 02770e5e2203f..2513841e043af 100644 --- a/datafusion/functions/src/string/lower.rs +++ b/datafusion/functions/src/string/lower.rs @@ -105,8 +105,7 @@ mod tests { fn to_lower(input: ArrayRef, expected: ArrayRef) -> Result<()> { let func = LowerFunc::new(); let args = vec![ColumnarValue::Array(input)]; - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch - let result = match func.invoke(&args)? { + let result = match func.invoke_batch(&args, 0)? { ColumnarValue::Array(result) => result, _ => unreachable!("lower"), }; diff --git a/datafusion/functions/src/string/upper.rs b/datafusion/functions/src/string/upper.rs index 1293e51fa9947..907d9a4b51351 100644 --- a/datafusion/functions/src/string/upper.rs +++ b/datafusion/functions/src/string/upper.rs @@ -105,8 +105,7 @@ mod tests { fn to_upper(input: ArrayRef, expected: ArrayRef) -> Result<()> { let func = UpperFunc::new(); let args = vec![ColumnarValue::Array(input)]; - #[allow(deprecated)] // TODO migrate UDF invoke to invoke_batch - let result = match func.invoke(&args)? { + let result = match func.invoke_batch(&args, 0)? { ColumnarValue::Array(result) => result, _ => unreachable!("upper"), }; From 9e72e1b4e06ba3b44e939e9d35fadeff4795ff67 Mon Sep 17 00:00:00 2001 From: irenjj Date: Mon, 11 Nov 2024 21:40:58 +0800 Subject: [PATCH 2/2] fix --- datafusion/functions/src/datetime/date_bin.rs | 36 ++++++++++--------- .../functions/src/datetime/date_trunc.rs | 8 +++-- .../functions/src/datetime/make_date.rs | 17 ++++----- datafusion/functions/src/datetime/to_char.rs | 15 ++++---- datafusion/functions/src/datetime/to_date.rs | 18 +++++----- datafusion/functions/src/string/concat.rs | 2 +- datafusion/functions/src/string/concat_ws.rs | 4 +-- datafusion/functions/src/string/contains.rs | 2 +- datafusion/functions/src/string/lower.rs | 3 +- datafusion/functions/src/string/upper.rs | 3 +- 10 files changed, 61 insertions(+), 47 deletions(-) diff --git a/datafusion/functions/src/datetime/date_bin.rs b/datafusion/functions/src/datetime/date_bin.rs index 4385b1fedf0b1..5538a85b54f31 100644 --- a/datafusion/functions/src/datetime/date_bin.rs +++ b/datafusion/functions/src/datetime/date_bin.rs @@ -480,7 +480,7 @@ mod tests { use crate::datetime::date_bin::{date_bin_nanos_interval, DateBinFunc}; use arrow::array::types::TimestampNanosecondType; - use arrow::array::{IntervalDayTimeArray, TimestampNanosecondArray}; + use arrow::array::{Array, IntervalDayTimeArray, TimestampNanosecondArray}; use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos; use arrow::datatypes::{DataType, TimeUnit}; @@ -503,11 +503,12 @@ mod tests { ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ], - 0, + 1, ); assert!(res.is_ok()); let timestamps = Arc::new((1..6).map(Some).collect::()); + let batch_size = timestamps.len(); let res = DateBinFunc::new().invoke_batch( &[ ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some( @@ -519,7 +520,7 @@ mod tests { ColumnarValue::Array(timestamps), ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ], - 0, + batch_size, ); assert!(res.is_ok()); @@ -533,7 +534,7 @@ mod tests { ))), ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ], - 0, + 1, ); assert!(res.is_ok()); @@ -550,7 +551,7 @@ mod tests { ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ], - 0, + 1, ); assert!(res.is_ok()); @@ -566,7 +567,7 @@ mod tests { milliseconds: 1, }, )))], - 0, + 1, ); assert_eq!( res.err().unwrap().strip_backtrace(), @@ -580,7 +581,7 @@ mod tests { ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ], - 0, + 1, ); assert_eq!( res.err().unwrap().strip_backtrace(), @@ -599,7 +600,7 @@ mod tests { ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ], - 0, + 1, ); assert_eq!( res.err().unwrap().strip_backtrace(), @@ -615,7 +616,7 @@ mod tests { ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ], - 0, + 1, ); assert_eq!( res.err().unwrap().strip_backtrace(), @@ -629,7 +630,7 @@ mod tests { ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ], - 0, + 1, ); assert_eq!( res.err().unwrap().strip_backtrace(), @@ -643,7 +644,7 @@ mod tests { ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ], - 0, + 1, ); assert_eq!( res.err().unwrap().strip_backtrace(), @@ -662,7 +663,7 @@ mod tests { ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(Some(1), None)), ], - 0, + 1, ); assert_eq!( res.err().unwrap().strip_backtrace(), @@ -680,7 +681,7 @@ mod tests { ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(Some(1), None)), ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ], - 0, + 1, ); assert!(res.is_ok()); @@ -695,13 +696,14 @@ mod tests { }) .collect::(), ); + let batch_size = intervals.len(); let res = DateBinFunc::new().invoke_batch( &[ ColumnarValue::Array(intervals), ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ], - 0, + batch_size, ); assert_eq!( res.err().unwrap().strip_backtrace(), @@ -710,6 +712,7 @@ mod tests { // unsupported array type for origin let timestamps = Arc::new((1..6).map(Some).collect::()); + let batch_size = timestamps.len(); let res = DateBinFunc::new().invoke_batch( &[ ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some( @@ -721,7 +724,7 @@ mod tests { ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ColumnarValue::Array(timestamps), ], - 0, + batch_size, ); assert_eq!( res.err().unwrap().strip_backtrace(), @@ -837,6 +840,7 @@ mod tests { .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) .collect::() .with_timezone_opt(tz_opt.clone()); + let batch_size = input.len(); let result = DateBinFunc::new() .invoke_batch( &[ @@ -847,7 +851,7 @@ mod tests { tz_opt.clone(), )), ], - 0, + batch_size, ) .unwrap(); if let ColumnarValue::Array(result) = result { diff --git a/datafusion/functions/src/datetime/date_trunc.rs b/datafusion/functions/src/datetime/date_trunc.rs index a0d472d81d8f2..5ec308ef9c818 100644 --- a/datafusion/functions/src/datetime/date_trunc.rs +++ b/datafusion/functions/src/datetime/date_trunc.rs @@ -484,7 +484,7 @@ mod tests { use arrow::array::cast::as_primitive_array; use arrow::array::types::TimestampNanosecondType; - use arrow::array::TimestampNanosecondArray; + use arrow::array::{Array, TimestampNanosecondArray}; use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos; use arrow::datatypes::{DataType, TimeUnit}; use datafusion_common::ScalarValue; @@ -724,13 +724,14 @@ mod tests { .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) .collect::() .with_timezone_opt(tz_opt.clone()); + let batch_size = input.len(); let result = DateTruncFunc::new() .invoke_batch( &[ ColumnarValue::Scalar(ScalarValue::from("day")), ColumnarValue::Array(Arc::new(input)), ], - 0, + batch_size, ) .unwrap(); if let ColumnarValue::Array(result) = result { @@ -885,13 +886,14 @@ mod tests { .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) .collect::() .with_timezone_opt(tz_opt.clone()); + let batch_size = input.len(); let result = DateTruncFunc::new() .invoke_batch( &[ ColumnarValue::Scalar(ScalarValue::from("hour")), ColumnarValue::Array(Arc::new(input)), ], - 0, + batch_size, ) .unwrap(); if let ColumnarValue::Array(result) = result { diff --git a/datafusion/functions/src/datetime/make_date.rs b/datafusion/functions/src/datetime/make_date.rs index 4960b8c25254e..a13511f33398c 100644 --- a/datafusion/functions/src/datetime/make_date.rs +++ b/datafusion/functions/src/datetime/make_date.rs @@ -241,7 +241,7 @@ mod tests { ColumnarValue::Scalar(ScalarValue::Int64(Some(1))), ColumnarValue::Scalar(ScalarValue::UInt32(Some(14))), ], - 0, + 1, ) .expect("that make_date parsed values without error"); @@ -258,7 +258,7 @@ mod tests { ColumnarValue::Scalar(ScalarValue::UInt64(Some(1))), ColumnarValue::Scalar(ScalarValue::UInt32(Some(14))), ], - 0, + 1, ) .expect("that make_date parsed values without error"); @@ -275,7 +275,7 @@ mod tests { ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some("1".to_string()))), ColumnarValue::Scalar(ScalarValue::Utf8(Some("14".to_string()))), ], - 0, + 1, ) .expect("that make_date parsed values without error"); @@ -288,6 +288,7 @@ mod tests { let years = Arc::new((2021..2025).map(Some).collect::()); let months = Arc::new((1..5).map(Some).collect::()); let days = Arc::new((11..15).map(Some).collect::()); + let batch_size = years.len(); let res = MakeDateFunc::new() .invoke_batch( &[ @@ -295,7 +296,7 @@ mod tests { ColumnarValue::Array(months), ColumnarValue::Array(days), ], - 0, + batch_size, ) .expect("that make_date parsed values without error"); @@ -317,7 +318,7 @@ mod tests { // invalid number of arguments let res = MakeDateFunc::new() - .invoke_batch(&[ColumnarValue::Scalar(ScalarValue::Int32(Some(1)))], 0); + .invoke_batch(&[ColumnarValue::Scalar(ScalarValue::Int32(Some(1)))], 1); assert_eq!( res.err().unwrap().strip_backtrace(), "Execution error: make_date function requires 3 arguments, got 1" @@ -330,7 +331,7 @@ mod tests { ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ], - 0, + 1, ); assert_eq!( res.err().unwrap().strip_backtrace(), @@ -344,7 +345,7 @@ mod tests { ColumnarValue::Scalar(ScalarValue::UInt64(Some(u64::MAX))), ColumnarValue::Scalar(ScalarValue::Int32(Some(22))), ], - 0, + 1, ); assert_eq!( res.err().unwrap().strip_backtrace(), @@ -358,7 +359,7 @@ mod tests { ColumnarValue::Scalar(ScalarValue::Int32(Some(22))), ColumnarValue::Scalar(ScalarValue::UInt32(Some(u32::MAX))), ], - 0, + 1, ); assert_eq!( res.err().unwrap().strip_backtrace(), diff --git a/datafusion/functions/src/datetime/to_char.rs b/datafusion/functions/src/datetime/to_char.rs index c31575afa8570..dd4ae7b8464ed 100644 --- a/datafusion/functions/src/datetime/to_char.rs +++ b/datafusion/functions/src/datetime/to_char.rs @@ -387,7 +387,7 @@ mod tests { let result = ToCharFunc::new() .invoke_batch( &[ColumnarValue::Scalar(value), ColumnarValue::Scalar(format)], - 0, + 1, ) .expect("that to_char parsed values without error"); @@ -461,13 +461,14 @@ mod tests { ]; for (value, format, expected) in scalar_array_data { + let batch_size = format.len(); let result = ToCharFunc::new() .invoke_batch( &[ ColumnarValue::Scalar(value), ColumnarValue::Array(Arc::new(format) as ArrayRef), ], - 0, + batch_size, ) .expect("that to_char parsed values without error"); @@ -589,13 +590,14 @@ mod tests { ]; for (value, format, expected) in array_scalar_data { + let batch_size = value.len(); let result = ToCharFunc::new() .invoke_batch( &[ ColumnarValue::Array(value as ArrayRef), ColumnarValue::Scalar(format), ], - 0, + batch_size, ) .expect("that to_char parsed values without error"); @@ -608,13 +610,14 @@ mod tests { } for (value, format, expected) in array_array_data { + let batch_size = value.len(); let result = ToCharFunc::new() .invoke_batch( &[ ColumnarValue::Array(value), ColumnarValue::Array(Arc::new(format) as ArrayRef), ], - 0, + batch_size, ) .expect("that to_char parsed values without error"); @@ -632,7 +635,7 @@ mod tests { // invalid number of arguments let result = ToCharFunc::new() - .invoke_batch(&[ColumnarValue::Scalar(ScalarValue::Int32(Some(1)))], 0); + .invoke_batch(&[ColumnarValue::Scalar(ScalarValue::Int32(Some(1)))], 1); assert_eq!( result.err().unwrap().strip_backtrace(), "Execution error: to_char function requires 2 arguments, got 1" @@ -644,7 +647,7 @@ mod tests { ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)), ], - 0, + 1, ); assert_eq!( result.err().unwrap().strip_backtrace(), diff --git a/datafusion/functions/src/datetime/to_date.rs b/datafusion/functions/src/datetime/to_date.rs index a654735a3d81b..ff322ce319605 100644 --- a/datafusion/functions/src/datetime/to_date.rs +++ b/datafusion/functions/src/datetime/to_date.rs @@ -214,7 +214,7 @@ mod tests { fn test_scalar(sv: ScalarValue, tc: &TestCase) { let to_date_result = - ToDateFunc::new().invoke_batch(&[ColumnarValue::Scalar(sv)], 0); + ToDateFunc::new().invoke_batch(&[ColumnarValue::Scalar(sv)], 1); match to_date_result { Ok(ColumnarValue::Scalar(ScalarValue::Date32(date_val))) => { @@ -234,8 +234,9 @@ mod tests { A: From> + Array + 'static, { let date_array = A::from(vec![tc.date_str]); + let batch_size = date_array.len(); let to_date_result = ToDateFunc::new() - .invoke_batch(&[ColumnarValue::Array(Arc::new(date_array))], 0); + .invoke_batch(&[ColumnarValue::Array(Arc::new(date_array))], batch_size); match to_date_result { Ok(ColumnarValue::Array(a)) => { @@ -329,7 +330,7 @@ mod tests { ColumnarValue::Scalar(sv), ColumnarValue::Scalar(format_scalar), ], - 0, + 1, ); match to_date_result { @@ -351,12 +352,13 @@ mod tests { let date_array = A::from(vec![tc.formatted_date]); let format_array = A::from(vec![tc.format_str]); + let batch_size = date_array.len(); let to_date_result = ToDateFunc::new().invoke_batch( &[ ColumnarValue::Array(Arc::new(date_array)), ColumnarValue::Array(Arc::new(format_array)), ], - 0, + batch_size, ); match to_date_result { @@ -395,7 +397,7 @@ mod tests { ColumnarValue::Scalar(format1_scalar), ColumnarValue::Scalar(format2_scalar), ], - 0, + 1, ); match to_date_result { @@ -421,7 +423,7 @@ mod tests { let formatted_date_scalar = ScalarValue::Utf8(Some(date_str.into())); let to_date_result = ToDateFunc::new() - .invoke_batch(&[ColumnarValue::Scalar(formatted_date_scalar)], 0); + .invoke_batch(&[ColumnarValue::Scalar(formatted_date_scalar)], 1); match to_date_result { Ok(ColumnarValue::Scalar(ScalarValue::Date32(date_val))) => { @@ -439,7 +441,7 @@ mod tests { let date_scalar = ScalarValue::Utf8(Some(date_str.into())); let to_date_result = - ToDateFunc::new().invoke_batch(&[ColumnarValue::Scalar(date_scalar)], 0); + ToDateFunc::new().invoke_batch(&[ColumnarValue::Scalar(date_scalar)], 1); match to_date_result { Ok(ColumnarValue::Scalar(ScalarValue::Date32(date_val))) => { @@ -460,7 +462,7 @@ mod tests { let date_scalar = ScalarValue::Utf8(Some(date_str.into())); let to_date_result = - ToDateFunc::new().invoke_batch(&[ColumnarValue::Scalar(date_scalar)], 0); + ToDateFunc::new().invoke_batch(&[ColumnarValue::Scalar(date_scalar)], 1); if let Ok(ColumnarValue::Scalar(ScalarValue::Date32(_))) = to_date_result { panic!( diff --git a/datafusion/functions/src/string/concat.rs b/datafusion/functions/src/string/concat.rs index 1c53c448a889f..f1e60004ddd00 100644 --- a/datafusion/functions/src/string/concat.rs +++ b/datafusion/functions/src/string/concat.rs @@ -408,7 +408,7 @@ mod tests { ]))); let args = &[c0, c1, c2]; - let result = ConcatFunc::new().invoke_batch(args, 0)?; + let result = ConcatFunc::new().invoke_batch(args, 3)?; let expected = Arc::new(StringArray::from(vec!["foo,x", "bar,", "baz,z"])) as ArrayRef; match &result { diff --git a/datafusion/functions/src/string/concat_ws.rs b/datafusion/functions/src/string/concat_ws.rs index f770f5ce41381..98a75f121c35f 100644 --- a/datafusion/functions/src/string/concat_ws.rs +++ b/datafusion/functions/src/string/concat_ws.rs @@ -467,7 +467,7 @@ mod tests { ]))); let args = &[c0, c1, c2]; - let result = ConcatWsFunc::new().invoke_batch(args, 0)?; + let result = ConcatWsFunc::new().invoke_batch(args, 3)?; let expected = Arc::new(StringArray::from(vec!["foo,x", "bar", "baz,z"])) as ArrayRef; match &result { @@ -492,7 +492,7 @@ mod tests { ]))); let args = &[c0, c1, c2]; - let result = ConcatWsFunc::new().invoke_batch(args, 0)?; + let result = ConcatWsFunc::new().invoke_batch(args, 3)?; let expected = Arc::new(StringArray::from(vec![Some("foo,x"), None, Some("baz+z")])) as ArrayRef; diff --git a/datafusion/functions/src/string/contains.rs b/datafusion/functions/src/string/contains.rs index 6700ddec3c966..3acd2464524da 100644 --- a/datafusion/functions/src/string/contains.rs +++ b/datafusion/functions/src/string/contains.rs @@ -145,7 +145,7 @@ mod test { Some("yyy?()"), ]))); let scalar = ColumnarValue::Scalar(ScalarValue::Utf8(Some("x?(".to_string()))); - let actual = udf.invoke_batch(&[array, scalar], 0).unwrap(); + let actual = udf.invoke_batch(&[array, scalar], 2).unwrap(); let expect = ColumnarValue::Array(Arc::new(BooleanArray::from(vec![ Some(true), Some(false), diff --git a/datafusion/functions/src/string/lower.rs b/datafusion/functions/src/string/lower.rs index 2513841e043af..78887fde0a8eb 100644 --- a/datafusion/functions/src/string/lower.rs +++ b/datafusion/functions/src/string/lower.rs @@ -104,8 +104,9 @@ mod tests { fn to_lower(input: ArrayRef, expected: ArrayRef) -> Result<()> { let func = LowerFunc::new(); + let batch_size = input.len(); let args = vec![ColumnarValue::Array(input)]; - let result = match func.invoke_batch(&args, 0)? { + let result = match func.invoke_batch(&args, batch_size)? { ColumnarValue::Array(result) => result, _ => unreachable!("lower"), }; diff --git a/datafusion/functions/src/string/upper.rs b/datafusion/functions/src/string/upper.rs index 907d9a4b51351..5039d094f2d6d 100644 --- a/datafusion/functions/src/string/upper.rs +++ b/datafusion/functions/src/string/upper.rs @@ -104,8 +104,9 @@ mod tests { fn to_upper(input: ArrayRef, expected: ArrayRef) -> Result<()> { let func = UpperFunc::new(); + let batch_size = input.len(); let args = vec![ColumnarValue::Array(input)]; - let result = match func.invoke_batch(&args, 0)? { + let result = match func.invoke_batch(&args, batch_size)? { ColumnarValue::Array(result) => result, _ => unreachable!("upper"), };