viirya commented on code in PR #7181:
URL: https://github.com/apache/datafusion-comet/pull/7181#discussion_r4391838906
##########
native/core/src/execution/python_udf.rs:
##########
@@ -297,14 +468,18 @@ impl ArrowPythonUdf {
kwargs.set_item(name, arg).map_err(python_error)?;
}
}
- let result = self
- .callable
- .bind(py)
- .call(
- PyTuple::new(py, positional).map_err(python_error)?,
- Some(&kwargs),
- )
- .map_err(|error| python_traceback_error(py, error))?;
+ let positional = PyTuple::new(py,
positional).map_err(python_error)?;
+ let call = self
+ .task
+ .as_deref()
+ .map(|task| RunningCall::enter(py, task))
+ .transpose()?;
+ let result = self.callable.bind(py).call(positional,
Some(&kwargs));
+ drop(call);
Review Comment:
Good catch, thanks. An extension type's `__arrow_ext_class__` can return an
`ExtensionArray` subclass that overrides `cast`, and the same goes for
`__len__`, the type comparison, and the type's `__str__` in the mismatch error.
All of them ran after the call was unrecorded and after the kill check. Fixed
in 69697372f. Everything from the callable to `_export_to_c` now runs in
`call_and_export` while the call is still recorded, so a kill interrupts any
user code in that span. The call is unrecorded only after the export, on
success and on error alike, and the kill check follows it.
`interrupts_a_user_defined_result_cast` returns an extension array whose `cast`
loops, interrupts it, and checks that the call fails with `TaskKilled`. It
fails when the call is unrecorded right after the callable.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
