feat: use datafusion python FFI for integration (native datafusion python support)

[milenkovicm]

Which issue does this PR close?

import ballista

from ballista import BallistaSessionContext

import datafusion
# BallistaSessionContext replaces SessionContext
# from datafusion import SessionContext
from datafusion import DataFrame

# replace datafusion SessionContext
# ctx = SessionContext()
# with BallistaSessionContext
ctx = BallistaSessionContext(address=f"df://localhost:50050")

# to run query distributed ob ballista
ctx.sql("SELECT 1 as r").show()

Closes #.

Rationale for this change

• update to datafusion python 55
• use FFI query planner released with datafusion python 55, to make ballista integration just a tin wrapper on top of datfusion-puthon

What changes are included in this PR?

• removed DistributedDataFrame and BallistaContext
• removed redefining meta classes for redefinition of DistributedDataFrame and BallistaContext
• removed methods which are tied to DistributedDataFrame and BallistaContext, i would argue they should not be there in the first place
• use FFI datafusion python query planner

Are there any user-facing changes?

yes api changed, some methods and classes disappeared, as python integration was experimental users should have expected

this PR builds on @timsaucer's PR #2252

Follow up

• support for distributed python UDFs

[milenkovicm]

hey @timsaucer this is follow up on your #2252 after datafusion python 55 released,

uv run --no-project pytest

FAILED python/tests/test_context.py::test_write_csv - Exception: DataFusion error: FFI error: This feature is not implemented: FFI does not support decode_file_format
FAILED python/tests/test_context.py::test_write_parquet - Exception: DataFusion error: FFI error: This feature is not implemented: FFI does not support decode_file_format
FAILED python/tests/test_context.py::test_write_json - Exception: DataFusion error: FFI error: This feature is not implemented: FFI does not support decode_file_format

I have removed good part of ballista specific code, happy day scenario works, it connects to ballista and runs the query, but write_ methods are failing, and i'm not sure if there is something else needed.

any ideas what could be done ?

cc @andygrove another pair of eyes would help (thanks)

[milenkovicm]

there is comment apache/datafusion-python#1679 (comment) but it makes no sense to me, perhaps @timsaucer may have better idea what claude wants to say :D

[andygrove]

Thanks for pushing this through @milenkovicm, and @timsaucer for the groundwork in #2252! Moving to a real DataFusion SessionContext is a big improvement. Every DataFrame method now runs on the cluster, where before only the few the metaclass wrapped did.

I built this branch and main locally and ran a few probes against setup_test_cluster(). I left comments on a config regression, EXPLAIN ANALYZE and some API cleanup. I'm requesting changes mainly for the config issue, since it makes queries fail that work on main today.

Since this removes DistributedDataFrame, turns BallistaSessionContext into a function and changes what config does, could we also add the api-change label and a Python client section to docs/source/upgrading/56.0.0.md?

[andygrove]

Moving the serde into PyBallistaPhysicalExtensionCodec makes sense given your earlier note. Would it make sense to move the proto message there too? Right now core's wire format has a node that its own codec rejects, and a library user who hits this arm gets an internal error with no hint about where the node comes from. If it needs to stay in core's oneof so the Python codec can tell it apart from Ballista nodes, a short comment on the message and on this arm saying only the Python client produces it would help. Keeping it out of core would also avoid adding a PhysicalPlanType variant, though tags 13 to 15 went in the same way.

[milenkovicm]

I suggest keeping this here, with the changed exception saying this is intended to be used in python only. I agree its not ideal but alternative is to have tonic infrastructure in python project. wdyt ?

[andygrove]

I think ballista.client.use_tls and ballista.client.pull stop working in this path. DistributedQueryExec::execute reads both from a native BallistaConfig on the task context. datafusion-ffi rebuilds that context on the Ballista side with only a generic FFI_ExtensionOptions, so both fall back to their defaults. The docs tell clients to set ballista.client.use_tls for TLS result fetches, and on main it works from cluster_config. I haven't tried this against a TLS cluster, so please double check. Could execute() fall back to self.config when the task context has no BallistaConfig? Or could the Python side rebuild one from the datafusion_ffi options?

[milenkovicm]

i can't reproduce this issue

ctx.sql("SET ballista.job.name='super job'").show()
ctx.sql("SET ballista.shuffle.force_remote_read=true").show()

ctx.sql("select id, count(*) from t group by id").show()

manually checking scheduler i see them set

same for

ctx = BallistaSessionContext(address=f"df://{host}:{port}", cluster_config={
    "ballista.job.name": "super job",
    "ballista.shuffle.force_remote_read": True
})

can you please re-check when u get chance

[milenkovicm]

thanks @andygrove comments make sense, will follow up over the weekend, converting to PR draft until then

[milenkovicm]

please have a look when you get chance @andygrove