Skip to content

Commit eb15075

Browse files
committed
Added support for ScalarUDFImpl::invoke_with_return_type where the invoke is passed the return type created for the udf instance
1 parent b0b6e44 commit eb15075

File tree

2 files changed

+26
-5
lines changed

2 files changed

+26
-5
lines changed

datafusion/expr/src/udf.rs

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -203,9 +203,6 @@ impl ScalarUDF {
203203
self.inner.simplify(args, info)
204204
}
205205

206-
/// Invoke the function on `args`, returning the appropriate result.
207-
///
208-
/// See [`ScalarUDFImpl::invoke`] for more details.
209206
#[deprecated(since = "42.1.0", note = "Use `invoke_batch` instead")]
210207
pub fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
211208
#[allow(deprecated)]
@@ -227,6 +224,19 @@ impl ScalarUDF {
227224
self.inner.invoke_batch(args, number_rows)
228225
}
229226

227+
/// Invoke the function on `args`, returning the appropriate result.
228+
///
229+
/// See [`ScalarUDFImpl::invoke_batch`] for more details.
230+
pub fn invoke_batch_with_return_type(
231+
&self,
232+
args: &[ColumnarValue],
233+
number_rows: usize,
234+
return_type: &DataType,
235+
) -> Result<ColumnarValue> {
236+
self.inner
237+
.invoke_batch_with_return_type(args, number_rows, return_type)
238+
}
239+
230240
/// Invoke the function without `args` but number of rows, returning the appropriate result.
231241
///
232242
/// See [`ScalarUDFImpl::invoke_no_args`] for more details.
@@ -356,7 +366,7 @@ where
356366
/// }
357367
/// }
358368
/// }
359-
///
369+
///
360370
/// static DOCUMENTATION: OnceLock<Documentation> = OnceLock::new();
361371
///
362372
/// fn get_doc() -> &'static Documentation {
@@ -537,6 +547,17 @@ pub trait ScalarUDFImpl: Debug + Send + Sync {
537547
}
538548
}
539549

550+
/// This function will be called with the evaluated children as in `invoke` however, the value
551+
/// returned previously from `ScalarUDFImpl::return_type` for this expr will be passed in.
552+
fn invoke_batch_with_return_type(
553+
&self,
554+
args: &[ColumnarValue],
555+
number_rows: usize,
556+
_return_type: &DataType,
557+
) -> Result<ColumnarValue> {
558+
self.invoke_batch(args, number_rows)
559+
}
560+
540561
/// Invoke the function without `args`, instead the number of rows are provided,
541562
/// returning the appropriate result.
542563
#[deprecated(since = "42.1.0", note = "Use `invoke_batch` instead")]

datafusion/physical-expr/src/scalar_function.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,7 @@ impl PhysicalExpr for ScalarFunctionExpr {
141141
.collect::<Result<Vec<_>>>()?;
142142

143143
// evaluate the function
144-
let output = self.fun.invoke_batch(&inputs, batch.num_rows())?;
144+
let output = self.fun.invoke_batch_with_return_type(&inputs, batch.num_rows(), self.return_type())?;
145145

146146
if let ColumnarValue::Array(array) = &output {
147147
if array.len() != batch.num_rows() {

0 commit comments

Comments
 (0)