feat: Prototype calling Python UDFs from rust (#559)

sparrow-py/Cargo.toml

@@ -13,6 +13,7 @@ Python library for building and executing Kaskada queries.
 
 [dependencies]
 itertools = "0.11.0"
+arrow = {version = "44.0.0", features = ["pyarrow"] }
 pyo3 = {version = "0.19.1", features = ["abi3-py37", "extension-module", "generate-import-lib"]}
 
 [lib]

sparrow-py/pyproject.toml

@@ -5,6 +5,7 @@ description = "Kaskada query builder and local execution engine."
 authors = []
 
 [tool.poetry.dependencies]
+pandas = "^2.0.3"
 python = ">=3.8,<4.0"
 
 [tool.poetry.dev-dependencies]
@@ -20,6 +21,7 @@ furo = ">=2021.11.12"
 isort = ">=5.10.1"
 mypy = ">=0.930"
 pep8-naming = ">=0.12.1"
+pyarrow = "^12.0.1"
 pytest = ">=6.2.5"
 pyupgrade = ">=2.29.1"
 safety = ">=1.10.3"

sparrow-py/pysrc/sparrow_py/expr.py

@@ -27,7 +27,7 @@ def __init__(self, name: str, *args: "Expr") -> None:
 
         Raises
         ------
-        TypeError # noqa: DAR402
+        TypeError
             If the argument types are invalid for the given function.
         """
         ffi_args = [arg.ffi for arg in args]

sparrow-py/pysrc/sparrow_py/ffi.pyi

@@ -1,4 +1,6 @@
 from typing import Sequence
+from sparrow_py.udf import Udf
+import pyarrow as pa
 
 class Session:
     def __init__(self) -> None: ...
@@ -12,3 +14,5 @@ class Expr:
     def __repr__(self) -> str: ...
     def __str__(self) -> str: ...
     def data_type_string(self) -> str: ...
+
+def call_udf(udf: Udf, result_type: pa.DataType, *args: pa.Array) -> pa.Array: ...

sparrow-py/pysrc/sparrow_py/udf.py

@@ -0,0 +1,33 @@
+import functools
+import pandas as pd
+import pyarrow as pa
+from typing import Callable
+
+# TODO: Allow functions to return `pd.DataFrame` for struct arrays.
+FuncType = Callable[..., pd.Series]
+
+class Udf(object):
+    def __init__(self, name, func: FuncType, signature: str) -> None:
+        functools.update_wrapper(self, func)
+        self.name = name
+        self.func = func
+        self.signature = signature
+
+    def run_pyarrow(self, result_type: pa.DataType, *args: pa.Array) -> pa.Array:
+        # TODO: I believe this will return a series for simple arrays, and a
+        # dataframe for struct arrays. We should explore how this handles
+        # different types.
+        pd_args = [arg.to_pandas() for arg in args]
+        pd_result = self.func(*pd_args)
+
+        if isinstance(pd_result, pd.Series):
+            return pa.Array.from_pandas(pd_result, type=result_type)
+        else:
+            raise TypeError(f'Unsupported result type: {type(pd_result)}')
+
+
+def fenl_udf(name: str, signature: str):
+    def decorator(func: FuncType):
+        print(type(func))
+        return Udf(name, func, signature)
+    return decorator

sparrow-py/pytests/udf_test.py

@@ -0,0 +1,27 @@
+from sparrow_py.udf import Udf, fenl_udf
+from sparrow_py.ffi import call_udf
+import pandas as pd
+import pyarrow as pa
+
+@fenl_udf("add", "add(x: number, y: number) -> number")
+def add(x: pd.Series, y: pd.Series) -> pd.Series:
+    return x + y
+
+def test_numeric_udf_pure_python() -> None:
+    """Test the python side of UDFs."""
+    assert type(add) == Udf
+
+    x = pa.array([1, 12, 17, 23, 28], type=pa.int8())
+    y = pa.array([1, 13, 18, 20, 4], type=pa.int8())
+    result = add.run_pyarrow(pa.int8(), x, y)
+    print(result)
+    assert result == pa.array([2, 25, 35, 43, 32], type=pa.int8())
+
+
+def test_numeric_udf_rust() -> None:
+    """Test the rust side of UDFs."""
+    x = pa.array([1, 12, 17, 23, 28], type=pa.int8())
+    y = pa.array([1, 13, 18, 20, 4], type=pa.int8())
+    result = call_udf(add, pa.int8(), x, y)
+    print(result)
+    assert result == pa.array([2, 25, 35, 43, 32], type=pa.int8())

sparrow-py/src/lib.rs

@@ -2,12 +2,14 @@ use pyo3::prelude::*;
 
 mod expr;
 mod session;
+mod udf;
 
 /// A Python module implemented in Rust. The name of this function must match
 /// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
 /// import the module.
 #[pymodule]
 fn ffi(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(udf::call_udf, m)?)?;
     m.add_class::<session::Session>()?;
     m.add_class::<expr::Expr>()?;
 

sparrow-py/src/udf.rs

@@ -0,0 +1,32 @@
+use arrow::array::ArrayData;
+use arrow::datatypes::DataType;
+use arrow::pyarrow::{FromPyArrow, ToPyArrow};
+use pyo3::prelude::*;
+use pyo3::types::PyTuple;
+
+#[pyfunction]
+#[pyo3(signature = (udf, result_type, *args))]
+pub(super) fn call_udf<'py>(
+    py: Python<'py>,
+    udf: &'py PyAny,
+    result_type: &'py PyAny,
+    args: &'py PyTuple,
+) -> PyResult<&'py PyAny> {
+    let result_type = DataType::from_pyarrow(result_type)?;
+
+    // 1. Make sure we can convert each input to and from arrow arrays.
+    let mut udf_args = Vec::with_capacity(args.len() + 1);
+    udf_args.push(result_type.to_pyarrow(py)?);
+    for arg in args {
+        let array_data = ArrayData::from_pyarrow(arg)?;
+        let py_array: PyObject = array_data.to_pyarrow(py)?;
+        udf_args.push(py_array);
+    }
+    let args = PyTuple::new(py, udf_args);
+    let result = udf.call_method("run_pyarrow", args, None)?;
+
+    let array_data: ArrayData = ArrayData::from_pyarrow(result)?;
+    assert_eq!(array_data.data_type(), &result_type);
+
+    Ok(result)
+}