chore(deps): upgrade to DataFusion 54 - #1906
Conversation
Bump the workspace to the published datafusion 54 (and arrow/arrow-flight 58.3, object_store 0.13.2, rustyline 18 in ballista-cli), replacing the branch-54 git-rev pin from the earlier draft. API migration to 54: - as_any removed from ExecutionPlan/DataSource/PhysicalExpr: downcast directly on the trait object (and upcast dyn ShuffleWriter to dyn ExecutionPlan); as_any retained where it still exists (Arrow arrays, UserDefinedLogicalNode, ExtensionOptions). - ExecutionPlan::partition_statistics returns Result<Arc<Statistics>>. - parse_protobuf_partitioning/_hash_partitioning take a PhysicalPlanDecodeContext; serialize_partitioning takes the codec plus the proto converter. - TaskContext::new and FunctionRegistry gain higher-order-function arguments (wired empty in Ballista). - BatchPartitioner::new_hash_partitioner is fallible. Test updates for 54 behavior changes: - approx_percentile_cont_with_weight result; shuffle-writer tests made robust to hash distribution; execution-graph/dot tests use Ballista session config so the small join stays Partitioned; regenerated AQE plan snapshots (projection folded into HashJoinExec).
|
There seems to be a major performance regression. This was caught by the |
|
The Parquet scans are reading too many rows. Still trying to find the root cause. |
DataFusion 54's DataSourceExec hands file groups to partition streams from a shared work-queue that is only divided across partitions when all partitions of one plan instance are polled concurrently. Ballista executes one partition per task on its own decoded plan instance, so a task that polls a single partition in isolation drains the whole queue and scans the entire table -- producing N-times-inflated results and N-times the work for every task dispatched after the first wave. Restrict each task's DataSourceExec to the file group for its partition_id before execution (other group slots emptied, partition count preserved) so the lone execute(partition_id) reads only its own group. Skip when partition_id is outside the source's file groups (an operator between the scan and the stage output changed the partition count). Closes apache#1907
…tion DataFusion 54 plans uncorrelated scalar subqueries as a physical ScalarSubqueryExec wrapping a ScalarSubqueryExpr that reads an in-process shared results container. That container cannot cross process or stage boundaries, and datafusion-proto can only deserialize the expr inside its surrounding exec, so when Ballista splits a plan into stages at shuffle boundaries the executor receives a bare ScalarSubqueryExpr and fails to decode it -- hanging queries such as TPC-H q11/q15/q22. Disable datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery in the Ballista restricted configuration so the optimizer rewrites uncorrelated scalar subqueries to joins, which Ballista distributes correctly. Verified via datafusion-cli that the rewrite produces results identical to the physical-execution path. Closes apache#1909
Add unit tests for the two DataFusion 54 distributed-execution fixes: - restrict_scan_to_partition keeps only the task's own file group, leaves the partition count intact, returns None for an out-of-range partition, and ignores non-file-backed plans. - the Ballista restricted configuration disables datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery so uncorrelated scalar subqueries are rewritten to joins.
|
Following up on my notes above (the perf regression / "parquet scans reading too many rows") — root-caused and fixed, and it turned out to be a correctness bug, not just performance. What was happening: DataFusion 54's Fix: restrict each task's While verifying the full suite I hit a second, independent DF54 regression: uncorrelated scalar subqueries (q11/q15/q22) failed to deserialize once the plan is split into stages, because DF54 plans them as a Both fixes are in this PR with unit tests, and the full 22-query TPC-H SF10 suite now runs end to end with correct results and no regression. I've updated the PR description accordingly. One related robustness issue I noticed but left out of this PR (it's pre-existing and orthogonal): a deterministic task decode failure makes the scheduler mark the executor dead, which with a single executor hangs the job instead of failing the query. Filed separately as #1908. |
| .set_bool( | ||
| "datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery", | ||
| false, | ||
| ) |
There was a problem hiding this comment.
I filed follow on issue #1910 to optimize this
martin-g
left a comment
martin-g
left a comment
There was a problem hiding this comment.
Should https://github.com/apache/datafusion-ballista/pull/1863/changes#diff-a52fcf388ba654263ab5f17dbf84452ffb97bda2c0a082fc347bb5ffe6c55591R422 be reverted with this PR too ? Or in a follow-up ?
| } | ||
|
|
||
| fn higher_order_function_names(&self) -> HashSet<String> { | ||
| HashSet::new() |
There was a problem hiding this comment.
This will break the built-in HOFs from DataFusion.
IMO it would be better to copy them from the state at
datafusion-ballista/ballista/core/src/registry.rs
Lines 154 to 166 in 85805db
| ) -> Option<Arc<dyn ExecutionPlan>> { | ||
| let exec = plan.downcast_ref::<DataSourceExec>()?; | ||
| let source: &dyn Any = exec.data_source().as_ref(); | ||
| let config = source.downcast_ref::<FileScanConfig>()?; |
There was a problem hiding this comment.
If the config is different type, e.g. MemorySourceConfig then Ok(Transformed::no(p)) will be executed and it may return duplicate results, no ?
Maybe log a WARN here before returning None ?!
| ) | ||
| })?; | ||
| let converter = DefaultPhysicalProtoConverter {}; | ||
| let decode_ctx = PhysicalPlanDecodeContext::new(ctx, self.default_codec.as_ref()); |
There was a problem hiding this comment.
| let decode_ctx = PhysicalPlanDecodeContext::new(ctx, self.default_codec.as_ref()); | |
| let decode_ctx = PhysicalPlanDecodeContext::new(ctx, self); |
DataFusion 54 ships built-in higher-order functions (array_any_match, array_filter, array_transform), but BallistaFunctionRegistry returned an empty set and every TaskContext the executor built passed Default::default() for higher-order functions, so plans using those functions failed to decode on the executor. Carry higher-order functions on BallistaFunctionRegistry (populated from the session state and the defaults) and thread them through get_task_definition, get_task_definition_vec, and all TaskContext construction sites.
restrict_scan_to_partition only restricts FileScanConfig, the source that hands out file groups from a shared work-queue. Other sources (e.g. MemorySourceConfig) isolate partitions in their own open(partition) and need no restriction, so they are left unchanged silently. Warn for any source type that is neither, so a future source that distributes work across partitions and could over-read in a single-partition task is surfaced rather than failing silently.
The physical plan decode context was given default_codec, which cannot decode Ballista extension nodes nested inside an expression. Pass the Ballista codec itself, matching DataFusion's convention of threading the full extension codec; it delegates to default_codec for non-Ballista nodes, so it is a strict superset.
datafusion-proto 53.1.0 could not distinguish a FilterExec with an empty projection (Some(vec![])) from a full projection (None), so the planner rewrote empty-projection filters into a ProjectionExec-wrapped equivalent before serialization (apache#1838). DataFusion 54 encodes None as the full column list and Some(vec![]) as an empty list and reconstructs None only when all columns are present, so the distinction round-trips correctly. Remove make_filter_projection_serde_safe and apply it nowhere; add a physical serde round-trip test asserting a zero-column FilterExec keeps its empty projection.
|
Thanks for the review @martin-g. I have addressed feedback so far. |
Reconcile this branch's DataFusion 54 work with the canonical DF54 upgrade (apache#1906) that landed upstream: - Keep enable_physical_uncorrelated_scalar_subquery at DataFusion's default (enabled); this branch materializes uncorrelated scalar subqueries in the client planner instead of disabling the option. - Adopt upstream's higher-order function support in the function registry and TaskContext construction. - Adopt upstream's recursive PhysicalPlanDecodeContext codec. - Adopt upstream's more robust restrict_scan_to_partition that handles MemorySourceConfig and warns on unrecognized DataSourceExec sources. - Drop the make_filter_projection_serde_safe workaround (no longer needed on DF54; upstream added regression coverage in serde).
2 participants
Which issue does this PR close?
Closes #1776. Supersedes the earlier draft #1777 (which pinned a
branch-54git rev); this is a fresh upgrade off the latestmainagainst the publisheddatafusion = "54".Also fixes two DataFusion 54 distributed-execution regressions surfaced by this upgrade: Closes #1907 and Closes #1909.
Rationale for this change
DataFusion 54.0.0 is now released on crates.io, so Ballista can track it with a normal version dependency instead of a git pin. Re-doing the upgrade fresh on current
main(rather than rebasing the stale draft) keeps the API migration consistent with all the code that has landed since, and validates against the released crate (whose API differs from the pre-releasebranch-54commit the draft targeted — e.g.parse_protobuf_partitioningnow takes aPhysicalPlanDecodeContextrather than the extra argument the draft removed).The upgrade also surfaced two distributed-execution regressions (caught by the
TPC-H SF10CI job) that are fixed here so the upgrade is correct and performant end to end — see "Distributed-execution fixes" below.What changes are included in this PR?
datafusion/datafusion-*54,arrow/arrow-flight58.3,object_store0.13.2, andrustyline18inballista-cli.as_anywas removed fromExecutionPlan/DataSource/PhysicalExpr; downcast directly on the trait object, and upcastdyn ShuffleWritertodyn ExecutionPlanbefore downcasting.as_anyis retained where it still exists (Arrow arrays,UserDefinedLogicalNode,ExtensionOptions).ExecutionPlan::partition_statisticsnow returnsResult<Arc<Statistics>>.parse_protobuf_partitioning/parse_protobuf_hash_partitioningtake aPhysicalPlanDecodeContext;serialize_partitioningtakes the codec plus the proto converter.TaskContext::newandFunctionRegistrygained higher-order-function parameters/methods, wired with empty defaults in Ballista.BatchPartitioner::new_hash_partitioneris now fallible.approx_percentile_cont_with_weightresult value.hash_join_single_partition_threshold = 0) so the small test join keeps itsPartitionedplan, matching production.ProjectionExecintoHashJoinExecas aprojection=field).Distributed-execution fixes
DataSourceExechands file groups to partition streams from a shared work-queue that is only divided across partitions when all partitions of one plan instance are polled concurrently. Ballista executes one partition per task on its own decoded plan instance, so a task polling a single partition in isolation drained the whole queue and scanned the entire table — producing N-times-inflated aggregates and N-times the work for every task after the first wave. The executor now restricts each task'sDataSourceExecto the file group for itspartition_idbefore execution.ScalarSubqueryExecwhoseScalarSubqueryExprcan only be deserialized inside that exec; Ballista's stage splitting separated them, so the executor could not decode the stage plan (TPC-H q11/q15/q22). The Ballista restricted configuration now disablesdatafusion.optimizer.enable_physical_uncorrelated_scalar_subquery, so the optimizer rewrites these to joins, which Ballista distributes correctly.Are there any user-facing changes?
Ballista now runs on DataFusion 54. No Ballista API or SQL semantics change.
Verified:
cargo build --workspace --all-targets,cargo clippy --all-targets --workspace -- -D warnings, andcargo test --workspaceall clean (including new unit tests for both distributed-execution fixes). The full 22-query TPC-H SF10 suite runs end to end on a distributed cluster with correct results and no regression.