User-Defined Functions

DataFusion provides powerful expressions and functions, reducing the need for custom Python functions. However you can still incorporate your own functions, i.e. User-Defined Functions (UDFs).

Scalar Functions

When writing a user-defined function that can operate on a row by row basis, these are called Scalar Functions. You can define your own scalar function by calling udf() .

The basic definition of a scalar UDF is a python function that takes one or more pyarrow arrays and returns a single array as output. DataFusion scalar UDFs operate on an entire batch of records at a time, though the evaluation of those records should be on a row by row basis. In the following example, we compute if the input array contains null values.

In [1]: import pyarrow

In [2]: import datafusion

In [3]: from datafusion import udf, col

In [4]: def is_null(array: pyarrow.Array) -> pyarrow.Array:
   ...:     return array.is_null()
   ...: 

In [5]: is_null_arr = udf(is_null, [pyarrow.int64()], pyarrow.bool_(), 'stable')

In [6]: ctx = datafusion.SessionContext()

In [7]: batch = pyarrow.RecordBatch.from_arrays(
   ...:     [pyarrow.array([1, None, 3]), pyarrow.array([4, 5, 6])],
   ...:     names=["a", "b"],
   ...: )
   ...: 

In [8]: df = ctx.create_dataframe([[batch]], name="batch_array")

In [9]: df.select(col("a"), is_null_arr(col("a")).alias("is_null")).show()
DataFrame()
+---+---------+
| a | is_null |
+---+---------+
| 1 | false   |
|   | true    |
| 3 | false   |
+---+---------+

In the previous example, we used the fact that pyarrow provides a variety of built in array functions such as is_null(). There are additional pyarrow compute functions available. When possible, it is highly recommended to use these functions because they can perform computations without doing any copy operations from the original arrays. This leads to greatly improved performance.

If you need to perform an operation in python that is not available with the pyarrow compute functions, you will need to convert the record batch into python values, perform your operation, and construct an array. This operation of converting the built in data type of the array into a python object can be one of the slowest operations in DataFusion, so it should be done sparingly.

The following example performs the same operation as before with is_null but demonstrates converting to Python objects to do the evaluation.

In [10]: import pyarrow

In [11]: import datafusion

In [12]: from datafusion import udf, col

In [13]: def is_null(array: pyarrow.Array) -> pyarrow.Array:
   ....:     return pyarrow.array([value.as_py() is None for value in array])
   ....: 

In [14]: is_null_arr = udf(is_null, [pyarrow.int64()], pyarrow.bool_(), 'stable')

In [15]: ctx = datafusion.SessionContext()

In [16]: batch = pyarrow.RecordBatch.from_arrays(
   ....:     [pyarrow.array([1, None, 3]), pyarrow.array([4, 5, 6])],
   ....:     names=["a", "b"],
   ....: )
   ....: 

In [17]: df = ctx.create_dataframe([[batch]], name="batch_array")

In [18]: df.select(col("a"), is_null_arr(col("a")).alias("is_null")).show()
DataFrame()
+---+---------+
| a | is_null |
+---+---------+
| 1 | false   |
|   | true    |
| 3 | false   |
+---+---------+

Aggregate Functions

The udaf() function allows you to define User-Defined Aggregate Functions (UDAFs). To use this you must implement an Accumulator that determines how the aggregation is performed.

When defining a UDAF there are four methods you need to implement. The update function takes the array(s) of input and updates the internal state of the accumulator. You should define this function to have as many input arguments as you will pass when calling the UDAF. Since aggregation may be split into multiple batches, we must have a method to combine multiple batches. For this, we have two functions, state and merge. state will return an array of scalar values that contain the current state of a single batch accumulation. Then we must merge the results of these different states. Finally evaluate is the call that will return the final result after the merge is complete.

In the following example we want to define a custom aggregate function that will return the difference between the sum of two columns. The state can be represented by a single value and we can also see how the inputs to update and merge differ.

import pyarrow
import pyarrow.compute
import datafusion
from datafusion import col, udaf, Accumulator
from typing import List

class MyAccumulator(Accumulator):
    """
    Interface of a user-defined accumulation.
    """
    def __init__(self):
        self._sum = 0.0

    def update(self, values_a: pyarrow.Array, values_b: pyarrow.Array) -> None:
        self._sum = self._sum + pyarrow.compute.sum(values_a).as_py() - pyarrow.compute.sum(values_b).as_py()

    def merge(self, states: List[pyarrow.Array]) -> None:
        self._sum = self._sum + pyarrow.compute.sum(states[0]).as_py()

    def state(self) -> pyarrow.Array:
        return pyarrow.array([self._sum])

    def evaluate(self) -> pyarrow.Scalar:
        return pyarrow.scalar(self._sum)

ctx = datafusion.SessionContext()
df = ctx.from_pydict(
    {
        "a": [4, 5, 6],
        "b": [1, 2, 3],
    }
)

my_udaf = udaf(MyAccumulator, [pyarrow.float64(), pyarrow.float64()], pyarrow.float64(), [pyarrow.float64()], 'stable')

df.aggregate([], [my_udaf(col("a"), col("b")).alias("col_diff")])

Window Functions

To implement a User-Defined Window Function (UDWF) you must call the udwf() function using a class that implements the abstract class WindowEvaluator.

There are three methods of evaluation of UDWFs.

• evaluate is the simplest case, where you are given an array and are expected to calculate the value for a single row of that array. This is the simplest case, but also the least performant.

• evaluate_all computes the values for all rows for an input array at a single time.

• evaluate_all_with_rank computes the values for all rows, but you only have the rank information for the rows.

Which methods you implement are based upon which of these options are set.

uses_window_frame	supports_bounded_execution	include_rank	function_to_implement
False (default)	False (default)	False (default)	evaluate_all
False	True	False	evaluate
False	True	False	evaluate_all_with_rank
True	True/False	True/False	evaluate

UDWF options

When you define your UDWF you can override the functions that return these values. They will determine which evaluate functions are called.

• uses_window_frame is set for functions that compute based on the specified window frame. If your function depends upon the specified frame, set this to True.

• supports_bounded_execution specifies if your function can be incrementally computed.

• include_rank is set to True for window functions that can be computed only using the rank information.

import pyarrow as pa
from datafusion import udwf, col, SessionContext
from datafusion.user_defined import WindowEvaluator

class ExponentialSmooth(WindowEvaluator):
    def __init__(self, alpha: float) -> None:
        self.alpha = alpha

    def evaluate_all(self, values: list[pa.Array], num_rows: int) -> pa.Array:
        results = []
        curr_value = 0.0
        values = values[0]
        for idx in range(num_rows):
            if idx == 0:
                curr_value = values[idx].as_py()
            else:
                curr_value = values[idx].as_py() * self.alpha + curr_value * (
                    1.0 - self.alpha
                )
            results.append(curr_value)

        return pa.array(results)

exp_smooth = udwf(
    ExponentialSmooth(0.9),
    pa.float64(),
    pa.float64(),
    volatility="immutable",
)

ctx = SessionContext()

df = ctx.from_pydict({
    "a": [1.0, 2.1, 2.9, 4.0, 5.1, 6.0, 6.9, 8.0]
})

df.select("a", exp_smooth(col("a")).alias("smooth_a")).show()

Table Functions

User Defined Table Functions are slightly different than the other functions described here. These functions take any number of Expr arguments, but only literal expressions are supported. Table functions must return a Table Provider as described in the ref:_io_custom_table_provider page.

Once you have a table function, you can register it with the session context by using datafusion.context.SessionContext.register_udtf().

There are examples of both rust backed and python based table functions in the examples folder of the repository. If you have a rust backed table function that you wish to expose via PyO3, you need to expose it as a PyCapsule.

#[pymethods]
impl MyTableFunction {
    fn __datafusion_table_function__<'py>(
        &self,
        py: Python<'py>,
    ) -> PyResult<Bound<'py, PyCapsule>> {
        let name = cr"datafusion_table_function".into();

        let func = self.clone();
        let provider = FFI_TableFunction::new(Arc::new(func), None);

        PyCapsule::new(py, provider, Some(name))
    }
}

previous

Window Functions

next

IO