diff --git a/content/blog/2026-10-29-ballista-python-udfs.md b/content/blog/2026-10-29-ballista-python-udfs.md new file mode 100644 index 00000000..2ea08a0d --- /dev/null +++ b/content/blog/2026-10-29-ballista-python-udfs.md @@ -0,0 +1,298 @@ +--- +layout: post +title: Distributed Python UDFs in Apache DataFusion Ballista +date: 2026-10-29 +author: pmc +categories: [subprojects] +--- + + +[TOC] + +The Python client for [Apache DataFusion Ballista] can now run Python user-defined functions (UDFs) across a +cluster. Scalar, aggregate, and window UDFs written in Python are serialized into the query plan, shipped to +the scheduler, and executed on Ballista executors, from both the DataFrame API and SQL. + +This comes alongside a larger change in how the Ballista Python client works. Instead of wrapping +[DataFusion Python], Ballista now plugs into a standard `datafusion.SessionContext` through DataFusion's FFI +layer. If you know the DataFusion Python API, you already know how to use Ballista. + +Neither change would have been possible without a long series of improvements to DataFusion Python and the +`datafusion-ffi` crate. This post explains what changed, how it works, and who made it happen. + +[Apache DataFusion Ballista]: https://datafusion.apache.org/ballista/ +[DataFusion Python]: https://datafusion.apache.org/python/ + +## What it looks like + +Python UDFs are most useful when logic is awkward to write with built-in SQL functions but natural to express +in Python, such as business rules kept in a Python data structure. To show what this looks like on a cluster, +we will rewrite a query from the [TPC-H] benchmark, the same benchmark Ballista uses to measure its own +performance, so you can try it on data you generate yourself. + +[TPC-H Query 19] computes the revenue from parts that qualify for a discount. The rules differ by brand, and +the reference SQL spells them out as three nearly identical blocks joined with `OR`. With a Python UDF, the +rules can live in an ordinary Python dictionary, and the query is reduced to a join, the shared filters, and a +single call to the UDF. + +To follow along, generate the two tables the query needs with [tpchgen-cli]: + +```shell +pip install tpchgen-cli +tpchgen-cli parquet --scale-factor 1 --tables lineitem,part --output-dir data +``` + +Then start a local cluster. The `ballista` Python package can run both the scheduler and the executor. Start +each in its own terminal, from the same directory as `data/`, so the executor can find the files: + +```shell +# terminal 1: scheduler, listening on localhost:50050 +python -c "from ballista import BallistaScheduler; s = BallistaScheduler(); s.start(); s.wait_for_termination()" + +# terminal 2: executor +python -c "from ballista import BallistaExecutor; e = BallistaExecutor(); e.start(); e.wait_for_termination()" +``` + +Alternatively, `ballista.setup_test_cluster()` starts both inside the current Python process and returns the +host and port to connect to. Until the next release, the `ballista` package needs to be built from source, as +described in [Trying it out](#trying-it-out). + +The UDF is written with [pyarrow.compute], so each batch stays in Arrow rather than being converted to Python +objects: + +```python +import functools +import os + +import pyarrow as pa +import pyarrow.compute as pc +from datafusion import col, lit, udf +from datafusion import functions as f +from ballista import BallistaSessionContext + +ctx = BallistaSessionContext("df://localhost:50050") +ctx.register_parquet("lineitem", "data/lineitem.parquet") +ctx.register_parquet("part", "data/part.parquet") + +# brand -> (allowed containers, minimum quantity, maximum size) +DISCOUNT_RULES = { + "Brand#12": (["SM CASE", "SM BOX", "SM PACK", "SM PKG"], 1, 5), + "Brand#23": (["MED BAG", "MED BOX", "MED PKG", "MED PACK"], 10, 10), + "Brand#34": (["LG CASE", "LG BOX", "LG PACK", "LG PKG"], 20, 15), +} + + +def is_discounted( + brand: pa.Array, container: pa.Array, quantity: pa.Array, size: pa.Array +) -> pa.Array: + print(f"is_discounted: {len(brand)} rows in process {os.getpid()}", flush=True) + matches = [ + functools.reduce( + pc.and_, + [ + pc.equal(brand, b), + pc.is_in(container, pa.array(containers)), + pc.greater_equal(quantity, min_qty), + pc.less_equal(quantity, min_qty + 10), + pc.greater_equal(size, 1), + pc.less_equal(size, max_size), + ], + ) + for b, (containers, min_qty, max_size) in DISCOUNT_RULES.items() + ] + return functools.reduce(pc.or_, matches) + + +is_discounted_udf = udf( + is_discounted, + [pa.string(), pa.string(), pa.decimal128(15, 2), pa.int32()], + pa.bool_(), + "immutable", + name="is_discounted", +) + +# DataFrame API +lineitem = ctx.table("lineitem").filter( + f.in_list(col("l_shipmode"), [lit("AIR"), lit("AIR REG")]), + col("l_shipinstruct") == lit("DELIVER IN PERSON"), +) + +revenue = ( + lineitem.join(ctx.table("part"), left_on="l_partkey", right_on="p_partkey") + .filter( + is_discounted_udf( + col("p_brand"), col("p_container"), col("l_quantity"), col("p_size") + ) + ) + .aggregate( + [], [f.sum(col("l_extendedprice") * (lit(1) - col("l_discount"))).alias("revenue")] + ) +) +revenue.show() + +# SQL +ctx.register_udf(is_discounted_udf) +ctx.sql(""" + select sum(l_extendedprice * (1 - l_discount)) as revenue + from lineitem join part on p_partkey = l_partkey + where l_shipmode in ('AIR', 'AIR REG') + and l_shipinstruct = 'DELIVER IN PERSON' + and is_discounted(p_brand, p_container, l_quantity, p_size) +""").show() +``` + +The join and filters run as native DataFusion operators on the executors, and only the rows that survive them +are handed to `is_discounted`. The UDF does not need to be registered on the cluster ahead of time. Its code, +along with the `DISCOUNT_RULES` dictionary it references, travels with the query. + +You can see this happen. The `print` call at the top of `is_discounted` writes one line per batch, and those +lines appear in the executor's terminal, not in the terminal where you ran the query: + +```text +is_discounted: 8192 rows in process 61867 +is_discounted: 8192 rows in process 61867 +... +``` + +The client only prints the query results. The function was pickled into the plan, sent to the executor, and +run there, in a different Python process. + +[TPC-H]: https://www.tpc.org/tpch/ +[TPC-H Query 19]: https://github.com/apache/datafusion-ballista/blob/main/benchmarks/queries/q19.sql +[pyarrow.compute]: https://arrow.apache.org/docs/python/api/compute.html +[tpchgen-cli]: https://github.com/datafusion-contrib/tpcgen-rs + +## A native DataFusion Python experience + +Earlier versions of the Ballista Python client subclassed DataFusion Python's `SessionContext` and `DataFrame`. +The main execution methods, such as `collect`, `show`, and `to_pandas`, were intercepted, the logical plan was +serialized, and the query was resubmitted through a separate Ballista session. This worked for the common +cases, but it had real drawbacks: + +- Methods the wrapper did not know about ran locally instead of on the cluster. +- Results came back as a Ballista-specific `DistributedDataFrame`, and whether a query ran on the cluster + depended on the wrapper re-wrapping every DataFrame it returned. Write options such as `partition_by` were + silently ignored. +- The wrapper had to track every change to the DataFusion Python API, which caused churn in the Ballista + repository each time a feature was added upstream. For example, it mirrored DataFusion Python's Parquet + writer options field by field. +- Python objects did not always cross the boundary between the two libraries cleanly. The issue that started + this effort, [#1142], reported a `TypeError: argument 'group_by': 'Expr' object cannot be converted to 'Expr'` + from a simple aggregate. + +DataFusion Python 55 can install a custom query planner and extension codecs through FFI. Ballista now uses +this to configure a standard `datafusion.SessionContext`. DataFusion Python plans each query as usual, and +Ballista's query planner sends the plan to the scheduler. Any operation that executes a plan, including +`collect`, `show`, `to_pandas`, `write_parquet`, and `EXPLAIN ANALYZE`, runs on the cluster. The object you get +back from `BallistaSessionContext` is a `datafusion.SessionContext`, and the DataFrames it creates are regular +`datafusion.DataFrame`s. + +Because these are the standard DataFusion Python objects, the full DataFusion Python API and its +documentation, including the [user guide], [API reference], and examples, apply to Ballista as-is. Previously, +any method the wrapper did not intercept quietly ran locally instead of on the cluster. One known gap remains: +`DataFrame.cache()` is not yet supported, because the in-memory table it creates cannot be serialized into the +plan sent to the scheduler. + +[user guide]: https://datafusion.apache.org/python/user-guide/introduction.html +[API reference]: https://datafusion.apache.org/python/autoapi/index.html + +[#1142]: https://github.com/apache/datafusion-ballista/issues/1142 + +## How Python UDFs travel with the plan + +Ballista exchanges plans between the client, scheduler, and executors as protobuf, using DataFusion's +extension codecs for anything the built-in serialization does not cover. A Python UDF is a reference to a live +Python object, so it cannot be described by name alone. + +DataFusion Python 55 added codecs that solve this. When a plan contains a Python UDF, the codec pickles the +function with [cloudpickle] and embeds the payload inline in the plan. On the receiving side, the codec +recognizes the payload and reconstructs the UDF. Ballista composes these codecs with its own: + +- The client, scheduler, and executors started from the `ballista` Python package all use Python-aware + logical and physical codecs, with Ballista's default codecs underneath. +- Plans that do not contain Python UDFs serialize exactly as before, so there is no overhead for queries that + do not use them. + +Because the function body is shipped as a pickle, a few requirements follow: + +- The scheduler and executors must be started from the `ballista` Python package (`BallistaScheduler` and + `BallistaExecutor`), since decoding a Python UDF requires a Python interpreter. +- They must run the same Python minor version as the client. +- Any module the UDF imports must be installed wherever it runs. + +One known limitation remains: Python UDFs used as shuffle partitioning keys are not yet inlined. + +[cloudpickle]: https://github.com/cloudpipe/cloudpickle + +## The building blocks + +Most of the work behind this feature happened upstream, in DataFusion Python and the `datafusion-ffi` crate. +[Tim Saucer] has led the DataFusion FFI effort for a long time, and over the DataFusion 54 and 55 cycles he +added the pieces Ballista needed: + +- [Logical and physical extension codecs][dfp-1541] in DataFusion Python +- [Pickling of expressions][dfp-1544], then of [aggregate and window UDFs][dfp-1545] +- [Per-session control over UDF inlining][dfp-1546], with a strict mode that refuses to serialize UDFs +- [FFI query planner support][dfp-1677] +- [Composable extension codecs][dfp-1678], so projects like Ballista can layer their own codecs on top +- [Hardening the extension API][dfp-1759] before it shipped in 55 +- [A user guide and runnable examples][dfp-1547] for distributing expressions + +[Marko Milenković] then brought these together in Ballista. He rebuilt the Python client on the FFI query +planner ([#2588]) and wired the Python-aware codecs through the client, scheduler, and executors to support +UDFs ([#2619]). + +[Tim Saucer]: https://github.com/timsaucer +[Marko Milenković]: https://github.com/milenkovicm +[dfp-1541]: https://github.com/apache/datafusion-python/pull/1541 +[dfp-1544]: https://github.com/apache/datafusion-python/pull/1544 +[dfp-1545]: https://github.com/apache/datafusion-python/pull/1545 +[dfp-1546]: https://github.com/apache/datafusion-python/pull/1546 +[dfp-1547]: https://github.com/apache/datafusion-python/pull/1547 +[dfp-1677]: https://github.com/apache/datafusion-python/pull/1677 +[dfp-1678]: https://github.com/apache/datafusion-python/pull/1678 +[dfp-1759]: https://github.com/apache/datafusion-python/pull/1759 +[#2588]: https://github.com/apache/datafusion-ballista/pull/2588 +[#2619]: https://github.com/apache/datafusion-ballista/pull/2619 + +This is a good example of how the DataFusion ecosystem compounds. Extension points built to make DataFusion +Python more composable turned out to be exactly what a distributed engine needed to support Python natively. + +## Trying it out + +Python UDF support is on the `main` branch of Ballista and will be included in the next release of the +`ballista` Python package. Until then, you can build the package from source by following the instructions in +the [Python client README]. The [example script] starts a local test cluster and walks through both the +DataFrame and SQL forms shown above. + +Progress on the Python client is tracked in [#2618]. If you try it out, we would love to hear how it goes. +Please report bugs and feature requests on [GitHub], or join the conversation on the +[DataFusion communication channels]. + +[Python client README]: https://github.com/apache/datafusion-ballista/blob/main/python/README.md +[example script]: https://github.com/apache/datafusion-ballista/blob/main/python/examples/ballista_execution.py +[#2618]: https://github.com/apache/datafusion-ballista/issues/2618 +[GitHub]: https://github.com/apache/datafusion-ballista/issues +[DataFusion communication channels]: https://datafusion.apache.org/contributor-guide/communication.html + +## Acknowledgements + +Thank you to Tim Saucer for the DataFusion FFI and DataFusion Python work this builds on, to Marko Milenković +for integrating it into Ballista, and to everyone who contributes to DataFusion Python and the FFI crates.