milenkovicm opened a new pull request, #2619:
URL: https://github.com/apache/datafusion-ballista/pull/2619
## What
Python UDFs (scalar, aggregate and window) now work with
`BallistaSessionContext`. They are serialized with the plan as cloudpickle
payloads and run on executors, both from the DataFrame API and from SQL after
`ctx.register_udf(...)`.
```python
is_null_udf = udf(is_null, [pa.int64()], pa.bool_(), "stable",
name="is_null_udf")
df.select(col("a"), is_null_udf(col("a")).alias("n")).show()
ctx.register_udf(is_null_udf)
ctx.sql("select a, is_null_udf(a) as n from t").show()
```
Closes #.
Relates to: #2618
Depends on: #2588
## How
- **Dependency:** enables `datafusion-python = "55.0.0"` (`default-features
= false`) so we can reuse its `PythonLogicalCodec` / `PythonPhysicalCodec`.
- **Codecs (`src/codec.rs`):** `PyBallistaLogicalExtensionCodec` and
`PyBallistaPhysicalExtensionCodec` now wrap datafusion-python's Python codecs,
with Ballista's default codecs as the terminal codec. UDF and physical-expr
encode/decode methods are delegated. The terminal codec writes unframed
payloads, so plans without Python UDFs serialize exactly as before. A test
covers this (`test_non_python_udf_encodes_as_ballista`).
- **Scheduler / executor (`src/cluster.rs`):** `BallistaScheduler`,
`BallistaExecutor` and the test cluster (`setup_test_cluster`) now use the
Python-aware codecs (`override_logical_codec` / `override_physical_codec`, new
`ballista_codec()` helper), so they can decode Python UDFs.
- **FFI planner (`src/ffi_planner_workaround.rs`):**
`FileFormatFallbackCodec` detects inline Python UDF payloads (`DFPYUDF` /
`DFPYUDA` / `DFPYUDW` prefixes) and decodes them with a local
`PythonLogicalCodec`. If they went through the session's codec they would come
back as foreign UDFs, which can't be inlined again, so they would reach the
scheduler by name only.
## Known limitation
`BallistaPhysicalExtensionCodec` serializes shuffle partitioning expressions
with its internal default codec, so Python UDFs used as **shuffle partitioning
keys** are not inlined.
## Tests
- Removed `XFAIL_UDF` from `test_udf_dataframe` and `test_udf_sql`; both now
pass.
- Added a Rust test checking that non-Python UDFs still encode the same as
with Ballista's codecs.
- Added a UDF example to `examples/ballista_execution.py`.
## Notes for reviewers
- The `DFPY*` prefixes copy constants that are `pub(crate)` in
datafusion-python. They need to be kept in sync until they are exposed upstream.
- `Cargo.lock` grows a lot (reqwest, rustls, quinn, datafusion-spark, …).
These are transitive dependencies of `datafusion-python`.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
