nevi-me commented on issue #552:
URL: https://github.com/apache/sedona-db/issues/552#issuecomment-3804033066
I think I understand what you're suggesting, but let me test it out below.
I've created a wrapper that would allow me to do type conversions, I'm still
reading through the code to understand how to do that. Below I have a hopefully
trivial example that uses `st_point`, which I'd expect to work with Float64
types.
```rust
use datafusion_expr::ScalarUDFImpl;
use sedona_expr::scalar_udf::SedonaScalarUDF;
use sedona_schema::datatypes::SedonaType;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct SedonaUDFWrapper {
inner: SedonaScalarUDF,
}
impl SedonaUDFWrapper {
pub fn new(inner: SedonaScalarUDF) -> Self {
Self { inner }
}
}
impl ScalarUDFImpl for SedonaUDFWrapper {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
self.inner.name()
}
fn signature(&self) -> &datafusion_expr::Signature {
self.inner.signature()
}
fn return_type(
&self,
arg_types: &[arrow::datatypes::DataType],
) -> datafusion::error::Result<arrow::datatypes::DataType> {
self.inner.return_type(arg_types)
}
fn invoke_with_args(
&self,
args: datafusion_expr::ScalarFunctionArgs,
) -> datafusion::error::Result<datafusion_expr::ColumnarValue> {
let mut args = args;
// TODO: convert args to Sedona types
self.inner.invoke_with_args(args)
}
}
#[cfg(test)]
mod tests {
use arrow_array::BinaryArray;
use datafusion::prelude::SessionContext;
use datafusion_expr::ScalarUDF;
use super::*;
#[ tokio::test]
async fn test_exec_sedona_udf_wrapper() {
let ctx = SessionContext::new();
let st_point =
sedona_functions::register::default_function_set().scalar_udf("st_point").unwrap().clone();
let st_centroid =
sedona_functions::register::default_function_set().scalar_udf("st_centroid").unwrap().clone();
ctx.register_udf(ScalarUDF::new_from_impl(SedonaUDFWrapper::new(st_point)));
ctx.register_udf(ScalarUDF::new_from_impl(SedonaUDFWrapper::new(st_centroid)));
let df = ctx.sql("select st_point(-64.36, 45.09) as
point").await.unwrap();
let result = df.collect().await.unwrap();
let batch = result.get(0).unwrap().clone();
arrow::util::pretty::print_batches(&[batch]).unwrap();
// let value =
batch.column(0).as_any().downcast_ref::<BinaryArray>().unwrap();
// assert_eq!(value.value(0), true);
}
}
```
Running the test however fails with
```
called `Result::unwrap()` on an `Err` value: Plan("Execution error: Function
'st_point' user-defined coercion failed with \"This feature is not implemented:
Function st_point does not implement coerce_types\" No function matches the
given name and argument types 'st_point(Float64, Float64)'. You might need to
add explicit type casts.\n\tCandidate functions:\n\tst_point(UserDefined)")
```
I would expect registering `st_point` to ?just work. Are you able to see
what I'm doing wrong with registering the UDF? I've seen that you have a
`SedonaContext`, so I'm also checking if there's additional setup in there,
which I might be missing in my plain `SessionContext`.
--
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]
