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/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 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/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 63486952eb..f1c482734d 100644 --- a/python/sedonadb/python/sedonadb/context.py +++ b/python/sedonadb/python/sedonadb/context.py @@ -170,6 +170,40 @@ def sql(self, sql: str) -> DataFrame: """ return DataFrame(self._impl, self._impl.sql(sql), self.options) + 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) + 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 new file mode 100644 index 0000000000..236243c728 --- /dev/null +++ b/python/sedonadb/python/sedonadb/udf.py @@ -0,0 +1,316 @@ +# 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 inspect +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: List[Union[TypeMatcher, Any]] = 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 │ + └────────┘ + + 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): + 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 + ) + + return decorator + + +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 = "geography" +"""Match a geography argument""" + +GEOMETRY: TypeMatcher = "geometry" +"""Match a geometry argument""" + +NUMERIC: TypeMatcher = "numeric" +"""Match any numeric argument""" + +STRING: TypeMatcher = "string" +"""Match any string argument (i.e., string, string view, large string)""" + + +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, + return_type, + input_types=None, + 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): + return return_type + + self._return_type = return_type_impl + else: + self._return_type = return_type + + self._invoke_batch = invoke_batch + self._input_types = input_types + if name is None and hasattr(invoke_batch, "__name__"): + self._name = invoke_batch.__name__ + else: + self._name = name + + self._volatility = volatility + + def __sedona_internal_udf__(self): + return sedona_scalar_udf( + self._invoke_batch, + self._return_type, + self._input_types, + self._volatility, + self._name, + ) + + def __datafusion_scalar_udf__(self): + return self.__sedona_internal_udf__().__datafusion_scalar_udf__() + + +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 0e39a7caca..4c480484e6 100644 --- a/python/sedonadb/src/context.rs +++ b/python/sedonadb/src/context.rs @@ -16,13 +16,17 @@ // under the License. use std::{collections::HashMap, sync::Arc}; +use datafusion_expr::ScalarUDFImpl; use pyo3::prelude::*; 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, + udf::PySedonaScalarUdf, }; #[pyclass] @@ -116,4 +120,47 @@ impl InternalContext { self.inner.ctx.deregister_table(table_ref)?; Ok(()) } + + 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::()?; + let name = py_scalar_udf.inner.name(); + self.inner + .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)?; + self.inner.ctx.register_udf(scalar_udf); + return Ok(()); + } + + Err(PySedonaError::SedonaPython( + "Expected an object implementing __sedona_internal_udf__ or __datafusion_scalar_udf__" + .to_string(), + )) + } } 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..e31b0bef52 100644 --- a/python/sedonadb/src/import_from.rs +++ b/python/sedonadb/src/import_from.rs @@ -20,18 +20,26 @@ 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 datafusion_expr::ScalarUDF; +use datafusion_ffi::{ + table_provider::{FFI_TableProvider, ForeignTableProvider}, + udf::{FFI_ScalarUDF, ForeignScalarUDF}, +}; use pyo3::{ types::{PyAnyMethods, PyCapsule, PyCapsuleMethods}, Bound, PyAny, Python, }; use sedona::record_batch_reader_provider::RecordBatchReaderProvider; +use sedona_schema::{ + datatypes::SedonaType, + matchers::{ArgMatcher, TypeMatcher}, +}; use crate::error::PySedonaError; @@ -63,6 +71,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, @@ -88,6 +103,59 @@ 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_array__")?.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_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 = + 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..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; +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}; @@ -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!(sedona_scalar_udf, m)?)?; m.add_class::()?; m.add_class::()?; 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 new file mode 100644 index 0000000000..eeb7cdd9bb --- /dev/null +++ b/python/sedonadb/src/udf.rs @@ -0,0 +1,341 @@ +// 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::{ + 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::{PyAnyMethods, PyCapsule, PyTuple}, + Bound, PyObject, Python, +}; +use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; +use sedona_schema::{datatypes::SedonaType, matchers::ArgMatcher}; + +use crate::{ + error::PySedonaError, + import_from::{check_pycapsule, import_arg_matcher, import_arrow_array, import_sedona_type}, + schema::PySedonaType, +}; + +#[pyfunction] +pub fn sedona_scalar_udf<'py>( + py: Python<'py>, + py_invoke_batch: PyObject, + py_return_type: PyObject, + py_input_types: Option>, + volatility: &str, + name: &str, +) -> Result { + let volatility = match volatility { + "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 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); + + Ok(PySedonaScalarUdf { + inner: sedona_scalar_udf, + }) +} + +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, + }; + + Ok(kernel_impl) +} + +#[derive(Debug)] +struct PySedonaScalarKernel { + matcher: Option, + py_return_field: PyObject, + py_invoke_batch: PyObject, +} + +impl SedonaScalarKernel for PySedonaScalarKernel { + fn return_type(&self, _args: &[SedonaType]) -> Result> { + 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], + scalar_args: &[Option<&ScalarValue>], + ) -> Result> { + if let Some(matcher) = &self.matcher { + let return_type = matcher.match_args(args)?; + return Ok(return_type); + } + + 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::>(); + + 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); + } + + let return_type = import_sedona_type(py_return_field.bind(py))?; + Ok(Some(return_type)) + })?; + + Ok(return_type) + } + + 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, + }) + .collect::>(); + + 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, num_rows), 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_bound)?; + let result_sedona_type = SedonaType::from_storage_field(&result_field)?; + + if return_type != &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 { + 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) + })?; + + if args.is_empty() { + return Ok(ColumnarValue::Array(result)); + } + + 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(Clone)] +pub struct PySedonaScalarUdf { + pub inner: SedonaScalarUDF, +} + +#[pymethods] +impl PySedonaScalarUdf { + fn __datafusion_scalar_udf__<'py>( + &self, + py: Python<'py>, + ) -> Result, PySedonaError> { + let capsule_name = CString::new("datafusion_scalar_udf").unwrap(); + 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))?) + } +} + +#[pyclass] +#[derive(Debug)] +pub struct PySedonaValue { + pub sedona_type: PySedonaType, + pub value: ColumnarValue, + pub num_rows: usize, +} + +#[pymethods] +impl PySedonaValue { + #[getter] + fn r#type(&self) -> Result { + Ok(self.sedona_type.clone()) + } + + fn is_scalar(&self) -> bool { + 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(), + 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 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))?) + } + + #[pyo3(signature = (requested_schema=None))] + fn __arrow_c_array__<'py>( + &self, + py: Python<'py>, + requested_schema: Option>, + ) -> Result<(Bound<'py, PyCapsule>, Bound<'py, PyCapsule>), PySedonaError> { + if let Some(requested_schema) = requested_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_schema, Some(schema_capsule_name))?, + PyCapsule::new(py, ffi_array, Some(array_capsule_name))?, + )) + } + + fn __repr__(&self) -> String { + let label = match &self.value { + ColumnarValue::Array(_) => "Array", + ColumnarValue::Scalar(_) => "Scalar", + }; + + format!( + "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 new file mode 100644 index 0000000000..3dea72673e --- /dev/null +++ b/python/sedonadb/tests/test_udf.py @@ -0,0 +1,264 @@ +# 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 sedonadb import udf + + +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"]}), + ) + + +def test_udf_return_type_fn(con): + 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'::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"]}), + ) + + +def test_udf_name(): + udf_impl = udf.arrow_udf(pa.binary(), name="foofy")(some_udf) + 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): + 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)))) + + 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): + 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 or its storage " + "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="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"]}), + ) 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] { 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