From 88d86e7fe26d6c3803c80ecc1b1c4235a58e36ca Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Fri, 17 Oct 2025 21:44:25 -0500 Subject: [PATCH 01/25] building --- python/sedonadb/src/error.rs | 6 + python/sedonadb/src/import_from.rs | 31 +++- python/sedonadb/src/lib.rs | 4 +- python/sedonadb/src/udf.rs | 222 +++++++++++++++++++++++++++++ 4 files changed, 259 insertions(+), 4 deletions(-) create mode 100644 python/sedonadb/src/udf.rs diff --git a/python/sedonadb/src/error.rs b/python/sedonadb/src/error.rs index c274ed186b..a3bc309907 100644 --- a/python/sedonadb/src/error.rs +++ b/python/sedonadb/src/error.rs @@ -51,6 +51,12 @@ impl From for PySedonaError { } } +impl From for DataFusionError { + fn from(other: PySedonaError) -> Self { + DataFusionError::External(Box::new(other)) + } +} + impl From for PySedonaError { fn from(other: PyErr) -> Self { PySedonaError::Py(other) diff --git a/python/sedonadb/src/import_from.rs b/python/sedonadb/src/import_from.rs index b6c694e26b..f2ccca4c28 100644 --- a/python/sedonadb/src/import_from.rs +++ b/python/sedonadb/src/import_from.rs @@ -20,11 +20,11 @@ use std::{ }; use arrow_array::{ - ffi::FFI_ArrowSchema, + ffi::{FFI_ArrowArray, FFI_ArrowSchema}, ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream}, - RecordBatchReader, + make_array, ArrayRef, RecordBatchReader, }; -use arrow_schema::Schema; +use arrow_schema::{Field, Schema}; use datafusion::catalog::TableProvider; use datafusion_ffi::table_provider::{FFI_TableProvider, ForeignTableProvider}; use pyo3::{ @@ -88,6 +88,31 @@ pub fn import_arrow_array_stream<'py>( Ok(Box::new(stream_reader)) } +pub fn import_arrow_array(obj: &Bound) -> Result<(Field, ArrayRef), PySedonaError> { + let schema_and_array = obj.getattr("__arrow_c_schema__")?.call0()?; + let (schema_capsule, array_capsule): (Bound, Bound) = + schema_and_array.extract()?; + + let ffi_schema = unsafe { + FFI_ArrowSchema::from_raw(check_pycapsule(&schema_capsule, "arrow_schema")? as _) + }; + let ffi_array = + unsafe { FFI_ArrowArray::from_raw(check_pycapsule(&array_capsule, "arrow_array")? as _) }; + + let result_field = Field::try_from(&ffi_schema)?; + let result_array_data = unsafe { arrow_array::ffi::from_ffi(ffi_array, &ffi_schema)? }; + + Ok((result_field, make_array(result_array_data))) +} + +pub fn import_arrow_field(obj: &Bound) -> Result { + let capsule = obj.getattr("__arrow_c_schema__")?.call0()?; + let schema = + unsafe { FFI_ArrowSchema::from_raw(check_pycapsule(&capsule, "arrow_schema")? as _) }; + + Ok(Field::try_from(&schema)?) +} + pub fn import_arrow_schema(obj: &Bound) -> Result { let capsule = obj.getattr("__arrow_c_schema__")?.call0()?; let schema = diff --git a/python/sedonadb/src/lib.rs b/python/sedonadb/src/lib.rs index ca09d87528..80f4df7e62 100644 --- a/python/sedonadb/src/lib.rs +++ b/python/sedonadb/src/lib.rs @@ -14,7 +14,7 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. -use crate::error::PySedonaError; +use crate::{error::PySedonaError, udf::register_sedona_scalar_udf}; use pyo3::{ffi::Py_uintptr_t, prelude::*}; use sedona_adbc::AdbcSedonadbDriverInit; use sedona_proj::register::{configure_global_proj_engine, ProjCrsEngineBuilder}; @@ -27,6 +27,7 @@ mod import_from; mod reader; mod runtime; mod schema; +mod udf; const VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -89,6 +90,7 @@ fn _lib(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(configure_proj_shared, m)?)?; m.add_function(wrap_pyfunction!(sedona_adbc_driver_init, m)?)?; m.add_function(wrap_pyfunction!(sedona_python_version, m)?)?; + m.add_function(wrap_pyfunction!(register_sedona_scalar_udf, m)?)?; m.add_class::()?; m.add_class::()?; diff --git a/python/sedonadb/src/udf.rs b/python/sedonadb/src/udf.rs new file mode 100644 index 0000000000..ba59e4e618 --- /dev/null +++ b/python/sedonadb/src/udf.rs @@ -0,0 +1,222 @@ +use std::{ffi::CString, iter::zip, sync::Arc}; + +use arrow_array::{ + ffi::{FFI_ArrowArray, FFI_ArrowSchema}, + ArrayRef, +}; +use arrow_schema::{DataType, FieldRef}; +use datafusion_common::{not_impl_datafusion_err, Result, ScalarValue}; +use datafusion_expr::{ + ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, Volatility +}; +use datafusion_ffi::udf::FFI_ScalarUDF; +use pyo3::{ + pyclass, pyfunction, pymethods, + types::{PyAnyMethods, PyCapsule, PyDict, PyTuple}, + Bound, PyObject, Python, +}; +use sedona_schema::datatypes::SedonaType; + +use crate::{ + error::PySedonaError, + import_from::{import_arrow_array, import_arrow_field}, + schema::{PySedonaField, PySedonaType}, +}; + +#[pyfunction] +pub fn register_sedona_scalar_udf<'py>( + py: Python<'py>, + name: &str, + py_return_field: PyObject, + py_invoke_batch: PyObject, + volatility: Option<&str>, +) -> Result, PySedonaError> { + let volatility = match volatility.unwrap_or("immutable") { + "immutable" => Volatility::Immutable, + "stable" => Volatility::Stable, + "volatile" => Volatility::Volatile, + v => { + return Err(PySedonaError::SedonaPython(format!( + "Expected one of 'immutable', 'stable', or 'volatile' but got '{v}'" + ))); + } + }; + + let udf_impl = PySedonaScalarUdf { + name: name.to_string(), + signature: Signature::new(TypeSignature::UserDefined, volatility), + py_return_field, + py_invoke_batch, + }; + + let name = cr"datafusion_scalar_udf".into(); + let udf = ScalarUDF::from(udf_impl); + let ffi_udf = FFI_ScalarUDF::from(Arc::new(udf)); + Ok(PyCapsule::new(py, ffi_udf, Some(name))?) +} + +#[derive(Debug)] +struct PySedonaScalarUdf { + name: String, + signature: Signature, + py_return_field: PyObject, + py_invoke_batch: PyObject, +} + +impl ScalarUDFImpl for PySedonaScalarUdf { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Err(PySedonaError::SedonaPython("Unexpected call to return_type()".to_string()).into()) + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + Ok(eval_return_field(&self.py_return_field, args.arg_fields)?) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + Ok(arg_types.to_vec()) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + Ok(eval_invoke_batch(&self.py_invoke_batch, args)?) + } +} + +fn eval_return_field(func: &PyObject, arg_fields: &[FieldRef]) -> Result { + let return_field = Python::with_gil(|py| -> Result { + let py_arg_fields = arg_fields + .iter() + .map(|f| PySedonaField::new(f.as_ref().clone())) + .collect::>(); + + let py_args = PyTuple::new(py, py_arg_fields)?; + let py_return_field = func.call(py, py_args, None)?; + if py_return_field.is_none(py) { + return Err(PySedonaError::DF(Box::new(not_impl_datafusion_err!( + "Python Udf does not apply to arguments" + )))); + } + + let return_field = import_arrow_field(py_return_field.bind(py))?; + Ok(Arc::new(return_field)) + })?; + + Ok(return_field) +} + +fn eval_invoke_batch( + func: &PyObject, + args: ScalarFunctionArgs, +) -> Result { + let result = Python::with_gil(|py| -> Result { + let py_values = zip(&args.arg_fields, &args.args) + .map(|(f, arg)| PySedonaValue { + field: f.clone(), + value: arg.clone(), + num_rows: args.number_rows, + }) + .collect::>(); + + let py_return_field = PySedonaField::new(args.return_field.as_ref().clone()); + let py_args = PyTuple::new(py, py_values)?; + let py_kwargs = PyDict::new(py); + py_kwargs.set_item("return_field", py_return_field)?; + + let result = func.call(py, py_args, Some(&py_kwargs))?; + + let (result_field, result_array) = import_arrow_array(result.bind(py))?; + let result_sedona_type = SedonaType::from_storage_field(&result_field)?; + + let expected_result_sedona_type = SedonaType::from_storage_field(&args.return_field)?; + if expected_result_sedona_type != result_sedona_type { + return Err(PySedonaError::SedonaPython(format!( + "Expected {expected_result_sedona_type} but got {result_sedona_type}" + ))); + } + + Ok(result_array) + })?; + + if args.args.is_empty() { + return Ok(ColumnarValue::Array(result)); + } + + for arg in &args.args { + match arg { + ColumnarValue::Array(_) => return Ok(ColumnarValue::Array(result)), + ColumnarValue::Scalar(_) => {} + } + } + + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)) +} + +#[pyclass] +#[derive(Debug)] +pub struct PySedonaValue { + pub field: FieldRef, + pub value: ColumnarValue, + pub num_rows: usize, +} + +#[pymethods] +impl PySedonaValue { + #[getter] + fn r#type(&self) -> Result { + Ok(PySedonaType::new(SedonaType::from_storage_field( + &self.field, + )?)) + } + + fn is_scalar(&self) -> bool { + matches!(&self.value, ColumnarValue::Scalar(_)) + } + + fn to_array(&self) -> Result { + Ok(PySedonaValue { + field: self.field.clone(), + value: ColumnarValue::Array(self.value.to_array(self.num_rows)?), + num_rows: self.num_rows, + }) + } + + fn __arrow_c_schema__<'py>( + &self, + py: Python<'py>, + ) -> Result, PySedonaError> { + let schema_capsule_name = CString::new("arrow_schema").unwrap(); + let ffi_schema = FFI_ArrowSchema::try_from(self.field.as_ref().clone())?; + Ok(PyCapsule::new(py, ffi_schema, Some(schema_capsule_name))?) + } + + fn __arrow_c_array__<'py>( + &self, + py: Python<'py>, + ) -> Result, PySedonaError> { + let schema_capsule_name = CString::new("arrow_array").unwrap(); + let out_size = match &self.value { + ColumnarValue::Array(array) => array.len(), + ColumnarValue::Scalar(_) => 1, + }; + let array = self.value.to_array(out_size)?; + let ffi_array = FFI_ArrowArray::new(&array.to_data()); + Ok(PyCapsule::new(py, ffi_array, Some(schema_capsule_name))?) + } + + fn __repr__(&self) -> String { + format!("PySedonaValue {self:?}") + } +} From b00bce710b119da6291138e29da1f05b65d0aa67 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Fri, 17 Oct 2025 22:13:52 -0500 Subject: [PATCH 02/25] tweak --- python/sedonadb/python/sedonadb/udf.py | 37 +++++++++++++++++++ python/sedonadb/src/lib.rs | 4 +- python/sedonadb/src/udf.rs | 51 +++++++++++++++++++------- 3 files changed, 77 insertions(+), 15 deletions(-) create mode 100644 python/sedonadb/python/sedonadb/udf.py diff --git a/python/sedonadb/python/sedonadb/udf.py b/python/sedonadb/python/sedonadb/udf.py new file mode 100644 index 0000000000..0eff543bf0 --- /dev/null +++ b/python/sedonadb/python/sedonadb/udf.py @@ -0,0 +1,37 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from sedonadb._lib import sedona_scalar_udf + +class ScalarUdfImpl: + + @property + def name(self): + raise NotImplementedError() + + @property + def volatility(self): + return "immutable" + + def return_type(self, arg_types, scalar_args): + raise NotImplementedError() + + def invoke_batch(self, args, return_type, num_rows): + raise NotImplementedError() + + def __datafusion_scalar_udf__(self): + return sedona_scalar_udf(self.name, self.return_type, self.invoke_batch, self.volatility) diff --git a/python/sedonadb/src/lib.rs b/python/sedonadb/src/lib.rs index 80f4df7e62..62a0cabaf8 100644 --- a/python/sedonadb/src/lib.rs +++ b/python/sedonadb/src/lib.rs @@ -14,7 +14,7 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. -use crate::{error::PySedonaError, udf::register_sedona_scalar_udf}; +use crate::{error::PySedonaError, udf::sedona_scalar_udf}; use pyo3::{ffi::Py_uintptr_t, prelude::*}; use sedona_adbc::AdbcSedonadbDriverInit; use sedona_proj::register::{configure_global_proj_engine, ProjCrsEngineBuilder}; @@ -90,7 +90,7 @@ fn _lib(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(configure_proj_shared, m)?)?; m.add_function(wrap_pyfunction!(sedona_adbc_driver_init, m)?)?; m.add_function(wrap_pyfunction!(sedona_python_version, m)?)?; - m.add_function(wrap_pyfunction!(register_sedona_scalar_udf, m)?)?; + m.add_function(wrap_pyfunction!(sedona_scalar_udf, m)?)?; m.add_class::()?; m.add_class::()?; diff --git a/python/sedonadb/src/udf.rs b/python/sedonadb/src/udf.rs index ba59e4e618..9b36840b1e 100644 --- a/python/sedonadb/src/udf.rs +++ b/python/sedonadb/src/udf.rs @@ -1,3 +1,20 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + use std::{ffi::CString, iter::zip, sync::Arc}; use arrow_array::{ @@ -7,12 +24,13 @@ use arrow_array::{ use arrow_schema::{DataType, FieldRef}; use datafusion_common::{not_impl_datafusion_err, Result, ScalarValue}; use datafusion_expr::{ - ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, Volatility + ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, + TypeSignature, Volatility, }; use datafusion_ffi::udf::FFI_ScalarUDF; use pyo3::{ pyclass, pyfunction, pymethods, - types::{PyAnyMethods, PyCapsule, PyDict, PyTuple}, + types::{PyCapsule, PyTuple}, Bound, PyObject, Python, }; use sedona_schema::datatypes::SedonaType; @@ -24,14 +42,14 @@ use crate::{ }; #[pyfunction] -pub fn register_sedona_scalar_udf<'py>( +pub fn sedona_scalar_udf<'py>( py: Python<'py>, name: &str, py_return_field: PyObject, py_invoke_batch: PyObject, - volatility: Option<&str>, + volatility: &str, ) -> Result, PySedonaError> { - let volatility = match volatility.unwrap_or("immutable") { + let volatility = match volatility { "immutable" => Volatility::Immutable, "stable" => Volatility::Stable, "volatile" => Volatility::Volatile, @@ -81,7 +99,7 @@ impl ScalarUDFImpl for PySedonaScalarUdf { } fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { - Ok(eval_return_field(&self.py_return_field, args.arg_fields)?) + Ok(eval_return_field(&self.py_return_field, args)?) } fn coerce_types(&self, arg_types: &[DataType]) -> Result> { @@ -93,15 +111,24 @@ impl ScalarUDFImpl for PySedonaScalarUdf { } } -fn eval_return_field(func: &PyObject, arg_fields: &[FieldRef]) -> Result { +fn eval_return_field(func: &PyObject, args: ReturnFieldArgs) -> Result { let return_field = Python::with_gil(|py| -> Result { - let py_arg_fields = arg_fields + let py_arg_fields = args + .arg_fields .iter() .map(|f| PySedonaField::new(f.as_ref().clone())) .collect::>(); + let py_scalar_values = zip(args.arg_fields, args.scalar_arguments) + .map(|(field, maybe_arg)| { + maybe_arg.map(|arg| PySedonaValue { + field: field.clone(), + value: ColumnarValue::Scalar(arg.clone()), + num_rows: 1, + }) + }) + .collect::>(); - let py_args = PyTuple::new(py, py_arg_fields)?; - let py_return_field = func.call(py, py_args, None)?; + let py_return_field = func.call(py, (py_arg_fields, py_scalar_values), None)?; if py_return_field.is_none(py) { return Err(PySedonaError::DF(Box::new(not_impl_datafusion_err!( "Python Udf does not apply to arguments" @@ -130,10 +157,8 @@ fn eval_invoke_batch( let py_return_field = PySedonaField::new(args.return_field.as_ref().clone()); let py_args = PyTuple::new(py, py_values)?; - let py_kwargs = PyDict::new(py); - py_kwargs.set_item("return_field", py_return_field)?; - let result = func.call(py, py_args, Some(&py_kwargs))?; + let result = func.call(py, (py_args, py_return_field, args.number_rows), None)?; let (result_field, result_array) = import_arrow_array(result.bind(py))?; let result_sedona_type = SedonaType::from_storage_field(&result_field)?; From ac0202e2fccc49767da26639bc9298486a82fc00 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Fri, 17 Oct 2025 22:50:05 -0500 Subject: [PATCH 03/25] it works --- python/sedonadb/src/context.rs | 12 ++++++++++-- python/sedonadb/src/import_from.rs | 15 +++++++++++++-- python/sedonadb/src/udf.rs | 23 ++++++++++++++--------- 3 files changed, 37 insertions(+), 13 deletions(-) diff --git a/python/sedonadb/src/context.rs b/python/sedonadb/src/context.rs index 0e39a7caca..79ee7eb17f 100644 --- a/python/sedonadb/src/context.rs +++ b/python/sedonadb/src/context.rs @@ -21,8 +21,10 @@ use sedona::context::SedonaContext; use tokio::runtime::Runtime; use crate::{ - dataframe::InternalDataFrame, error::PySedonaError, - import_from::import_table_provider_from_any, runtime::wait_for_future, + dataframe::InternalDataFrame, + error::PySedonaError, + import_from::{import_ffi_scalar_udf, import_table_provider_from_any}, + runtime::wait_for_future, }; #[pyclass] @@ -116,4 +118,10 @@ impl InternalContext { self.inner.ctx.deregister_table(table_ref)?; Ok(()) } + + pub fn register_udf<'py>(&self, py: Python<'py>, udf: PyObject) -> Result<(), PySedonaError> { + let udf = import_ffi_scalar_udf(udf.bind(py))?; + self.inner.ctx.register_udf(udf); + Ok(()) + } } diff --git a/python/sedonadb/src/import_from.rs b/python/sedonadb/src/import_from.rs index f2ccca4c28..ded17caa3d 100644 --- a/python/sedonadb/src/import_from.rs +++ b/python/sedonadb/src/import_from.rs @@ -26,7 +26,11 @@ use arrow_array::{ }; use arrow_schema::{Field, Schema}; use datafusion::catalog::TableProvider; -use datafusion_ffi::table_provider::{FFI_TableProvider, ForeignTableProvider}; +use datafusion_expr::ScalarUDF; +use datafusion_ffi::{ + table_provider::{FFI_TableProvider, ForeignTableProvider}, + udf::{FFI_ScalarUDF, ForeignScalarUDF}, +}; use pyo3::{ types::{PyAnyMethods, PyCapsule, PyCapsuleMethods}, Bound, PyAny, Python, @@ -63,6 +67,13 @@ pub fn import_ffi_table_provider( Ok(Arc::new(provider)) } +pub fn import_ffi_scalar_udf(obj: &Bound) -> Result { + let capsule = obj.getattr("__datafusion_scalar_udf__")?.call0()?; + let udf_ptr = check_pycapsule(&capsule, "datafusion_scalar_udf")? as *mut FFI_ScalarUDF; + let udf: ForeignScalarUDF = unsafe { udf_ptr.as_ref().unwrap().try_into()? }; + Ok(udf.into()) +} + pub fn import_arrow_array_stream<'py>( py: Python<'py>, obj: &Bound, @@ -89,7 +100,7 @@ pub fn import_arrow_array_stream<'py>( } pub fn import_arrow_array(obj: &Bound) -> Result<(Field, ArrayRef), PySedonaError> { - let schema_and_array = obj.getattr("__arrow_c_schema__")?.call0()?; + let schema_and_array = obj.getattr("__arrow_c_array__")?.call0()?; let (schema_capsule, array_capsule): (Bound, Bound) = schema_and_array.extract()?; diff --git a/python/sedonadb/src/udf.rs b/python/sedonadb/src/udf.rs index 9b36840b1e..b70592018b 100644 --- a/python/sedonadb/src/udf.rs +++ b/python/sedonadb/src/udf.rs @@ -38,7 +38,7 @@ use sedona_schema::datatypes::SedonaType; use crate::{ error::PySedonaError, import_from::{import_arrow_array, import_arrow_field}, - schema::{PySedonaField, PySedonaType}, + schema::PySedonaType, }; #[pyfunction] @@ -116,8 +116,10 @@ fn eval_return_field(func: &PyObject, args: ReturnFieldArgs) -> Result>(); + .map(|f| -> Result<_, PySedonaError> { + Ok(PySedonaType::new(SedonaType::from_storage_field(f)?)) + }) + .collect::, _>>()?; let py_scalar_values = zip(args.arg_fields, args.scalar_arguments) .map(|(field, maybe_arg)| { maybe_arg.map(|arg| PySedonaValue { @@ -155,18 +157,18 @@ fn eval_invoke_batch( }) .collect::>(); - let py_return_field = PySedonaField::new(args.return_field.as_ref().clone()); + let expected_return_type = SedonaType::from_storage_field(&args.return_field)?; + let py_return_type = PySedonaType::new(expected_return_type.clone()); let py_args = PyTuple::new(py, py_values)?; - let result = func.call(py, (py_args, py_return_field, args.number_rows), None)?; + let result = func.call(py, (py_args, py_return_type, args.number_rows), None)?; let (result_field, result_array) = import_arrow_array(result.bind(py))?; let result_sedona_type = SedonaType::from_storage_field(&result_field)?; - let expected_result_sedona_type = SedonaType::from_storage_field(&args.return_field)?; - if expected_result_sedona_type != result_sedona_type { + if expected_return_type != result_sedona_type { return Err(PySedonaError::SedonaPython(format!( - "Expected {expected_result_sedona_type} but got {result_sedona_type}" + "Expected {expected_return_type} but got {result_sedona_type}" ))); } @@ -242,6 +244,9 @@ impl PySedonaValue { } fn __repr__(&self) -> String { - format!("PySedonaValue {self:?}") + let sedona_type = SedonaType::from_storage_field(&self.field) + .map(|t| t.to_string()) + .unwrap_or("".to_string()); + format!("PySedonaValue {}[{}]", sedona_type, self.num_rows) } } From a1372f50ad0d8ca99eed10b453f2cdc80a012cfc Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Fri, 17 Oct 2025 22:50:29 -0500 Subject: [PATCH 04/25] format --- python/sedonadb/python/sedonadb/udf.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/sedonadb/python/sedonadb/udf.py b/python/sedonadb/python/sedonadb/udf.py index 0eff543bf0..6a16ec1578 100644 --- a/python/sedonadb/python/sedonadb/udf.py +++ b/python/sedonadb/python/sedonadb/udf.py @@ -17,8 +17,8 @@ from sedonadb._lib import sedona_scalar_udf -class ScalarUdfImpl: +class ScalarUdfImpl: @property def name(self): raise NotImplementedError() @@ -34,4 +34,6 @@ def invoke_batch(self, args, return_type, num_rows): raise NotImplementedError() def __datafusion_scalar_udf__(self): - return sedona_scalar_udf(self.name, self.return_type, self.invoke_batch, self.volatility) + return sedona_scalar_udf( + self.name, self.return_type, self.invoke_batch, self.volatility + ) From 613317c705f0ad2521a9e2a17009c853aab9141c Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 20 Oct 2025 14:48:40 -0500 Subject: [PATCH 05/25] builds again --- Cargo.lock | 1 + python/sedonadb/Cargo.toml | 1 + python/sedonadb/python/sedonadb/udf.py | 36 ++-- python/sedonadb/src/import_from.rs | 32 ++++ python/sedonadb/src/schema.rs | 1 + python/sedonadb/src/udf.rs | 249 +++++++++++++------------ rust/sedona-schema/src/matchers.rs | 9 +- 7 files changed, 193 insertions(+), 136 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 94550ec3d3..59ce6bbcac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5206,6 +5206,7 @@ dependencies = [ "pyo3", "sedona", "sedona-adbc", + "sedona-expr", "sedona-geoparquet", "sedona-proj", "sedona-schema", diff --git a/python/sedonadb/Cargo.toml b/python/sedonadb/Cargo.toml index 98379bde25..939a48e4da 100644 --- a/python/sedonadb/Cargo.toml +++ b/python/sedonadb/Cargo.toml @@ -42,6 +42,7 @@ futures = { workspace = true } pyo3 = { version = "0.25.1" } sedona = { path = "../../rust/sedona" } sedona-adbc = { path = "../../rust/sedona-adbc" } +sedona-expr = { path = "../../rust/sedona-expr" } sedona-geoparquet = { path = "../../rust/sedona-geoparquet" } sedona-schema = { path = "../../rust/sedona-schema" } sedona-proj = { path = "../../c/sedona-proj", default-features = false } diff --git a/python/sedonadb/python/sedonadb/udf.py b/python/sedonadb/python/sedonadb/udf.py index 6a16ec1578..9a376e54e1 100644 --- a/python/sedonadb/python/sedonadb/udf.py +++ b/python/sedonadb/python/sedonadb/udf.py @@ -15,25 +15,39 @@ # specific language governing permissions and limitations # under the License. +from typing import Literal, Optional + from sedonadb._lib import sedona_scalar_udf class ScalarUdfImpl: - @property - def name(self): - raise NotImplementedError() + def __init__( + self, + invoke_batch, + return_type, + input_types=None, + volatility: Literal["immutable", "stable", "volatile"] = "immutable", + name: Optional[str] = None, + ): + if input_types is None and not callable(return_type): - @property - def volatility(self): - return "immutable" + def return_type_impl(*args, **kwargs): + return return_type - def return_type(self, arg_types, scalar_args): - raise NotImplementedError() + self._return_type = return_type_impl + else: + self._return_type = return_type - def invoke_batch(self, args, return_type, num_rows): - raise NotImplementedError() + self._invoke_batch = invoke_batch + self._input_types = input_types + self._name = name + self._volatility = volatility def __datafusion_scalar_udf__(self): return sedona_scalar_udf( - self.name, self.return_type, self.invoke_batch, self.volatility + self._invoke_batch, + self._return_type, + self._input_types, + self._volatility, + self.name, ) diff --git a/python/sedonadb/src/import_from.rs b/python/sedonadb/src/import_from.rs index ded17caa3d..e31b0bef52 100644 --- a/python/sedonadb/src/import_from.rs +++ b/python/sedonadb/src/import_from.rs @@ -36,6 +36,10 @@ use pyo3::{ Bound, PyAny, Python, }; use sedona::record_batch_reader_provider::RecordBatchReaderProvider; +use sedona_schema::{ + datatypes::SedonaType, + matchers::{ArgMatcher, TypeMatcher}, +}; use crate::error::PySedonaError; @@ -116,6 +120,34 @@ pub fn import_arrow_array(obj: &Bound) -> Result<(Field, ArrayRef), PySed Ok((result_field, make_array(result_array_data))) } +pub fn import_arg_matcher( + obj: &Bound, +) -> Result, PySedonaError> { + if let Ok(string_value) = obj.extract::() { + match string_value.as_str() { + "geometry" => return Ok(ArgMatcher::is_geometry()), + "geography" => return Ok(ArgMatcher::is_geography()), + "numeric" => return Ok(ArgMatcher::is_numeric()), + "string" => return Ok(ArgMatcher::is_string()), + "binary" => return Ok(ArgMatcher::is_binary()), + "boolean" => return Ok(ArgMatcher::is_boolean()), + v => { + return Err(PySedonaError::SedonaPython(format!( + "Can't interpret literal string '{v}' as ArgMatcher" + ))) + } + } + } + + let sedona_type = import_sedona_type(obj)?; + Ok(ArgMatcher::is_exact(sedona_type)) +} + +pub fn import_sedona_type(obj: &Bound) -> Result { + let field = import_arrow_field(obj)?; + Ok(SedonaType::from_storage_field(&field)?) +} + pub fn import_arrow_field(obj: &Bound) -> Result { let capsule = obj.getattr("__arrow_c_schema__")?.call0()?; let schema = diff --git a/python/sedonadb/src/schema.rs b/python/sedonadb/src/schema.rs index d9466ea455..d261043ca4 100644 --- a/python/sedonadb/src/schema.rs +++ b/python/sedonadb/src/schema.rs @@ -171,6 +171,7 @@ impl PySedonaField { } #[pyclass] +#[derive(Clone, Debug)] pub struct PySedonaType { pub inner: SedonaType, } diff --git a/python/sedonadb/src/udf.rs b/python/sedonadb/src/udf.rs index b70592018b..175254a5a4 100644 --- a/python/sedonadb/src/udf.rs +++ b/python/sedonadb/src/udf.rs @@ -21,33 +21,31 @@ use arrow_array::{ ffi::{FFI_ArrowArray, FFI_ArrowSchema}, ArrayRef, }; -use arrow_schema::{DataType, FieldRef}; -use datafusion_common::{not_impl_datafusion_err, Result, ScalarValue}; -use datafusion_expr::{ - ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, - TypeSignature, Volatility, -}; +use datafusion_common::{Result, ScalarValue}; +use datafusion_expr::{ColumnarValue, ScalarUDF, Volatility}; use datafusion_ffi::udf::FFI_ScalarUDF; use pyo3::{ pyclass, pyfunction, pymethods, - types::{PyCapsule, PyTuple}, + types::{PyCapsule, PyNone, PyTuple}, Bound, PyObject, Python, }; -use sedona_schema::datatypes::SedonaType; +use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; +use sedona_schema::{datatypes::SedonaType, matchers::ArgMatcher}; use crate::{ error::PySedonaError, - import_from::{import_arrow_array, import_arrow_field}, + import_from::{import_arg_matcher, import_arrow_array, import_sedona_type}, schema::PySedonaType, }; #[pyfunction] pub fn sedona_scalar_udf<'py>( py: Python<'py>, - name: &str, - py_return_field: PyObject, py_invoke_batch: PyObject, + py_return_type: PyObject, + py_input_types: Option>, volatility: &str, + name: &str, ) -> Result, PySedonaError> { let volatility = match volatility { "immutable" => Volatility::Immutable, @@ -60,141 +58,149 @@ pub fn sedona_scalar_udf<'py>( } }; - let udf_impl = PySedonaScalarUdf { - name: name.to_string(), - signature: Signature::new(TypeSignature::UserDefined, volatility), + let scalar_kernel = sedona_scalar_kernel(py, py_input_types, py_return_type, py_invoke_batch)?; + let sedona_scalar_udf = + SedonaScalarUDF::new(name, vec![Arc::new(scalar_kernel)], volatility, None); + let scalar_udf: ScalarUDF = sedona_scalar_udf.into(); + + let name = cr"datafusion_scalar_udf".into(); + let ffi_udf = FFI_ScalarUDF::from(Arc::new(scalar_udf)); + Ok(PyCapsule::new(py, ffi_udf, Some(name))?) +} + +fn sedona_scalar_kernel<'py>( + py: Python<'py>, + input_types: Option>, + py_return_field: PyObject, + py_invoke_batch: PyObject, +) -> Result { + let matcher = if let Some(input_types) = input_types { + let arg_matchers = input_types + .iter() + .map(|obj| import_arg_matcher(obj.bind(py))) + .collect::, _>>()?; + let return_type = import_sedona_type(py_return_field.bind(py))?; + Some(ArgMatcher::new(arg_matchers, return_type)) + } else { + None + }; + + let kernel_impl = PySedonaScalarKernel { + matcher, py_return_field, py_invoke_batch, }; - let name = cr"datafusion_scalar_udf".into(); - let udf = ScalarUDF::from(udf_impl); - let ffi_udf = FFI_ScalarUDF::from(Arc::new(udf)); - Ok(PyCapsule::new(py, ffi_udf, Some(name))?) + Ok(kernel_impl) } #[derive(Debug)] -struct PySedonaScalarUdf { - name: String, - signature: Signature, +struct PySedonaScalarKernel { + matcher: Option, py_return_field: PyObject, py_invoke_batch: PyObject, } -impl ScalarUDFImpl for PySedonaScalarUdf { - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn name(&self) -> &str { - &self.name +impl SedonaScalarKernel for PySedonaScalarKernel { + fn return_type(&self, _args: &[SedonaType]) -> Result> { + Err(PySedonaError::SedonaPython("Unexpected call to return_type()".to_string()).into()) } - fn signature(&self) -> &Signature { - &self.signature - } + fn return_type_from_args_and_scalars( + &self, + args: &[SedonaType], + scalar_args: &[Option<&ScalarValue>], + ) -> Result> { + if let Some(matcher) = &self.matcher { + let return_type = matcher.match_args(args)?; + return Ok(return_type); + } - fn return_type(&self, _arg_types: &[DataType]) -> Result { - Err(PySedonaError::SedonaPython("Unexpected call to return_type()".to_string()).into()) - } + let return_type = Python::with_gil(|py| -> Result, PySedonaError> { + let py_sedona_types = args + .iter() + .map(|arg| -> Result<_, PySedonaError> { Ok(PySedonaType::new(arg.clone())) }) + .collect::, _>>()?; + let py_scalar_values = zip(&py_sedona_types, scalar_args) + .map(|(sedona_type, maybe_arg)| { + maybe_arg.map(|arg| PySedonaValue { + sedona_type: sedona_type.clone(), + value: ColumnarValue::Scalar(arg.clone()), + num_rows: 1, + }) + }) + .collect::>(); - fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { - Ok(eval_return_field(&self.py_return_field, args)?) - } + let py_return_field = + self.py_return_field + .call(py, (py_sedona_types, py_scalar_values), None)?; + if py_return_field.is_none(py) { + return Ok(None); + } - fn coerce_types(&self, arg_types: &[DataType]) -> Result> { - Ok(arg_types.to_vec()) - } + let return_type = import_sedona_type(py_return_field.bind(py))?; + Ok(Some(return_type)) + })?; - fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - Ok(eval_invoke_batch(&self.py_invoke_batch, args)?) + Ok(return_type) } -} -fn eval_return_field(func: &PyObject, args: ReturnFieldArgs) -> Result { - let return_field = Python::with_gil(|py| -> Result { - let py_arg_fields = args - .arg_fields - .iter() - .map(|f| -> Result<_, PySedonaError> { - Ok(PySedonaType::new(SedonaType::from_storage_field(f)?)) - }) - .collect::, _>>()?; - let py_scalar_values = zip(args.arg_fields, args.scalar_arguments) - .map(|(field, maybe_arg)| { - maybe_arg.map(|arg| PySedonaValue { - field: field.clone(), - value: ColumnarValue::Scalar(arg.clone()), - num_rows: 1, + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result { + let result = Python::with_gil(|py| -> Result { + let py_values = zip(arg_types, args) + .map(|(sedona_type, arg)| PySedonaValue { + sedona_type: PySedonaType::new(sedona_type.clone()), + value: arg.clone(), + num_rows: 0, }) - }) - .collect::>(); - - let py_return_field = func.call(py, (py_arg_fields, py_scalar_values), None)?; - if py_return_field.is_none(py) { - return Err(PySedonaError::DF(Box::new(not_impl_datafusion_err!( - "Python Udf does not apply to arguments" - )))); - } + .collect::>(); - let return_field = import_arrow_field(py_return_field.bind(py))?; - Ok(Arc::new(return_field)) - })?; + // let expected_return_type = SedonaType::from_storage_field(&args.return_field)?; + // let py_return_type = PySedonaType::new(expected_return_type.clone()); + let py_return_type = PyNone::get(py); + let py_args = PyTuple::new(py, py_values)?; - Ok(return_field) -} + let result = self + .py_invoke_batch + .call(py, (py_args, py_return_type, 0), None)?; -fn eval_invoke_batch( - func: &PyObject, - args: ScalarFunctionArgs, -) -> Result { - let result = Python::with_gil(|py| -> Result { - let py_values = zip(&args.arg_fields, &args.args) - .map(|(f, arg)| PySedonaValue { - field: f.clone(), - value: arg.clone(), - num_rows: args.number_rows, - }) - .collect::>(); - - let expected_return_type = SedonaType::from_storage_field(&args.return_field)?; - let py_return_type = PySedonaType::new(expected_return_type.clone()); - let py_args = PyTuple::new(py, py_values)?; - - let result = func.call(py, (py_args, py_return_type, args.number_rows), None)?; - - let (result_field, result_array) = import_arrow_array(result.bind(py))?; - let result_sedona_type = SedonaType::from_storage_field(&result_field)?; - - if expected_return_type != result_sedona_type { - return Err(PySedonaError::SedonaPython(format!( - "Expected {expected_return_type} but got {result_sedona_type}" - ))); - } + let (_, result_array) = import_arrow_array(result.bind(py))?; + // let result_sedona_type = SedonaType::from_storage_field(&result_field)?; - Ok(result_array) - })?; + // if expected_return_type != result_sedona_type { + // return Err(PySedonaError::SedonaPython(format!( + // "Expected {expected_return_type} but got {result_sedona_type}" + // ))); + // } - if args.args.is_empty() { - return Ok(ColumnarValue::Array(result)); - } + Ok(result_array) + })?; - for arg in &args.args { - match arg { - ColumnarValue::Array(_) => return Ok(ColumnarValue::Array(result)), - ColumnarValue::Scalar(_) => {} + if args.is_empty() { + return Ok(ColumnarValue::Array(result)); } - } - Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( - &result, 0, - )?)) + for arg in args { + match arg { + ColumnarValue::Array(_) => return Ok(ColumnarValue::Array(result)), + ColumnarValue::Scalar(_) => {} + } + } + + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)) + } } #[pyclass] #[derive(Debug)] pub struct PySedonaValue { - pub field: FieldRef, + pub sedona_type: PySedonaType, pub value: ColumnarValue, pub num_rows: usize, } @@ -203,9 +209,7 @@ pub struct PySedonaValue { impl PySedonaValue { #[getter] fn r#type(&self) -> Result { - Ok(PySedonaType::new(SedonaType::from_storage_field( - &self.field, - )?)) + Ok(self.sedona_type.clone()) } fn is_scalar(&self) -> bool { @@ -214,7 +218,7 @@ impl PySedonaValue { fn to_array(&self) -> Result { Ok(PySedonaValue { - field: self.field.clone(), + sedona_type: self.sedona_type.clone(), value: ColumnarValue::Array(self.value.to_array(self.num_rows)?), num_rows: self.num_rows, }) @@ -225,7 +229,8 @@ impl PySedonaValue { py: Python<'py>, ) -> Result, PySedonaError> { let schema_capsule_name = CString::new("arrow_schema").unwrap(); - let ffi_schema = FFI_ArrowSchema::try_from(self.field.as_ref().clone())?; + let storage_field = self.sedona_type.inner.to_storage_field("", true)?; + let ffi_schema = FFI_ArrowSchema::try_from(storage_field)?; Ok(PyCapsule::new(py, ffi_schema, Some(schema_capsule_name))?) } @@ -244,9 +249,9 @@ impl PySedonaValue { } fn __repr__(&self) -> String { - let sedona_type = SedonaType::from_storage_field(&self.field) - .map(|t| t.to_string()) - .unwrap_or("".to_string()); - format!("PySedonaValue {}[{}]", sedona_type, self.num_rows) + format!( + "PySedonaValue {}[{}]", + self.sedona_type.inner, self.num_rows + ) } } diff --git a/rust/sedona-schema/src/matchers.rs b/rust/sedona-schema/src/matchers.rs index 57a74ddcd0..2992b054f0 100644 --- a/rust/sedona-schema/src/matchers.rs +++ b/rust/sedona-schema/src/matchers.rs @@ -150,9 +150,12 @@ impl ArgMatcher { /// Matches the given Arrow type using PartialEq pub fn is_arrow(data_type: DataType) -> Arc { - Arc::new(IsExact { - exact_type: SedonaType::Arrow(data_type), - }) + Self::is_exact(SedonaType::Arrow(data_type)) + } + + /// Matches the given [SedonaType] using PartialEq + pub fn is_exact(exact_type: SedonaType) -> Arc { + Arc::new(IsExact { exact_type }) } /// Matches any geography or geometry argument without considering Crs From eed58004d74b694cb026d03d12eb7feeead7cc8b Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 20 Oct 2025 15:22:36 -0500 Subject: [PATCH 06/25] working through tests --- python/sedonadb/python/sedonadb/context.py | 3 +++ python/sedonadb/python/sedonadb/udf.py | 25 ++++++++++++++++++++-- python/sedonadb/src/udf.rs | 1 + python/sedonadb/tests/test_udf.py | 18 ++++++++++++++++ 4 files changed, 45 insertions(+), 2 deletions(-) create mode 100644 python/sedonadb/tests/test_udf.py diff --git a/python/sedonadb/python/sedonadb/context.py b/python/sedonadb/python/sedonadb/context.py index 63486952eb..888aaaf8a8 100644 --- a/python/sedonadb/python/sedonadb/context.py +++ b/python/sedonadb/python/sedonadb/context.py @@ -170,6 +170,9 @@ def sql(self, sql: str) -> DataFrame: """ return DataFrame(self._impl, self._impl.sql(sql), self.options) + def register_udf(self, udf): + self._impl.register_udf(udf) + def connect() -> SedonaContext: """Create a new [SedonaContext][sedonadb.context.SedonaContext]""" diff --git a/python/sedonadb/python/sedonadb/udf.py b/python/sedonadb/python/sedonadb/udf.py index 9a376e54e1..a31cb1e0a4 100644 --- a/python/sedonadb/python/sedonadb/udf.py +++ b/python/sedonadb/python/sedonadb/udf.py @@ -20,6 +20,23 @@ from sedonadb._lib import sedona_scalar_udf +def arrow_udf( + return_type, + input_types=None, + volatility: Literal["immutable", "stable", "volatile"] = "immutable", + name: Optional[str] = None, +): + def decorator(func): + def func_wrapper(args, return_type, num_rows): + return func(*args) + + name = func.__name__ if hasattr(func, "__name__") else None + return ScalarUdfImpl(func_wrapper, return_type, input_types, volatility, name) + + # Decorator must always be used with parentheses + return decorator + + class ScalarUdfImpl: def __init__( self, @@ -40,7 +57,11 @@ def return_type_impl(*args, **kwargs): self._invoke_batch = invoke_batch self._input_types = input_types - self._name = name + if name is None and hasattr(invoke_batch, "__name__"): + self._name = invoke_batch.__name__ + else: + self._name = name + self._volatility = volatility def __datafusion_scalar_udf__(self): @@ -49,5 +70,5 @@ def __datafusion_scalar_udf__(self): self._return_type, self._input_types, self._volatility, - self.name, + self._name, ) diff --git a/python/sedonadb/src/udf.rs b/python/sedonadb/src/udf.rs index 175254a5a4..544fb5fddb 100644 --- a/python/sedonadb/src/udf.rs +++ b/python/sedonadb/src/udf.rs @@ -237,6 +237,7 @@ impl PySedonaValue { fn __arrow_c_array__<'py>( &self, py: Python<'py>, + requsted_schema: PyObject ) -> Result, PySedonaError> { let schema_capsule_name = CString::new("arrow_array").unwrap(); let out_size = match &self.value { diff --git a/python/sedonadb/tests/test_udf.py b/python/sedonadb/tests/test_udf.py new file mode 100644 index 0000000000..9c36bae6f5 --- /dev/null +++ b/python/sedonadb/tests/test_udf.py @@ -0,0 +1,18 @@ + +import pyarrow as pa +import sedonadb +from sedonadb import udf + +def test_basic_udf(con): + @udf.arrow_udf(pa.binary(), ["string", "numeric"]) + def some_udf(arg0, arg1): + arg0, arg1 = (pa.array(arg0.to_array()).to_pylist(), pa.array(arg1.to_array()).to_pylist()) + return pa.array( + (f"{item0} / {item1}".encode() for item0, item1 in zip(arg0, arg1)), + pa.binary() + ) + + assert some_udf._name == "some_udf" + + con.register_udf(some_udf) + con.sql("SELECT some_udf('abcd', 123)").show() From 34d045d9fe97f6edb3800e269482b42326e98e38 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 20 Oct 2025 16:47:40 -0500 Subject: [PATCH 07/25] passing test --- python/sedonadb/src/udf.rs | 66 ++++++++++++++++++++++-------- python/sedonadb/tests/test_udf.py | 17 +++++--- rust/sedona-expr/src/scalar_udf.rs | 14 ++++++- 3 files changed, 72 insertions(+), 25 deletions(-) diff --git a/python/sedonadb/src/udf.rs b/python/sedonadb/src/udf.rs index 544fb5fddb..209860664c 100644 --- a/python/sedonadb/src/udf.rs +++ b/python/sedonadb/src/udf.rs @@ -21,12 +21,13 @@ use arrow_array::{ ffi::{FFI_ArrowArray, FFI_ArrowSchema}, ArrayRef, }; +use arrow_schema::Field; use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ColumnarValue, ScalarUDF, Volatility}; use datafusion_ffi::udf::FFI_ScalarUDF; use pyo3::{ pyclass, pyfunction, pymethods, - types::{PyCapsule, PyNone, PyTuple}, + types::{PyCapsule, PyTuple}, Bound, PyObject, Python, }; use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; @@ -34,7 +35,7 @@ use sedona_schema::{datatypes::SedonaType, matchers::ArgMatcher}; use crate::{ error::PySedonaError, - import_from::{import_arg_matcher, import_arrow_array, import_sedona_type}, + import_from::{check_pycapsule, import_arg_matcher, import_arrow_array, import_sedona_type}, schema::PySedonaType, }; @@ -106,6 +107,14 @@ impl SedonaScalarKernel for PySedonaScalarKernel { Err(PySedonaError::SedonaPython("Unexpected call to return_type()".to_string()).into()) } + fn invoke_batch( + &self, + _arg_types: &[SedonaType], + _args: &[ColumnarValue], + ) -> Result { + Err(PySedonaError::SedonaPython("Unexpected call to invoke_batch()".to_string()).into()) + } + fn return_type_from_args_and_scalars( &self, args: &[SedonaType], @@ -145,37 +154,37 @@ impl SedonaScalarKernel for PySedonaScalarKernel { Ok(return_type) } - fn invoke_batch( + fn invoke_batch_from_args( &self, arg_types: &[SedonaType], args: &[ColumnarValue], + return_type: &SedonaType, + num_rows: usize, ) -> Result { let result = Python::with_gil(|py| -> Result { let py_values = zip(arg_types, args) .map(|(sedona_type, arg)| PySedonaValue { sedona_type: PySedonaType::new(sedona_type.clone()), value: arg.clone(), - num_rows: 0, + num_rows, }) .collect::>(); - // let expected_return_type = SedonaType::from_storage_field(&args.return_field)?; - // let py_return_type = PySedonaType::new(expected_return_type.clone()); - let py_return_type = PyNone::get(py); + let py_return_type = PySedonaType::new(return_type.clone()); let py_args = PyTuple::new(py, py_values)?; let result = self .py_invoke_batch .call(py, (py_args, py_return_type, 0), None)?; - let (_, result_array) = import_arrow_array(result.bind(py))?; - // let result_sedona_type = SedonaType::from_storage_field(&result_field)?; + let (result_field, result_array) = import_arrow_array(result.bind(py))?; + let result_sedona_type = SedonaType::from_storage_field(&result_field)?; - // if expected_return_type != result_sedona_type { - // return Err(PySedonaError::SedonaPython(format!( - // "Expected {expected_return_type} but got {result_sedona_type}" - // ))); - // } + if return_type != &result_sedona_type { + return Err(PySedonaError::SedonaPython(format!( + "Expected {return_type} but got {result_sedona_type}" + ))); + } Ok(result_array) })?; @@ -237,16 +246,37 @@ impl PySedonaValue { fn __arrow_c_array__<'py>( &self, py: Python<'py>, - requsted_schema: PyObject - ) -> Result, PySedonaError> { - let schema_capsule_name = CString::new("arrow_array").unwrap(); + requsted_schema: Option>, + ) -> Result<(Bound<'py, PyCapsule>, Bound<'py, PyCapsule>), PySedonaError> { + if let Some(requested_schema) = requsted_schema { + let ffi_requested_schema = unsafe { + FFI_ArrowSchema::from_raw(check_pycapsule(&requested_schema, "arrow_schema")? as _) + }; + let requested_type = + SedonaType::from_storage_field(&Field::try_from(&ffi_requested_schema)?)?; + if requested_type != self.sedona_type.inner { + return Err(PySedonaError::SedonaPython( + "requested type is not implemented for PySedonaValue".to_string(), + )); + } + } + + let schema_capsule_name = CString::new("arrow_schema").unwrap(); + let field = self.sedona_type.inner.to_storage_field("", true)?; + let ffi_schema = FFI_ArrowSchema::try_from(&field)?; + + let array_capsule_name = CString::new("arrow_array").unwrap(); let out_size = match &self.value { ColumnarValue::Array(array) => array.len(), ColumnarValue::Scalar(_) => 1, }; let array = self.value.to_array(out_size)?; let ffi_array = FFI_ArrowArray::new(&array.to_data()); - Ok(PyCapsule::new(py, ffi_array, Some(schema_capsule_name))?) + + Ok(( + PyCapsule::new(py, ffi_schema, Some(schema_capsule_name))?, + PyCapsule::new(py, ffi_array, Some(array_capsule_name))?, + )) } fn __repr__(&self) -> String { diff --git a/python/sedonadb/tests/test_udf.py b/python/sedonadb/tests/test_udf.py index 9c36bae6f5..d60f0b7ff8 100644 --- a/python/sedonadb/tests/test_udf.py +++ b/python/sedonadb/tests/test_udf.py @@ -1,18 +1,25 @@ - import pyarrow as pa -import sedonadb from sedonadb import udf +import pandas as pd + def test_basic_udf(con): @udf.arrow_udf(pa.binary(), ["string", "numeric"]) def some_udf(arg0, arg1): - arg0, arg1 = (pa.array(arg0.to_array()).to_pylist(), pa.array(arg1.to_array()).to_pylist()) + arg0, arg1 = ( + pa.array(arg0.to_array()).to_pylist(), + pa.array(arg1.to_array()).to_pylist(), + ) return pa.array( (f"{item0} / {item1}".encode() for item0, item1 in zip(arg0, arg1)), - pa.binary() + pa.binary(), ) assert some_udf._name == "some_udf" con.register_udf(some_udf) - con.sql("SELECT some_udf('abcd', 123)").show() + + pd.testing.assert_frame_equal( + con.sql("SELECT some_udf('abcd', 123) as col").to_pandas(), + pd.DataFrame({"col": [b"abcd / 123"]}), + ) diff --git a/rust/sedona-expr/src/scalar_udf.rs b/rust/sedona-expr/src/scalar_udf.rs index 7f4c187d52..4fca4f1f8f 100644 --- a/rust/sedona-expr/src/scalar_udf.rs +++ b/rust/sedona-expr/src/scalar_udf.rs @@ -83,6 +83,16 @@ pub trait SedonaScalarKernel: Debug { arg_types: &[SedonaType], args: &[ColumnarValue], ) -> Result; + + fn invoke_batch_from_args( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + _return_type: &SedonaType, + _num_rows: usize, + ) -> Result { + self.invoke_batch(arg_types, args) + } } /// Type definition for a Scalar kernel implementation function @@ -259,8 +269,8 @@ impl ScalarUDFImpl for SedonaScalarUDF { }) .collect::>(); - let (kernel, _) = self.return_type_impl(&arg_types, &arg_scalars)?; - kernel.invoke_batch(&arg_types, &args.args) + let (kernel, return_type) = self.return_type_impl(&arg_types, &arg_scalars)?; + kernel.invoke_batch_from_args(&arg_types, &args.args, &return_type, args.number_rows) } fn aliases(&self) -> &[String] { From 05275900cd0c2e389049deee6e6e4d0f9464ed5f Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 20 Oct 2025 17:35:56 -0500 Subject: [PATCH 08/25] better errors and tests --- python/sedonadb/python/sedonadb/udf.py | 38 ++++++++- python/sedonadb/src/udf.rs | 14 +++- python/sedonadb/tests/test_udf.py | 106 +++++++++++++++++++++---- 3 files changed, 139 insertions(+), 19 deletions(-) diff --git a/python/sedonadb/python/sedonadb/udf.py b/python/sedonadb/python/sedonadb/udf.py index a31cb1e0a4..48a4382747 100644 --- a/python/sedonadb/python/sedonadb/udf.py +++ b/python/sedonadb/python/sedonadb/udf.py @@ -30,13 +30,47 @@ def decorator(func): def func_wrapper(args, return_type, num_rows): return func(*args) - name = func.__name__ if hasattr(func, "__name__") else None - return ScalarUdfImpl(func_wrapper, return_type, input_types, volatility, name) + name_arg = func.__name__ if name is None and hasattr(func, "__name__") else name + return ScalarUdfImpl( + func_wrapper, return_type, input_types, volatility, name_arg + ) # Decorator must always be used with parentheses return decorator +class TypeMatcher(str): + """Helper class to mark type matchers that can be used as the `input_types` for + user-defined functions + + Note that the internal storage of the type matcher (currently a string) is + arbitrary and may change in a future release. Use the constants provided by + the `udf` module. + """ + + pass + + +BINARY: TypeMatcher = "binary" +"""Match any binary argument (i.e., binary, binary view, large binary, +fixed-size binary)""" + +BOOLEAN: TypeMatcher = "boolean" +"""Match a boolean argument""" + +GEOGRAPHY: TypeMatcher = "geometry" +"""Match a geometry argument""" + +GEOMETRY: TypeMatcher = "geography" +"""Match a geography argument""" + +NUMERIC: TypeMatcher = "numeric" +"""Match any numeric argument""" + +STRING: TypeMatcher = "string" +"""Match any string argument (i.e., string, string view, large string)""" + + class ScalarUdfImpl: def __init__( self, diff --git a/python/sedonadb/src/udf.rs b/python/sedonadb/src/udf.rs index 209860664c..bddee40417 100644 --- a/python/sedonadb/src/udf.rs +++ b/python/sedonadb/src/udf.rs @@ -27,7 +27,7 @@ use datafusion_expr::{ColumnarValue, ScalarUDF, Volatility}; use datafusion_ffi::udf::FFI_ScalarUDF; use pyo3::{ pyclass, pyfunction, pymethods, - types::{PyCapsule, PyTuple}, + types::{PyAnyMethods, PyCapsule, PyTuple}, Bound, PyObject, Python, }; use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; @@ -176,13 +176,21 @@ impl SedonaScalarKernel for PySedonaScalarKernel { let result = self .py_invoke_batch .call(py, (py_args, py_return_type, 0), None)?; + let result_bound = result.bind(py); + if !result_bound.hasattr("__arrow_c_array__")? { + return Err( + PySedonaError::SedonaPython( + "Expected result of user-defined function to return an object implementing __arrow_c_array__()".to_string() + ) + ); + } - let (result_field, result_array) = import_arrow_array(result.bind(py))?; + let (result_field, result_array) = import_arrow_array(result_bound)?; let result_sedona_type = SedonaType::from_storage_field(&result_field)?; if return_type != &result_sedona_type { return Err(PySedonaError::SedonaPython(format!( - "Expected {return_type} but got {result_sedona_type}" + "Expected result of user-defined function to return array of type {return_type} but got {result_sedona_type}" ))); } diff --git a/python/sedonadb/tests/test_udf.py b/python/sedonadb/tests/test_udf.py index d60f0b7ff8..66a14ffc77 100644 --- a/python/sedonadb/tests/test_udf.py +++ b/python/sedonadb/tests/test_udf.py @@ -1,25 +1,103 @@ +import pandas as pd import pyarrow as pa +import pytest from sedonadb import udf -import pandas as pd -def test_basic_udf(con): - @udf.arrow_udf(pa.binary(), ["string", "numeric"]) - def some_udf(arg0, arg1): - arg0, arg1 = ( - pa.array(arg0.to_array()).to_pylist(), - pa.array(arg1.to_array()).to_pylist(), - ) - return pa.array( - (f"{item0} / {item1}".encode() for item0, item1 in zip(arg0, arg1)), - pa.binary(), - ) +def some_udf(arg0, arg1): + arg0, arg1 = ( + pa.array(arg0.to_array()).to_pylist(), + pa.array(arg1.to_array()).to_pylist(), + ) + return pa.array( + (f"{item0} / {item1}".encode() for item0, item1 in zip(arg0, arg1)), + pa.binary(), + ) + + +def test_udf_matchers(con): + udf_impl = udf.arrow_udf(pa.binary(), [udf.STRING, udf.NUMERIC])(some_udf) + assert udf_impl._name == "some_udf" + + con.register_udf(udf_impl) + pd.testing.assert_frame_equal( + con.sql("SELECT some_udf('abcd', 123) as col").to_pandas(), + pd.DataFrame({"col": [b"abcd / 123"]}), + ) + + +def test_udf_types(con): + udf_impl = udf.arrow_udf(pa.binary(), [pa.string(), pa.int64()])(some_udf) + assert udf_impl._name == "some_udf" + + con.register_udf(udf_impl) + pd.testing.assert_frame_equal( + con.sql("SELECT some_udf('abcd', 123) as col").to_pandas(), + pd.DataFrame({"col": [b"abcd / 123"]}), + ) + + +def test_udf_any_input(con): + udf_impl = udf.arrow_udf(pa.binary())(some_udf) + assert udf_impl._name == "some_udf" + + con.register_udf(udf_impl) + pd.testing.assert_frame_equal( + con.sql("SELECT some_udf('abcd', 123) as col").to_pandas(), + pd.DataFrame({"col": [b"abcd / 123"]}), + ) - assert some_udf._name == "some_udf" - con.register_udf(some_udf) +def test_udf_return_type_fn(con): + udf_impl = udf.arrow_udf(lambda arg_types, arg_scalars: pa.binary())(some_udf) + assert udf_impl._name == "some_udf" + con.register_udf(udf_impl) pd.testing.assert_frame_equal( con.sql("SELECT some_udf('abcd', 123) as col").to_pandas(), pd.DataFrame({"col": [b"abcd / 123"]}), ) + + +def test_udf_name(): + udf_impl = udf.arrow_udf(pa.binary(), name="foofy")(some_udf) + assert udf_impl._name == "foofy" + + +def test_udf_bad_return_object(con): + @udf.arrow_udf(pa.binary()) + def questionable_udf(arg): + return None + + con.register_udf(questionable_udf) + with pytest.raises( + ValueError, + match="Expected result of user-defined function to return an object implementing __arrow_c_array__", + ): + con.sql("SELECT questionable_udf(123) as col").to_pandas() + + +def test_udf_bad_return_type(con): + @udf.arrow_udf(pa.binary()) + def questionable_udf(arg): + return pa.array(["abc"], pa.string()) + + con.register_udf(questionable_udf) + with pytest.raises( + ValueError, + match="Expected result of user-defined function to return array of type Binary but got Utf8", + ): + con.sql("SELECT questionable_udf(123) as col").to_pandas() + + +def test_udf_bad_return_length(con): + @udf.arrow_udf(pa.binary()) + def questionable_udf(arg): + return pa.array([b"abc", b"def"], pa.binary()) + + con.register_udf(questionable_udf) + with pytest.raises( + ValueError, + match="UDF questionable_udf returned a different number of rows than expected. Expected: 1, Got: 2.", + ): + con.sql("SELECT questionable_udf(123) as col").to_pandas() From 482c8ad71e1b94db4dbf95b9f12081fe7d5e3e24 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 20 Oct 2025 17:44:11 -0500 Subject: [PATCH 09/25] test arrays --- python/sedonadb/tests/test_udf.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/python/sedonadb/tests/test_udf.py b/python/sedonadb/tests/test_udf.py index 66a14ffc77..5179d17ca4 100644 --- a/python/sedonadb/tests/test_udf.py +++ b/python/sedonadb/tests/test_udf.py @@ -49,13 +49,26 @@ def test_udf_any_input(con): def test_udf_return_type_fn(con): - udf_impl = udf.arrow_udf(lambda arg_types, arg_scalars: pa.binary())(some_udf) + udf_impl = udf.arrow_udf(lambda arg_types, arg_scalars: arg_types[0])(some_udf) assert udf_impl._name == "some_udf" con.register_udf(udf_impl) pd.testing.assert_frame_equal( - con.sql("SELECT some_udf('abcd', 123) as col").to_pandas(), - pd.DataFrame({"col": [b"abcd / 123"]}), + con.sql("SELECT some_udf('abcd'::BYTEA, 123) as col").to_pandas(), + pd.DataFrame({"col": [b"b'abcd' / 123"]}), + ) + + +def test_udf_array_input(con): + udf_impl = udf.arrow_udf(pa.binary(), [udf.STRING, udf.NUMERIC])(some_udf) + assert udf_impl._name == "some_udf" + + con.register_udf(udf_impl) + pd.testing.assert_frame_equal( + con.sql( + "SELECT some_udf(x, 123) as col FROM (VALUES ('a'), ('b'), ('c')) as t(x)" + ).to_pandas(), + pd.DataFrame({"col": [b"a / 123", b"b / 123", b"c / 123"]}), ) From 7bd708347a59bc3d4d83b634ee29681e5aec57e0 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 20 Oct 2025 17:53:14 -0500 Subject: [PATCH 10/25] fix matcher, license --- python/sedonadb/python/sedonadb/udf.py | 8 ++++---- python/sedonadb/src/udf.rs | 1 + python/sedonadb/tests/test_udf.py | 17 +++++++++++++++++ 3 files changed, 22 insertions(+), 4 deletions(-) diff --git a/python/sedonadb/python/sedonadb/udf.py b/python/sedonadb/python/sedonadb/udf.py index 48a4382747..9d75b2ab9f 100644 --- a/python/sedonadb/python/sedonadb/udf.py +++ b/python/sedonadb/python/sedonadb/udf.py @@ -58,12 +58,12 @@ class TypeMatcher(str): BOOLEAN: TypeMatcher = "boolean" """Match a boolean argument""" -GEOGRAPHY: TypeMatcher = "geometry" -"""Match a geometry argument""" - -GEOMETRY: TypeMatcher = "geography" +GEOGRAPHY: TypeMatcher = "geography" """Match a geography argument""" +GEOMETRY: TypeMatcher = "geometry" +"""Match a geometry argument""" + NUMERIC: TypeMatcher = "numeric" """Match any numeric argument""" diff --git a/python/sedonadb/src/udf.rs b/python/sedonadb/src/udf.rs index bddee40417..f486208183 100644 --- a/python/sedonadb/src/udf.rs +++ b/python/sedonadb/src/udf.rs @@ -251,6 +251,7 @@ impl PySedonaValue { Ok(PyCapsule::new(py, ffi_schema, Some(schema_capsule_name))?) } + #[pyo3(signature = (requsted_schema=None))] fn __arrow_c_array__<'py>( &self, py: Python<'py>, diff --git a/python/sedonadb/tests/test_udf.py b/python/sedonadb/tests/test_udf.py index 5179d17ca4..b9fa744ea1 100644 --- a/python/sedonadb/tests/test_udf.py +++ b/python/sedonadb/tests/test_udf.py @@ -1,3 +1,20 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + import pandas as pd import pyarrow as pa import pytest From 4c02bffd7a50f110c32d3c6244ba7dbfc5c2d4ba Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 20 Oct 2025 22:45:44 -0500 Subject: [PATCH 11/25] tests --- python/sedonadb/python/sedonadb/udf.py | 28 ++++++++++++++++-- python/sedonadb/src/udf.rs | 13 ++++++--- python/sedonadb/tests/test_udf.py | 40 ++++++++++++++++++++++++++ 3 files changed, 74 insertions(+), 7 deletions(-) diff --git a/python/sedonadb/python/sedonadb/udf.py b/python/sedonadb/python/sedonadb/udf.py index 9d75b2ab9f..c82e4d9f54 100644 --- a/python/sedonadb/python/sedonadb/udf.py +++ b/python/sedonadb/python/sedonadb/udf.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +import inspect from typing import Literal, Optional from sedonadb._lib import sedona_scalar_udf @@ -27,15 +28,29 @@ def arrow_udf( name: Optional[str] = None, ): def decorator(func): - def func_wrapper(args, return_type, num_rows): - return func(*args) + kwarg_names = callable_kwarg_only_names(func) + if "return_type" in kwarg_names and "num_rows" in kwarg_names: + + def func_wrapper(args, return_type, num_rows): + return func(*args, return_type=return_type, num_rows=num_rows) + elif "return_type" in kwarg_names: + + def func_wrapper(args, return_type, num_rows): + return func(*args, return_type=return_type) + elif "num_rows" in kwarg_names: + + def func_wrapper(args, return_type, num_rows): + return func(*args, num_rows=num_rows) + else: + + def func_wrapper(args, return_type, num_rows): + return func(*args) name_arg = func.__name__ if name is None and hasattr(func, "__name__") else name return ScalarUdfImpl( func_wrapper, return_type, input_types, volatility, name_arg ) - # Decorator must always be used with parentheses return decorator @@ -106,3 +121,10 @@ def __datafusion_scalar_udf__(self): self._volatility, self._name, ) + + +def callable_kwarg_only_names(f): + sig = inspect.signature(f) + return [ + k for k, p in sig.parameters.items() if p.kind == inspect.Parameter.KEYWORD_ONLY + ] diff --git a/python/sedonadb/src/udf.rs b/python/sedonadb/src/udf.rs index f486208183..61c91955e7 100644 --- a/python/sedonadb/src/udf.rs +++ b/python/sedonadb/src/udf.rs @@ -173,9 +173,9 @@ impl SedonaScalarKernel for PySedonaScalarKernel { let py_return_type = PySedonaType::new(return_type.clone()); let py_args = PyTuple::new(py, py_values)?; - let result = self - .py_invoke_batch - .call(py, (py_args, py_return_type, 0), None)?; + let result = + self.py_invoke_batch + .call(py, (py_args, py_return_type, num_rows), None)?; let result_bound = result.bind(py); if !result_bound.hasattr("__arrow_c_array__")? { return Err( @@ -289,8 +289,13 @@ impl PySedonaValue { } fn __repr__(&self) -> String { + let label = match &self.value { + ColumnarValue::Array(_) => "Array", + ColumnarValue::Scalar(_) => "Scalar", + }; + format!( - "PySedonaValue {}[{}]", + "PySedonaValue {label} {}[{}]", self.sedona_type.inner, self.num_rows ) } diff --git a/python/sedonadb/tests/test_udf.py b/python/sedonadb/tests/test_udf.py index b9fa744ea1..bc9b9c5b06 100644 --- a/python/sedonadb/tests/test_udf.py +++ b/python/sedonadb/tests/test_udf.py @@ -94,6 +94,46 @@ def test_udf_name(): assert udf_impl._name == "foofy" +def test_py_sedona_value(con): + @udf.arrow_udf(pa.int64()) + def fn_arg_only(arg): + assert repr(arg) == "PySedonaValue Array Int64[1]" + assert arg.is_scalar() is False + assert repr(arg.type) == "SedonaType int64" + + return pa.array(range(len(pa.array(arg)))) + + con.register_udf(fn_arg_only) + con.sql("SELECT fn_arg_only(123)").to_arrow_table() + + +def test_udf_kwargs(con): + @udf.arrow_udf(pa.int64()) + def fn_return_type(arg, *, return_type=None): + assert repr(return_type) == "SedonaType int64" + return pa.array(range(len(pa.array(arg)))) + + con.register_udf(fn_return_type) + con.sql("SELECT fn_return_type('123')").to_arrow_table() + + @udf.arrow_udf(pa.int64()) + def fn_num_rows(arg, *, num_rows=None): + assert num_rows == 1 + return pa.array(range(len(pa.array(arg)))) + + con.register_udf(fn_num_rows) + con.sql("SELECT fn_num_rows('123')").to_arrow_table() + + @udf.arrow_udf(pa.int64()) + def fn_num_rows_and_return_type(arg, *, num_rows=None, return_type=None): + assert repr(return_type) == "SedonaType int64" + assert num_rows == 1 + return pa.array(range(len(pa.array(arg)))) + + con.register_udf(fn_num_rows_and_return_type) + con.sql("SELECT fn_num_rows_and_return_type('123')").to_arrow_table() + + def test_udf_bad_return_object(con): @udf.arrow_udf(pa.binary()) def questionable_udf(arg): From b03c5b80e004c6525a12f287608a59ebd4130efa Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 20 Oct 2025 23:13:53 -0500 Subject: [PATCH 12/25] don't go through ffi --- python/sedonadb/python/sedonadb/context.py | 2 +- python/sedonadb/python/sedonadb/udf.py | 2 +- python/sedonadb/src/context.rs | 11 ++++------- python/sedonadb/src/udf.rs | 20 +++++++++++++++----- python/sedonadb/tests/test_udf.py | 6 +++--- 5 files changed, 24 insertions(+), 17 deletions(-) diff --git a/python/sedonadb/python/sedonadb/context.py b/python/sedonadb/python/sedonadb/context.py index 888aaaf8a8..b434bead55 100644 --- a/python/sedonadb/python/sedonadb/context.py +++ b/python/sedonadb/python/sedonadb/context.py @@ -171,7 +171,7 @@ def sql(self, sql: str) -> DataFrame: return DataFrame(self._impl, self._impl.sql(sql), self.options) def register_udf(self, udf): - self._impl.register_udf(udf) + self._impl.register_udf(udf.__sedona_internal_udf__()) def connect() -> SedonaContext: diff --git a/python/sedonadb/python/sedonadb/udf.py b/python/sedonadb/python/sedonadb/udf.py index c82e4d9f54..0d42e1fd19 100644 --- a/python/sedonadb/python/sedonadb/udf.py +++ b/python/sedonadb/python/sedonadb/udf.py @@ -113,7 +113,7 @@ def return_type_impl(*args, **kwargs): self._volatility = volatility - def __datafusion_scalar_udf__(self): + def __sedona_internal_udf__(self): return sedona_scalar_udf( self._invoke_batch, self._return_type, diff --git a/python/sedonadb/src/context.rs b/python/sedonadb/src/context.rs index 79ee7eb17f..6ffe0698bf 100644 --- a/python/sedonadb/src/context.rs +++ b/python/sedonadb/src/context.rs @@ -21,10 +21,8 @@ use sedona::context::SedonaContext; use tokio::runtime::Runtime; use crate::{ - dataframe::InternalDataFrame, - error::PySedonaError, - import_from::{import_ffi_scalar_udf, import_table_provider_from_any}, - runtime::wait_for_future, + dataframe::InternalDataFrame, error::PySedonaError, + import_from::import_table_provider_from_any, runtime::wait_for_future, udf::PySedonaScalarUdf, }; #[pyclass] @@ -119,9 +117,8 @@ impl InternalContext { Ok(()) } - pub fn register_udf<'py>(&self, py: Python<'py>, udf: PyObject) -> Result<(), PySedonaError> { - let udf = import_ffi_scalar_udf(udf.bind(py))?; - self.inner.ctx.register_udf(udf); + pub fn register_udf(&self, udf: PySedonaScalarUdf) -> Result<(), PySedonaError> { + self.inner.ctx.register_udf(udf.inner); Ok(()) } } diff --git a/python/sedonadb/src/udf.rs b/python/sedonadb/src/udf.rs index 61c91955e7..9566e9302d 100644 --- a/python/sedonadb/src/udf.rs +++ b/python/sedonadb/src/udf.rs @@ -24,7 +24,6 @@ use arrow_array::{ use arrow_schema::Field; use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ColumnarValue, ScalarUDF, Volatility}; -use datafusion_ffi::udf::FFI_ScalarUDF; use pyo3::{ pyclass, pyfunction, pymethods, types::{PyAnyMethods, PyCapsule, PyTuple}, @@ -47,7 +46,7 @@ pub fn sedona_scalar_udf<'py>( py_input_types: Option>, volatility: &str, name: &str, -) -> Result, PySedonaError> { +) -> Result { let volatility = match volatility { "immutable" => Volatility::Immutable, "stable" => Volatility::Stable, @@ -64,9 +63,7 @@ pub fn sedona_scalar_udf<'py>( SedonaScalarUDF::new(name, vec![Arc::new(scalar_kernel)], volatility, None); let scalar_udf: ScalarUDF = sedona_scalar_udf.into(); - let name = cr"datafusion_scalar_udf".into(); - let ffi_udf = FFI_ScalarUDF::from(Arc::new(scalar_udf)); - Ok(PyCapsule::new(py, ffi_udf, Some(name))?) + Ok(PySedonaScalarUdf { inner: scalar_udf }) } fn sedona_scalar_kernel<'py>( @@ -194,6 +191,13 @@ impl SedonaScalarKernel for PySedonaScalarKernel { ))); } + if result_array.len() != num_rows { + return Err(PySedonaError::SedonaPython(format!( + "Expected result of user-defined function to return array of length {num_rows} but got {}", + result_array.len() + ))); + } + Ok(result_array) })?; @@ -214,6 +218,12 @@ impl SedonaScalarKernel for PySedonaScalarKernel { } } +#[pyclass] +#[derive(Clone)] +pub struct PySedonaScalarUdf { + pub inner: ScalarUDF, +} + #[pyclass] #[derive(Debug)] pub struct PySedonaValue { diff --git a/python/sedonadb/tests/test_udf.py b/python/sedonadb/tests/test_udf.py index bc9b9c5b06..cdb49365d1 100644 --- a/python/sedonadb/tests/test_udf.py +++ b/python/sedonadb/tests/test_udf.py @@ -97,8 +97,8 @@ def test_udf_name(): def test_py_sedona_value(con): @udf.arrow_udf(pa.int64()) def fn_arg_only(arg): - assert repr(arg) == "PySedonaValue Array Int64[1]" - assert arg.is_scalar() is False + assert repr(arg) == "PySedonaValue Scalar Int64[1]" + assert arg.is_scalar() is True assert repr(arg.type) == "SedonaType int64" return pa.array(range(len(pa.array(arg)))) @@ -168,6 +168,6 @@ def questionable_udf(arg): con.register_udf(questionable_udf) with pytest.raises( ValueError, - match="UDF questionable_udf returned a different number of rows than expected. Expected: 1, Got: 2.", + match="Expected result of user-defined function to return array of length 1 but got 2", ): con.sql("SELECT questionable_udf(123) as col").to_pandas() From e10d875ff4993edab8e4fd49424b7b726cc119c7 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Tue, 21 Oct 2025 10:05:45 -0500 Subject: [PATCH 13/25] fix clippy --- python/sedonadb/src/context.rs | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/python/sedonadb/src/context.rs b/python/sedonadb/src/context.rs index 6ffe0698bf..16ee4e36df 100644 --- a/python/sedonadb/src/context.rs +++ b/python/sedonadb/src/context.rs @@ -21,8 +21,11 @@ use sedona::context::SedonaContext; use tokio::runtime::Runtime; use crate::{ - dataframe::InternalDataFrame, error::PySedonaError, - import_from::import_table_provider_from_any, runtime::wait_for_future, udf::PySedonaScalarUdf, + dataframe::InternalDataFrame, + error::PySedonaError, + import_from::{import_ffi_scalar_udf, import_table_provider_from_any}, + runtime::wait_for_future, + udf::PySedonaScalarUdf, }; #[pyclass] @@ -117,8 +120,20 @@ impl InternalContext { Ok(()) } - pub fn register_udf(&self, udf: PySedonaScalarUdf) -> Result<(), PySedonaError> { - self.inner.ctx.register_udf(udf.inner); + pub fn register_udf(&self, udf: Bound) -> Result<(), PySedonaError> { + if udf.hasattr("__sedona_internal_udf__")? { + let py_scalar_udf = udf + .getattr("__sedona_internal_udf__")? + .call0()? + .extract::()?; + self.inner.ctx.register_udf(py_scalar_udf.inner); + return Ok(()); + } else if udf.hasattr("__datafusion_scalar_udf__")? { + let scalar_udf = import_ffi_scalar_udf(&udf)?; + self.inner.ctx.register_udf(scalar_udf); + return Ok(()); + } + Ok(()) } } From a636572a81016bff0c43bf70d8b966d7ac325df7 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Tue, 21 Oct 2025 10:32:40 -0500 Subject: [PATCH 14/25] test again --- python/sedonadb/python/sedonadb/udf.py | 7 +++++-- python/sedonadb/src/context.rs | 8 +++++--- python/sedonadb/src/udf.rs | 23 +++++++++++++++++++---- 3 files changed, 29 insertions(+), 9 deletions(-) diff --git a/python/sedonadb/python/sedonadb/udf.py b/python/sedonadb/python/sedonadb/udf.py index 0d42e1fd19..6a9512c5ac 100644 --- a/python/sedonadb/python/sedonadb/udf.py +++ b/python/sedonadb/python/sedonadb/udf.py @@ -28,7 +28,7 @@ def arrow_udf( name: Optional[str] = None, ): def decorator(func): - kwarg_names = callable_kwarg_only_names(func) + kwarg_names = _callable_kwarg_only_names(func) if "return_type" in kwarg_names and "num_rows" in kwarg_names: def func_wrapper(args, return_type, num_rows): @@ -122,8 +122,11 @@ def __sedona_internal_udf__(self): self._name, ) + def __datafusion_scalar_udf__(self): + return self.__sedona_internal_udf__().__datafusion_scalar_udf__() -def callable_kwarg_only_names(f): + +def _callable_kwarg_only_names(f): sig = inspect.signature(f) return [ k for k, p in sig.parameters.items() if p.kind == inspect.Parameter.KEYWORD_ONLY diff --git a/python/sedonadb/src/context.rs b/python/sedonadb/src/context.rs index 16ee4e36df..0ac2fd197c 100644 --- a/python/sedonadb/src/context.rs +++ b/python/sedonadb/src/context.rs @@ -25,7 +25,7 @@ use crate::{ error::PySedonaError, import_from::{import_ffi_scalar_udf, import_table_provider_from_any}, runtime::wait_for_future, - udf::PySedonaScalarUdf, + udf::PyScalarUdf, }; #[pyclass] @@ -125,8 +125,10 @@ impl InternalContext { let py_scalar_udf = udf .getattr("__sedona_internal_udf__")? .call0()? - .extract::()?; - self.inner.ctx.register_udf(py_scalar_udf.inner); + .extract::()?; + self.inner + .ctx + .register_udf(py_scalar_udf.inner.as_ref().clone()); return Ok(()); } else if udf.hasattr("__datafusion_scalar_udf__")? { let scalar_udf = import_ffi_scalar_udf(&udf)?; diff --git a/python/sedonadb/src/udf.rs b/python/sedonadb/src/udf.rs index 9566e9302d..55698bdb7e 100644 --- a/python/sedonadb/src/udf.rs +++ b/python/sedonadb/src/udf.rs @@ -24,6 +24,7 @@ use arrow_array::{ use arrow_schema::Field; use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ColumnarValue, ScalarUDF, Volatility}; +use datafusion_ffi::udf::FFI_ScalarUDF; use pyo3::{ pyclass, pyfunction, pymethods, types::{PyAnyMethods, PyCapsule, PyTuple}, @@ -46,7 +47,7 @@ pub fn sedona_scalar_udf<'py>( py_input_types: Option>, volatility: &str, name: &str, -) -> Result { +) -> Result { let volatility = match volatility { "immutable" => Volatility::Immutable, "stable" => Volatility::Stable, @@ -63,7 +64,9 @@ pub fn sedona_scalar_udf<'py>( SedonaScalarUDF::new(name, vec![Arc::new(scalar_kernel)], volatility, None); let scalar_udf: ScalarUDF = sedona_scalar_udf.into(); - Ok(PySedonaScalarUdf { inner: scalar_udf }) + Ok(PyScalarUdf { + inner: Arc::new(scalar_udf), + }) } fn sedona_scalar_kernel<'py>( @@ -220,8 +223,20 @@ impl SedonaScalarKernel for PySedonaScalarKernel { #[pyclass] #[derive(Clone)] -pub struct PySedonaScalarUdf { - pub inner: ScalarUDF, +pub struct PyScalarUdf { + pub inner: Arc, +} + +#[pymethods] +impl PyScalarUdf { + fn __datafusion_scalar_udf__<'py>( + &self, + py: Python<'py>, + ) -> Result, PySedonaError> { + let capsule_name = CString::new("datafusion_scalar_udf").unwrap(); + let ffi_scalar_udf = FFI_ScalarUDF::from(self.inner.clone()); + Ok(PyCapsule::new(py, ffi_scalar_udf, Some(capsule_name))?) + } } #[pyclass] From bf8691f233b8340202088b7d4918e4e100a488ac Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Tue, 21 Oct 2025 12:10:01 -0500 Subject: [PATCH 15/25] test datafusion interop --- python/sedonadb/pyproject.toml | 1 + python/sedonadb/python/sedonadb/context.py | 2 +- python/sedonadb/src/context.rs | 37 +++++++++++--- python/sedonadb/src/udf.rs | 16 +++--- python/sedonadb/tests/test_udf.py | 58 ++++++++++++++++++++++ 5 files changed, 99 insertions(+), 15 deletions(-) diff --git a/python/sedonadb/pyproject.toml b/python/sedonadb/pyproject.toml index d857b3cb39..ff949dfd6c 100644 --- a/python/sedonadb/pyproject.toml +++ b/python/sedonadb/pyproject.toml @@ -33,6 +33,7 @@ dynamic = ["version"] test = [ "adbc-driver-manager[dbapi]", "adbc-driver-postgresql", + "datafusion", "duckdb", "geoarrow-pyarrow", "geopandas", diff --git a/python/sedonadb/python/sedonadb/context.py b/python/sedonadb/python/sedonadb/context.py index b434bead55..888aaaf8a8 100644 --- a/python/sedonadb/python/sedonadb/context.py +++ b/python/sedonadb/python/sedonadb/context.py @@ -171,7 +171,7 @@ def sql(self, sql: str) -> DataFrame: return DataFrame(self._impl, self._impl.sql(sql), self.options) def register_udf(self, udf): - self._impl.register_udf(udf.__sedona_internal_udf__()) + self._impl.register_udf(udf) def connect() -> SedonaContext: diff --git a/python/sedonadb/src/context.rs b/python/sedonadb/src/context.rs index 0ac2fd197c..4c480484e6 100644 --- a/python/sedonadb/src/context.rs +++ b/python/sedonadb/src/context.rs @@ -16,6 +16,7 @@ // under the License. use std::{collections::HashMap, sync::Arc}; +use datafusion_expr::ScalarUDFImpl; use pyo3::prelude::*; use sedona::context::SedonaContext; use tokio::runtime::Runtime; @@ -25,7 +26,7 @@ use crate::{ error::PySedonaError, import_from::{import_ffi_scalar_udf, import_table_provider_from_any}, runtime::wait_for_future, - udf::PyScalarUdf, + udf::PySedonaScalarUdf, }; #[pyclass] @@ -120,15 +121,36 @@ impl InternalContext { Ok(()) } - pub fn register_udf(&self, udf: Bound) -> Result<(), PySedonaError> { + pub fn scalar_udf(&self, name: &str) -> Result { + if let Some(sedona_scalar_udf) = self.inner.functions.scalar_udf(name) { + Ok(PySedonaScalarUdf { + inner: sedona_scalar_udf.clone(), + }) + } else { + Err(PySedonaError::SedonaPython(format!( + "Sedona scalar UDF with name {name} was not found" + ))) + } + } + + pub fn register_udf(&mut self, udf: Bound) -> Result<(), PySedonaError> { if udf.hasattr("__sedona_internal_udf__")? { let py_scalar_udf = udf .getattr("__sedona_internal_udf__")? .call0()? - .extract::()?; + .extract::()?; + let name = py_scalar_udf.inner.name(); self.inner - .ctx - .register_udf(py_scalar_udf.inner.as_ref().clone()); + .functions + .insert_scalar_udf(py_scalar_udf.inner.clone()); + self.inner.ctx.register_udf( + self.inner + .functions + .scalar_udf(name) + .unwrap() + .clone() + .into(), + ); return Ok(()); } else if udf.hasattr("__datafusion_scalar_udf__")? { let scalar_udf = import_ffi_scalar_udf(&udf)?; @@ -136,6 +158,9 @@ impl InternalContext { return Ok(()); } - Ok(()) + Err(PySedonaError::SedonaPython( + "Expected an object implementing __sedona_internal_udf__ or __datafusion_scalar_udf__" + .to_string(), + )) } } diff --git a/python/sedonadb/src/udf.rs b/python/sedonadb/src/udf.rs index 55698bdb7e..47796cfbf1 100644 --- a/python/sedonadb/src/udf.rs +++ b/python/sedonadb/src/udf.rs @@ -47,7 +47,7 @@ pub fn sedona_scalar_udf<'py>( py_input_types: Option>, volatility: &str, name: &str, -) -> Result { +) -> Result { let volatility = match volatility { "immutable" => Volatility::Immutable, "stable" => Volatility::Stable, @@ -62,10 +62,9 @@ pub fn sedona_scalar_udf<'py>( let scalar_kernel = sedona_scalar_kernel(py, py_input_types, py_return_type, py_invoke_batch)?; let sedona_scalar_udf = SedonaScalarUDF::new(name, vec![Arc::new(scalar_kernel)], volatility, None); - let scalar_udf: ScalarUDF = sedona_scalar_udf.into(); - Ok(PyScalarUdf { - inner: Arc::new(scalar_udf), + Ok(PySedonaScalarUdf { + inner: sedona_scalar_udf, }) } @@ -223,18 +222,19 @@ impl SedonaScalarKernel for PySedonaScalarKernel { #[pyclass] #[derive(Clone)] -pub struct PyScalarUdf { - pub inner: Arc, +pub struct PySedonaScalarUdf { + pub inner: SedonaScalarUDF, } #[pymethods] -impl PyScalarUdf { +impl PySedonaScalarUdf { fn __datafusion_scalar_udf__<'py>( &self, py: Python<'py>, ) -> Result, PySedonaError> { let capsule_name = CString::new("datafusion_scalar_udf").unwrap(); - let ffi_scalar_udf = FFI_ScalarUDF::from(self.inner.clone()); + let scalar_udf: ScalarUDF = self.inner.clone().into(); + let ffi_scalar_udf = FFI_ScalarUDF::from(Arc::new(scalar_udf)); Ok(PyCapsule::new(py, ffi_scalar_udf, Some(capsule_name))?) } } diff --git a/python/sedonadb/tests/test_udf.py b/python/sedonadb/tests/test_udf.py index cdb49365d1..4f6ff3ceed 100644 --- a/python/sedonadb/tests/test_udf.py +++ b/python/sedonadb/tests/test_udf.py @@ -171,3 +171,61 @@ def questionable_udf(arg): match="Expected result of user-defined function to return array of length 1 but got 2", ): con.sql("SELECT questionable_udf(123) as col").to_pandas() + + +def test_udf_datafusion_to_sedonadb(con): + udf_impl = udf.arrow_udf( + pa.binary(), [udf.STRING, udf.NUMERIC], name="some_external_udf" + )(some_udf) + + class UdfWrapper: + def __init__(self, obj): + self.obj = obj + + def __datafusion_scalar_udf__(self): + return self.obj.__datafusion_scalar_udf__() + + con.register_udf(UdfWrapper(udf_impl)) + pd.testing.assert_frame_equal( + con.sql("SELECT some_external_udf('abcd', 123) as col").to_pandas(), + pd.DataFrame({"col": [b"abcd / 123"]}), + ) + + +def test_udf_sedonadb_registry_function_to_datafusion(con): + datafusion = pytest.importorskip("datafusion") + udf_impl = udf.arrow_udf(pa.binary(), [udf.STRING, udf.NUMERIC])(some_udf) + + # Register with our session + con.register_udf(udf_impl) + + # Create a datafusion session, fetch our udf and register with the other session + datafusion_ctx = datafusion.SessionContext() + datafusion_ctx.register_udf( + datafusion.ScalarUDF.from_pycapsule(con._impl.scalar_udf("some_udf")) + ) + + # Can't quite use to_pandas() because there is a schema/batch nullability mismatch + batches = datafusion_ctx.sql("SELECT some_udf('abcd', 123) as col").collect() + assert len(batches) == 1 + pd.testing.assert_frame_equal( + batches[0].to_pandas(), + pd.DataFrame({"col": [b"abcd / 123"]}), + ) + + +def test_udf_sedonadb_to_datafusion(): + datafusion = pytest.importorskip("datafusion") + udf_impl = udf.arrow_udf(pa.binary(), [udf.STRING, udf.NUMERIC])(some_udf) + + # Create a datafusion session, register udf_impl directly + datafusion_ctx = datafusion.SessionContext() + datafusion_ctx.register_udf(datafusion.ScalarUDF.from_pycapsule(udf_impl)) + + # Can't quite use to_pandas() because there is a schema/batch nullability mismatch + batches = datafusion_ctx.sql("SELECT some_udf('abcd', 123) as col").collect() + assert len(batches) == 1 + pd.testing.assert_frame_equal( + batches[0].to_pandas(), + pd.DataFrame({"col": [b"abcd / 123"]}), + ) From 49d33b1e9e2d96a5c34697f93c9a69250e1f6516 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Tue, 21 Oct 2025 14:30:26 -0500 Subject: [PATCH 16/25] document arrow udf fn --- python/sedonadb/python/sedonadb/udf.py | 155 ++++++++++++++++++++++++- python/sedonadb/src/udf.rs | 20 +++- 2 files changed, 170 insertions(+), 5 deletions(-) diff --git a/python/sedonadb/python/sedonadb/udf.py b/python/sedonadb/python/sedonadb/udf.py index 6a9512c5ac..79eeb4fd47 100644 --- a/python/sedonadb/python/sedonadb/udf.py +++ b/python/sedonadb/python/sedonadb/udf.py @@ -16,17 +16,160 @@ # under the License. import inspect -from typing import Literal, Optional +from typing import Any, Literal, Optional from sedonadb._lib import sedona_scalar_udf +from sedonadb.utility import sedona # noqa: F401 def arrow_udf( - return_type, + return_type: Any, input_types=None, volatility: Literal["immutable", "stable", "volatile"] = "immutable", name: Optional[str] = None, ): + """Generic Arrow-based user-defined scalar function decorator + + This decorator may be used to annotate a function that accepts arguments as + Arrow array wrappers implementing the + [Arrow PyCapsule Interface](https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html). + The annotated function must return a value of a consistent length of the + appropriate type. + + !!! warning + SedonaDB will call the provided function from multiple threads. Attempts + to modify shared state from the body of the function may crash or cause + unusual behaviour. + + SedonaDB Python UDFs are experimental and this interface may change based on + user feedback. + + Args: + return_type: One of: + - A data type (e.g., pyarrow.DataType, arro3.core.DataType, nanoarrow.Schema) + if this function returns the same type regardless of its inputs. + - A function of `arg_types` (list of data types) and `scalar_args` (list of + optional scalars) that returns a data type. This function is also + responsible for returning `None` if this function does not apply to the + input types. + input_types: One of: + - A list where each member is a data type or a `TypeMatcher`. The + `udf.GEOMETRY` and `udf.GEOGRAPHY` type matchers are the most useful + because otherwise the function will only match spatial data types whose + coordinate reference system (CRS) also matches (i.e., based on simple + equality). Using these type matchers will also ensure input CRS consistency + and will automatically propagate input CRSes into the output. + - `None`, indicating that this function can accept any number of arguments + of any type. Usually this is paired with a functional `return_type` that + dynamically computes a return type or returns `None` if the number or + types of arguments do not match. + volatility: Use "immutable" for functions whose output is always consistent + for the same inputs (even between queries); use "stable" for functions + whose output is always consistent for the same inputs but only within + the same query, and use "volatile" for functions that generate random + or otherwise non-deterministic output. + name: An optional name for the UDF. If not given, it will be derived from + the name of the provided function. + + Examples: + + >>> import pyarrow as pa + >>> from sedonadb import udf + >>> sd = sedona.db.connect() + + The simplest scalar UDF only specifies return types. This implies that + the function can handle input of any type. + + >>> @udf.arrow_udf(pa.string()) + ... def some_udf(arg0, arg1): + ... arg0, arg1 = ( + ... pa.array(arg0.to_array()).to_pylist(), + ... pa.array(arg1.to_array()).to_pylist(), + ... ) + ... return pa.array( + ... (f"{item0} / {item1}" for item0, item1 in zip(arg0, arg1)), + ... pa.string(), + ... ) + ... + >>> sd.register_udf(some_udf) + >>> sd.sql("SELECT some_udf(123, 'abc') as col").show() + ┌───────────┐ + │ col │ + │ utf8 │ + ╞═══════════╡ + │ 123 / abc │ + └───────────┘ + + Use the `TypeMatcher` constants where possible to specify input. + This ensures that the function can handle the usual range of input + types that might exist for a given input. + + >>> @udf.arrow_udf(pa.int64(), [udf.STRING]) + ... def char_count(arg0): + ... arg0 = pa.array(arg0.to_array()) + ... + ... return pa.array( + ... (len(item) for item in arg0.to_pylist()), + ... pa.int64() + ... ) + ... + >>> sd.register_udf(char_count) + >>> sd.sql("SELECT char_count('abcde') as col").show() + ┌───────┐ + │ col │ + │ int64 │ + ╞═══════╡ + │ 5 │ + └───────┘ + + In this case, the type matcher ensures we can also use the function + for string view input which is the usual type SedonaDB emits when + reading Parquet files. + + >>> sd.sql("SELECT char_count(arrow_cast('abcde', 'Utf8View')) as col").show() + ┌───────┐ + │ col │ + │ int64 │ + ╞═══════╡ + │ 5 │ + └───────┘ + + Geometry UDFs are best written using Shapely because pyproj (including its use + in GeoPandas) is not thread safe and can crash when attempting to look up + CRSes when importing an Arrow array. The UDF framework supports returning + geometry storage to make this possible. Coordinate reference system metadata + is propagated automatically from the input. + + >>> import shapely + >>> import geoarrow.pyarrow as ga + >>> @udf.arrow_udf(ga.wkb(), [udf.GEOMETRY, udf.NUMERIC]) + ... def shapely_udf(geom, distance): + ... geom_wkb = pa.array(geom.storage.to_array()) + ... distance = pa.array(distance.to_array()) + ... geom = shapely.from_wkb(geom_wkb) + ... result_shapely = shapely.buffer(geom, distance) + ... return pa.array(shapely.to_wkb(result_shapely)) + ... + >>> + >>> sd.register_udf(shapely_udf) + >>> sd.sql("SELECT ST_SRID(shapely_udf(ST_Point(0, 0), 2.0)) as col").show() + ┌────────┐ + │ col │ + │ uint32 │ + ╞════════╡ + │ 0 │ + └────────┘ + + >>> sd.sql("SELECT ST_SRID(shapely_udf(ST_SetSRID(ST_Point(0, 0), 3857), 2.0)) as col").show() + ┌────────┐ + │ col │ + │ uint32 │ + ╞════════╡ + │ 3857 │ + └────────┘ + + """ + def decorator(func): kwarg_names = _callable_kwarg_only_names(func) if "return_type" in kwarg_names and "num_rows" in kwarg_names: @@ -87,6 +230,14 @@ class TypeMatcher(str): class ScalarUdfImpl: + """Scalar user-defined function wrapper + + This class is a wrapper class used as the return value for user-defined + function constructors. This wrapper allows the UDF to be registered with + a SedonaDB context or any context that accepts DataFusion Python + Scalar UDFs. This object is not intended to be used to call a UDF. + """ + def __init__( self, invoke_batch, diff --git a/python/sedonadb/src/udf.rs b/python/sedonadb/src/udf.rs index 47796cfbf1..c92ea2ee9f 100644 --- a/python/sedonadb/src/udf.rs +++ b/python/sedonadb/src/udf.rs @@ -188,9 +188,12 @@ impl SedonaScalarKernel for PySedonaScalarKernel { let result_sedona_type = SedonaType::from_storage_field(&result_field)?; if return_type != &result_sedona_type { - return Err(PySedonaError::SedonaPython(format!( - "Expected result of user-defined function to return array of type {return_type} but got {result_sedona_type}" - ))); + let return_type_storage = SedonaType::Arrow(return_type.storage_type().clone()); + if return_type_storage != result_sedona_type { + return Err(PySedonaError::SedonaPython(format!( + "Expected result of user-defined function to return array of type {return_type} or its storage but got {result_sedona_type}" + ))); + } } if result_array.len() != num_rows { @@ -258,6 +261,17 @@ impl PySedonaValue { matches!(&self.value, ColumnarValue::Scalar(_)) } + #[getter] + fn storage(&self) -> Result { + Ok(PySedonaValue { + sedona_type: PySedonaType { + inner: SedonaType::Arrow(self.sedona_type.inner.storage_type().clone()), + }, + value: self.value.clone(), + num_rows: self.num_rows, + }) + } + fn to_array(&self) -> Result { Ok(PySedonaValue { sedona_type: self.sedona_type.clone(), From 4fb2d001c1b5b14f22726d2a0049706c931c3cca Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Tue, 21 Oct 2025 14:39:15 -0500 Subject: [PATCH 17/25] geometry udf test --- python/sedonadb/tests/test_udf.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/python/sedonadb/tests/test_udf.py b/python/sedonadb/tests/test_udf.py index 4f6ff3ceed..8ea6fcd961 100644 --- a/python/sedonadb/tests/test_udf.py +++ b/python/sedonadb/tests/test_udf.py @@ -94,6 +94,35 @@ def test_udf_name(): assert udf_impl._name == "foofy" +def test_shapely_udf(con): + import shapely + import geoarrow.pyarrow as ga + import numpy as np + + @udf.arrow_udf(ga.wkb(), [udf.GEOMETRY, udf.NUMERIC]) + def shapely_udf(geom, distance): + geom_wkb = pa.array(geom.storage.to_array()) + distance = pa.array(distance.to_array()) + geom = shapely.from_wkb(geom_wkb) + result_shapely = shapely.buffer(geom, distance) + return pa.array(shapely.to_wkb(result_shapely)) + + con.register_udf(shapely_udf) + + pd.testing.assert_frame_equal( + con.sql("SELECT ST_Area(shapely_udf(ST_Point(0, 0), 2.0)) as col").to_pandas(), + pd.DataFrame({"col": [12.485780609032208]}), + ) + + # Ensure we can propagate a crs + pd.testing.assert_frame_equal( + con.sql( + "SELECT ST_SRID(shapely_udf(ST_SetSRID(ST_Point(0, 0), 3857), 2.0)) as col" + ).to_pandas(), + pd.DataFrame({"col": [3857]}, dtype=np.uint32), + ) + + def test_py_sedona_value(con): @udf.arrow_udf(pa.int64()) def fn_arg_only(arg): From dc7978f907a662ca33718ea61f687615c1eaefe2 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Tue, 21 Oct 2025 14:40:38 -0500 Subject: [PATCH 18/25] test geometry udf --- python/sedonadb/tests/test_udf.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/sedonadb/tests/test_udf.py b/python/sedonadb/tests/test_udf.py index 8ea6fcd961..3dea72673e 100644 --- a/python/sedonadb/tests/test_udf.py +++ b/python/sedonadb/tests/test_udf.py @@ -184,7 +184,11 @@ def questionable_udf(arg): con.register_udf(questionable_udf) with pytest.raises( ValueError, - match="Expected result of user-defined function to return array of type Binary but got Utf8", + match=( + "Expected result of user-defined function to " + "return array of type Binary or its storage " + "but got Utf8" + ), ): con.sql("SELECT questionable_udf(123) as col").to_pandas() From 8ebeaa3df35419945a4600a30d589b145d770e93 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Tue, 21 Oct 2025 14:55:17 -0500 Subject: [PATCH 19/25] Update python/sedonadb/src/udf.rs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- python/sedonadb/src/udf.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/sedonadb/src/udf.rs b/python/sedonadb/src/udf.rs index c92ea2ee9f..eeb7cdd9bb 100644 --- a/python/sedonadb/src/udf.rs +++ b/python/sedonadb/src/udf.rs @@ -290,13 +290,13 @@ impl PySedonaValue { Ok(PyCapsule::new(py, ffi_schema, Some(schema_capsule_name))?) } - #[pyo3(signature = (requsted_schema=None))] + #[pyo3(signature = (requested_schema=None))] fn __arrow_c_array__<'py>( &self, py: Python<'py>, - requsted_schema: Option>, + requested_schema: Option>, ) -> Result<(Bound<'py, PyCapsule>, Bound<'py, PyCapsule>), PySedonaError> { - if let Some(requested_schema) = requsted_schema { + if let Some(requested_schema) = requested_schema { let ffi_requested_schema = unsafe { FFI_ArrowSchema::from_raw(check_pycapsule(&requested_schema, "arrow_schema")? as _) }; From 91e2d2546742f78ebe516aa2b7c283910be4d045 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Tue, 21 Oct 2025 15:13:53 -0500 Subject: [PATCH 20/25] document register_udf --- python/sedonadb/python/sedonadb/context.py | 33 +++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/python/sedonadb/python/sedonadb/context.py b/python/sedonadb/python/sedonadb/context.py index 888aaaf8a8..f1c482734d 100644 --- a/python/sedonadb/python/sedonadb/context.py +++ b/python/sedonadb/python/sedonadb/context.py @@ -170,7 +170,38 @@ def sql(self, sql: str) -> DataFrame: """ return DataFrame(self._impl, self._impl.sql(sql), self.options) - def register_udf(self, udf): + def register_udf(self, udf: Any): + """Register a user-defined function + + Args: + udf: An object implementing the DataFusion PyCapsule protocol + (i.e., `__datafusion_scalar_udf__`) or a function annotated + with [arrow_udf][sedonadb.udf.arrow_udf]. + + Examples: + + >>> import pyarrow as pa + >>> from sedonadb import udf + >>> sd = sedona.db.connect() + >>> @udf.arrow_udf(pa.int64(), [udf.STRING]) + ... def char_count(arg0): + ... arg0 = pa.array(arg0.to_array()) + ... + ... return pa.array( + ... (len(item) for item in arg0.to_pylist()), + ... pa.int64() + ... ) + ... + >>> sd.register_udf(char_count) + >>> sd.sql("SELECT char_count('abcde') as col").show() + ┌───────┐ + │ col │ + │ int64 │ + ╞═══════╡ + │ 5 │ + └───────┘ + + """ self._impl.register_udf(udf) From 1f08f43743e8a2d5fdc6886d1d23074f9c0f6dbd Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Tue, 21 Oct 2025 15:18:13 -0500 Subject: [PATCH 21/25] fix formatting --- python/sedonadb/python/sedonadb/udf.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sedonadb/python/sedonadb/udf.py b/python/sedonadb/python/sedonadb/udf.py index 79eeb4fd47..a9f545fdf3 100644 --- a/python/sedonadb/python/sedonadb/udf.py +++ b/python/sedonadb/python/sedonadb/udf.py @@ -45,14 +45,14 @@ def arrow_udf( user feedback. Args: - return_type: One of: + return_type: One of - A data type (e.g., pyarrow.DataType, arro3.core.DataType, nanoarrow.Schema) if this function returns the same type regardless of its inputs. - A function of `arg_types` (list of data types) and `scalar_args` (list of optional scalars) that returns a data type. This function is also responsible for returning `None` if this function does not apply to the input types. - input_types: One of: + input_types: One of - A list where each member is a data type or a `TypeMatcher`. The `udf.GEOMETRY` and `udf.GEOGRAPHY` type matchers are the most useful because otherwise the function will only match spatial data types whose From 075531f15db61adac929b15ae965f02e899d2789 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Tue, 21 Oct 2025 15:41:55 -0500 Subject: [PATCH 22/25] add udf to docs --- docs/reference/python.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/reference/python.md b/docs/reference/python.md index b1b6cc45b2..5a93ab5aa4 100644 --- a/docs/reference/python.md +++ b/docs/reference/python.md @@ -25,3 +25,5 @@ ::: sedonadb.testing ::: sedonadb.dbapi + +::: sedonadb.udf From 97de908d05d536bea1c75f239b46b7056cc00d48 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Wed, 22 Oct 2025 14:56:39 -0500 Subject: [PATCH 23/25] type annotation --- python/sedonadb/python/sedonadb/udf.py | 28 +++++++++++++------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/python/sedonadb/python/sedonadb/udf.py b/python/sedonadb/python/sedonadb/udf.py index a9f545fdf3..909404b295 100644 --- a/python/sedonadb/python/sedonadb/udf.py +++ b/python/sedonadb/python/sedonadb/udf.py @@ -16,15 +16,27 @@ # under the License. import inspect -from typing import Any, Literal, Optional +from typing import Any, Literal, Optional, List, Union from sedonadb._lib import sedona_scalar_udf from sedonadb.utility import sedona # noqa: F401 +class TypeMatcher(str): + """Helper class to mark type matchers that can be used as the `input_types` for + user-defined functions + + Note that the internal storage of the type matcher (currently a string) is + arbitrary and may change in a future release. Use the constants provided by + the `udf` module. + """ + + pass + + def arrow_udf( return_type: Any, - input_types=None, + input_types: List[Union[TypeMatcher, Any]] = None, volatility: Literal["immutable", "stable", "volatile"] = "immutable", name: Optional[str] = None, ): @@ -197,18 +209,6 @@ def func_wrapper(args, return_type, num_rows): return decorator -class TypeMatcher(str): - """Helper class to mark type matchers that can be used as the `input_types` for - user-defined functions - - Note that the internal storage of the type matcher (currently a string) is - arbitrary and may change in a future release. Use the constants provided by - the `udf` module. - """ - - pass - - BINARY: TypeMatcher = "binary" """Match any binary argument (i.e., binary, binary view, large binary, fixed-size binary)""" From c7a61f42c649e7fdc59722c0460b25ad0e6040f1 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 23 Oct 2025 13:33:14 -0500 Subject: [PATCH 24/25] document the extra kwargs --- python/sedonadb/python/sedonadb/udf.py | 28 ++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/python/sedonadb/python/sedonadb/udf.py b/python/sedonadb/python/sedonadb/udf.py index 909404b295..d2126f1171 100644 --- a/python/sedonadb/python/sedonadb/udf.py +++ b/python/sedonadb/python/sedonadb/udf.py @@ -180,6 +180,34 @@ def arrow_udf( │ 3857 │ └────────┘ + Annotated functions may also declare keyword arguments `return_type` and/or `num_rows`, + which will be passed the appropriate value by the UDF framework. This facilitates writing + generic UDFs and/or UDFs with no arguments. + + >>> import numpy as np + >>> def random_impl(return_type, num_rows): + ... pa_type = pa.field(return_type).type + ... return pa.array(np.random.random(num_rows), pa_type) + ... + >>> @udf.arrow_udf(pa.float32(), []) + ... def random_f32(*, return_type=None, num_rows=None): + ... return random_impl(return_type, num_rows) + ... + >>> @udf.arrow_udf(pa.float64(), []) + ... def random_f64(*, return_type=None, num_rows=None): + ... return random_impl(return_type, num_rows) + ... + >>> np.random.seed(487) + >>> sd.register_udf(random_f32) + >>> sd.register_udf(random_f64) + >>> sd.sql("SELECT random_f32() AS f32, random_f64() as f64;").show() + ┌────────────┬─────────────────────┐ + │ f32 ┆ f64 │ + │ float32 ┆ float64 │ + ╞════════════╪═════════════════════╡ + │ 0.35385555 ┆ 0.24793247139474195 │ + └────────────┴─────────────────────┘ + """ def decorator(func): From 7efd4200bf6015ad321d3c32a5afad9e49e27bfb Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 23 Oct 2025 13:45:08 -0500 Subject: [PATCH 25/25] clarify why we need to sometimes wrap return type as a callable --- python/sedonadb/python/sedonadb/udf.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/sedonadb/python/sedonadb/udf.py b/python/sedonadb/python/sedonadb/udf.py index d2126f1171..236243c728 100644 --- a/python/sedonadb/python/sedonadb/udf.py +++ b/python/sedonadb/python/sedonadb/udf.py @@ -274,6 +274,10 @@ def __init__( volatility: Literal["immutable", "stable", "volatile"] = "immutable", name: Optional[str] = None, ): + # If the input_types are None, the return_type must be callable when passed + # to the internals. In the Python API we allow a data type as the return type + # to the argument easier to understand, which means we may have to wrap + # it in a callable here. if input_types is None and not callable(return_type): def return_type_impl(*args, **kwargs):