From 5a1ec8508444eeb45d89473c4c2f66ed6b64ba31 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 29 Oct 2026 11:01:37 -0600 Subject: [PATCH 1/5] blog: add post on distributed Python UDFs in Ballista --- .../blog/2026-10-29-ballista-python-udfs.md | 176 ++++++++++++++++++ 1 file changed, 176 insertions(+) create mode 100644 content/blog/2026-10-29-ballista-python-udfs.md 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..4871ada2 --- /dev/null +++ b/content/blog/2026-10-29-ballista-python-udfs.md @@ -0,0 +1,176 @@ +--- +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 + +Define a Python function that operates on Arrow arrays, wrap it as a UDF, and use it as you would with +DataFusion Python: + +```python +import pyarrow as pa +from datafusion import col, udf +from ballista import BallistaSessionContext + +ctx = BallistaSessionContext("df://localhost:50050") +ctx.sql("create external table t stored as parquet location 'data/t.parquet'") + + +def is_null(array: pa.Array) -> pa.Array: + return array.is_null() + + +is_null_udf = udf(is_null, [pa.int64()], pa.bool_(), "stable", name="is_null_udf") + +# DataFrame API +df = ctx.table("t") +df.select(col("a"), is_null_udf(col("a")).alias("n")).show() + +# SQL +ctx.register_udf(is_null_udf) +ctx.sql("select a, is_null_udf(a) as n from t").show() +``` + +The UDF does not need to be registered on the cluster ahead of time. Its code travels with the query. + +## A native DataFusion Python experience + +Earlier versions of the Ballista Python client subclassed DataFusion Python's `SessionContext` and `DataFrame`. +Every method that executed a plan was 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. +- Users got a Ballista-specific `DistributedDataFrame` rather than a regular `datafusion.DataFrame`. +- The wrapper had to track every change to the DataFusion Python API. +- 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. + +[#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. From 759af23355761973de02be55eec0c78d4569a6b4 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 29 Oct 2026 16:35:03 -0400 Subject: [PATCH 2/5] Use TPC-H Q19 as the Python UDF example Replace the is_null UDF with one that encodes the TPC-H Query 19 discount rules in a Python dictionary and evaluates them with pyarrow.compute kernels, so batches stay in Arrow. - Show the UDF from both the DataFrame API and SQL - Add steps to generate data with tpchgen-cli and start a local scheduler and executor from the ballista Python package - Print from the UDF so readers can see it run in the executor process Tested against a local Ballista cluster built from main. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../blog/2026-10-29-ballista-python-udfs.md | 138 ++++++++++++++++-- 1 file changed, 123 insertions(+), 15 deletions(-) diff --git a/content/blog/2026-10-29-ballista-python-udfs.md b/content/blog/2026-10-29-ballista-python-udfs.md index 4871ada2..f17a2ebd 100644 --- a/content/blog/2026-10-29-ballista-python-udfs.md +++ b/content/blog/2026-10-29-ballista-python-udfs.md @@ -42,34 +42,142 @@ Neither change would have been possible without a long series of improvements to ## What it looks like -Define a Python function that operates on Arrow arrays, wrap it as a UDF, and use it as you would with -DataFusion Python: +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 -from datafusion import col, udf +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.sql("create external table t stored as parquet location 'data/t.parquet'") - +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", +) -def is_null(array: pa.Array) -> pa.Array: - return array.is_null() +# 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() +``` -is_null_udf = udf(is_null, [pa.int64()], pa.bool_(), "stable", name="is_null_udf") +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. -# DataFrame API -df = ctx.table("t") -df.select(col("a"), is_null_udf(col("a")).alias("n")).show() +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: -# SQL -ctx.register_udf(is_null_udf) -ctx.sql("select a, is_null_udf(a) as n from t").show() +```text +is_discounted: 8192 rows in process 61867 +is_discounted: 8192 rows in process 61867 +... ``` -The UDF does not need to be registered on the cluster ahead of time. Its code travels with the query. +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 From 939004f6c39ed9a8784da921da122a7d6edae401 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 29 Oct 2026 16:40:51 -0400 Subject: [PATCH 3/5] Note DataFusion Python docs apply to Ballista and cache() gap Since Ballista now returns standard DataFusion Python objects, point readers to the DataFusion Python user guide and API reference, and contrast with the old wrapper, where methods it did not intercept ran locally. Note that DataFrame.cache() is not yet supported because its in-memory table cannot be serialized into the plan. Co-Authored-By: Claude Opus 5.5 (1M context) --- content/blog/2026-10-29-ballista-python-udfs.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/content/blog/2026-10-29-ballista-python-udfs.md b/content/blog/2026-10-29-ballista-python-udfs.md index f17a2ebd..8e6f84cf 100644 --- a/content/blog/2026-10-29-ballista-python-udfs.md +++ b/content/blog/2026-10-29-ballista-python-udfs.md @@ -199,6 +199,15 @@ Ballista's query planner sends the plan to the scheduler. Any operation that exe 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 From 6b1453c04b96bce579eac0c5e48dc9d560bbfedd Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 29 Oct 2026 16:45:22 -0400 Subject: [PATCH 4/5] Clarify drawbacks of the old DistributedDataFrame wrapper Explain the concrete cost of the Ballista-specific DataFrame: running on the cluster depended on re-wrapping every returned DataFrame, and write options were silently ignored. Note that tracking the DataFusion Python API caused churn in the Ballista repository, such as mirroring the Parquet writer options field by field. Co-Authored-By: Claude Opus 5.5 (1M context) --- content/blog/2026-10-29-ballista-python-udfs.md | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/content/blog/2026-10-29-ballista-python-udfs.md b/content/blog/2026-10-29-ballista-python-udfs.md index 8e6f84cf..2c612fcd 100644 --- a/content/blog/2026-10-29-ballista-python-udfs.md +++ b/content/blog/2026-10-29-ballista-python-udfs.md @@ -186,8 +186,12 @@ Every method that executed a plan was intercepted, the logical plan was serializ 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. -- Users got a Ballista-specific `DistributedDataFrame` rather than a regular `datafusion.DataFrame`. -- The wrapper had to track every change to the DataFusion Python API. +- 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. From fb03f31f7b5573dba130a498c13dd6f701a35433 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 29 Oct 2026 16:45:50 -0400 Subject: [PATCH 5/5] Correct which methods the old wrapper intercepted Only the main execution methods were intercepted, not every method that executed a plan, which is why others ran locally. Co-Authored-By: Claude Opus 5.5 (1M context) --- content/blog/2026-10-29-ballista-python-udfs.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/content/blog/2026-10-29-ballista-python-udfs.md b/content/blog/2026-10-29-ballista-python-udfs.md index 2c612fcd..2ea08a0d 100644 --- a/content/blog/2026-10-29-ballista-python-udfs.md +++ b/content/blog/2026-10-29-ballista-python-udfs.md @@ -182,8 +182,9 @@ run there, in a different Python process. ## A native DataFusion Python experience Earlier versions of the Ballista Python client subclassed DataFusion Python's `SessionContext` and `DataFrame`. -Every method that executed a plan was 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: +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