(extension_functions)=

Functions and table functions

Four hooks contribute functions written in Rust. Users can also define functions in pure Python — see {doc}../user-guide/common-operations/udf-and-udfa — and the two roads meet at the same registration methods.

HookContributesWrapped byRegistered with
__datafusion_scalar_udf__scalar function{py:func}datafusion.udf{py:meth}~datafusion.SessionContext.register_udf
__datafusion_aggregate_udf__aggregate function{py:func}datafusion.udaf{py:meth}~datafusion.SessionContext.register_udaf
__datafusion_window_udf__window function{py:func}datafusion.udwf{py:meth}~datafusion.SessionContext.register_udwf
__datafusion_table_function__function returning a table{py:func}datafusion.udtf{py:meth}~datafusion.SessionContext.register_udtf

All four are implemented in datafusion-ffi-example, one per file.

The three scalar-shaped hooks

The scalar, aggregate, and window getters take no argument beyond py — they need no codec and no task-context provider, because a function is identified by name and signature rather than by anything session-scoped:

#[pymethods]
impl MyScalarUDF {
    fn __datafusion_scalar_udf__<'py>(
        &self,
        py: Python<'py>,
    ) -> PyResult<Bound<'py, PyCapsule>> {
        let udf = Arc::new(ScalarUDF::from(self.clone()));
        let ffi = FFI_ScalarUDF::from(udf);

        PyCapsule::new_with_value(py, ffi, cr"datafusion_scalar_udf")
    }
}

Aggregate and window follow identically with FFI_AggregateUDF / FFI_WindowUDF and the matching capsule names.

Your users wrap the object once and register the result:

from datafusion import udf

ctx.register_udf(udf(my_library.MyScalarUDF()))

Table functions

A table function takes literal Expr arguments and returns a table provider, so it needs the host's logical codec the way a {ref}table provider <extension_providers> does:

fn __datafusion_table_function__<'py>(
    &self,
    py: Python<'py>,
    session: Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyCapsule>> {
    let func = self.clone();
    let codec = ffi_logical_codec_from_pycapsule(session, None)?;
    let provider = FFI_TableFunction::new_with_ffi_codec(Arc::new(func), None, codec);

    PyCapsule::new_with_value(py, provider, cr"datafusion_table_function")
}

Only literal expressions are supported as arguments. The Python side is described under {doc}Table Functions <../user-guide/common-operations/udf-and-udfa>.

Serializing functions

A function that appears in a plan leaving the process has to be reconstructible on the far side. Functions are the one case where a codec often needs no payload at all: the name is the whole encoding, try_encode_udf writes nothing, and try_decode_udf rebuilds the function from name. See {ref}extension_codecs for how that works and for the one obligation it puts on your decoder — with an empty payload there is no id to route on, so your try_decode_udf can be called with a name belonging to another library.

NameOnlyUdfCodec in datafusion-ffi-example is the worked case.