feat(python): integrate Ballista through DataFusion FFI - #2252
Draft
timsaucer wants to merge 13 commits into
Draft
feat(python): integrate Ballista through DataFusion FFI#2252timsaucer wants to merge 13 commits into
timsaucer wants to merge 13 commits into
Conversation
Pin the datafusion* crates to a git rev of apache/datafusion main instead of the 54 crates.io release so Ballista tracks unreleased DataFusion and surfaces breaking changes between releases. Bump arrow/arrow-flight to 59 to match. Port to the DataFusion main APIs: - TableReference moved from datafusion::sql to datafusion::common. - PhysicalExtensionCodec::try_decode/try_encode gained a PhysicalProtoConverterExtension argument. - MetricValue gained a PeakMemoryUsage variant; serialize it as a named gauge. - Partitioning gained a Range variant; format it in the DOT graph. - EnforceSorting and EnforceDistribution were unified into EnsureRequirements; use it for post-stage-split re-optimization. - ExecutionPlan::partition_statistics is deprecated in favour of statistics_with_args(&StatisticsArgs); migrate all call sites and operator overrides. This is a correctness fix: DataFusion operators that now override only statistics_with_args return unknown stats through the deprecated method, which would silently disable Ballista's broadcast-threshold checks. - ListingOptions dropped its target_partitions/collect_stat fields; build it via the builder (target_partitions comes from the session config). - Distribution::HashPartitioned is deprecated; use KeyPartitioned. Disable dynamic filter pushdown in the AQE mock session config to match the Ballista production default so join-input swaps during re-optimization remain legal. Refresh the EXPLAIN ANALYZE and AQE plan snapshots for DataFusion's dropped peak_mem_used metric and its new RightSemi join lowering.
…on rev
Merged latest apache/main into the branch and bumped the datafusion* git
dependencies from b79d3965 to efa84e83 (current apache/datafusion main HEAD).
That bump forced a re-port of the physical-plan and proto-codec APIs:
- statistics_with_args + StatisticsArgs::compute_child_statistics were
refactored into StatisticsContext::compute + ExecutionPlan
::statistics_from_inputs(input_stats, args) +
ExecutionPlan::child_stats_requests(partition). All Ballista operator
overrides and call sites were migrated. Custom operators
(ChaosExec, ShuffleWriterExec, SortShuffleWriterExec, BufferExec,
RuntimeStatsExec, executor CollectExec) now declare a single
ChildStats::At(partition) request and return their sole child's
precomputed stats, restoring the correctness fix the original PR made
when statistics_with_args replaced partition_statistics.
- PhysicalExtensionCodec::try_decode/try_encode gained a
PhysicalProtoConverterExtension argument; all test call sites now pass
&DefaultPhysicalProtoConverter {}.
- Partitioning::Range is a new enum variant; the shuffle-writer
partition-mapping walker treats it as a K-space (like Hash), and the
task-builder shuffle-reader rewrite falls back to Unknown for it.
- ExecutionPlan::required_input_distribution is deprecated in favour of
input_distribution_requirements(); task_builder was migrated.
- apache/datafusion#23642 is fixed upstream, so the pin test that
asserted the bug's presence was inverted into a regression test that
the partition count survives round-trip. make_empty_exec_serde_safe
itself is kept as belt-and-suspenders for legacy stage plans decoded
from persisted state.
- Plan-structure test expectations were updated for the new physical
optimizer ordering (ProjectionExec now sits above SortExec in the
aggregate/window fan-out) and the new RoundRobinBatch(2) repartition
the optimizer inserts above the SortMergeJoin.
- context_checks.rs's EXPLAIN ANALYZE snapshot merged the new
ShuffleReaderExec fetch metrics from main with the peak_mem_used drop
from the datafusion bump.
Advance the pinned apache/datafusion main rev from efa84e83 to 70c26a06. Port to the APIs that changed in between: - QueryPlanner::create_physical_plan takes &dyn Session instead of &SessionState. - datafusion_substrait::serializer::deserialize_bytes is now synchronous and takes &[u8]. Sync the vendored datafusion_common.proto for the new max_in_list_size field.
AI Disclosure: This code was written in part by an AI agent.
AI Disclosure: This code was written in part by an AI agent.
AI Disclosure: This code was written in part by an AI agent.:
milenkovicm
reviewed
Aug 7, 2026
milenkovicm
left a comment
milenkovicm
left a comment
Contributor
There was a problem hiding this comment.
After many years! Thanks @timsaucer awesome as always!
| "address = f\"df://{host}:{port}\"\n", | ||
| "\n", | ||
| "ballista_config = ballista_datafusion_config_defaults()\n", | ||
| "ctx = SessionContext(SessionConfig(ballista_config)).with_extensions(\n", |
Contributor
There was a problem hiding this comment.
Can we use ballista context to hide next few lines and simplify session context creation ?
| ] | ||
|
|
||
|
|
||
| class RedefiningDataFrameMeta(type): |
Contributor
There was a problem hiding this comment.
Can we remove this once we get FFI planner ?
| session_id = args[0].session_id | ||
| df = func(*args, **kwargs) | ||
| return DistributedDataFrame(df, session_id, address) | ||
| return DistributedDataFrame( |
Contributor
There was a problem hiding this comment.
DistributedDataFrame, should be removed with FFI, should it ?
| name: &str, | ||
| buf: &[u8], | ||
| ) -> Result<Arc<datafusion::logical_expr::ScalarUDF>> { | ||
| self.default_codec.try_decode_udf(name, buf) |
Contributor
There was a problem hiding this comment.
Do we need to do something extra to get python udf serde?
martin-g
reviewed
Aug 10, 2026
| @@ -77,10 +82,25 @@ | |||
| "# ctx = BallistaSessionContext(\"df://your-scheduler:50050\")\n", | |||
Member
There was a problem hiding this comment.
This should be removed too, right ?
Suggested change
| "# ctx = BallistaSessionContext(\"df://your-scheduler:50050\")\n", |
An error occurred while trying to automatically change base from
datafusion-55
to
main
August 14, 2026 15:49
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
BallistaExtensionbundle that installs its components atomically withSessionContext.with_extensionsTesting
cargo checkcargo fmt -- --checkuv run pytest python/tests/test_context.py -q(21 passed)