NoahKusaba opened a new pull request, #2852: URL: https://github.com/apache/datafusion-ballista/pull/2852
# Which issue does this PR close? Follows up #2679 / #2721. Found through andygrove's `UNION ALL` repro on #2638 (https://github.com/apache/datafusion-ballista/pull/2638#discussion_r4505755569). # Rationale for this change When the scheduler splits a stage between tasks, a `UnionExec` gives each branch the partitions the task owns, so a task that owns only one branch restricts the others to no partitions (`[]`). Since #2721, a leaf the scheduler can't restrict fails the job unless the task runs all of it. That includes a branch the task runs none of, such as one holding a `DataSourceExec` over a custom source like Iceberg's, so `SELECT … FROM t UNION ALL SELECT … FROM t` fails whenever its stage is split. Before #2721 it returned duplicate rows. A task never executes a leaf it runs none of the partitions of, so the scheduler can replace any such leaf with an empty one, whatever it is. That needs no resolver, so it also works in standalone mode, which can't take one. # What changes are included in this PR? - **A leaf nothing can restrict becomes an empty leaf when the task runs none of it.** It stays an error for any other partial slice, which needs a resolver. The empty leaf is a zero-partition `MemorySourceConfig`, not an `EmptyExec`: DataFusion's proto codec decodes an `EmptyExec` with zero partitions as one, so it would reach the executor with a partition the task doesn't own. The existing `EmptyExec` / `PlaceholderRowExec` case already worked around this; both now use one helper, `no_partitions`. - **A `UnionExec` branch that keeps another partition count after restriction is refused**, instead of making the union run the wrong branch. This catches a branch that collapses its input, such as through the `SortExec` an `ORDER BY` puts on each branch, which reports one partition whatever its leaves do, and a node of your own that keeps its own count. - **The static planner keeps more union branches inline.** It decides whether a branch needs a stage of its own by restricting it to `[]` (`restricts_to_nothing`). A branch over an `UnresolvedShuffleExec`, which becomes a restrictable `ShuffleReaderExec` at task time, used to fail that trial and got an extra stage; it now stays in the union's stage. `losing_own_shuffle_output_restores_the_cache_reader` hardcoded the old stage numbering and is updated, still testing the same recovery. # Are these changes tested? Yes: - A leaf nobody can restrict becomes an empty in-memory source with its schema for `[]`, and is still refused for other partial slices. - A union branch the task doesn't own restricts to nothing without a resolver. - A restricted union keeps its empty branch at zero partitions through `physical_plan_to_bytes` / `physical_plan_from_bytes`. With an `EmptyExec` instead, this test fails: the decoded union has 2 partitions instead of 1. - A union branch whose node keeps its own partition count is refused, with the full error checked. - A union branch over a leaf the scheduler doesn't know stays inline in the static planner. `ballista-scheduler` unit tests (630) and all `ballista` client tests pass; clippy is clean. End to end, standalone Ballista with `ballista.scheduler.max_partitions_per_task = 1`, through `ballista-iceberg` (#2638) temporarily pointed at this branch, over two-row tables: | Query | this PR | without the union check | |---|---|---| | Iceberg `UNION ALL` | 4 rows ✓ | 4 rows ✓ | | Iceberg `UNION ALL … ORDER BY id` | fails: the branch's `SortExec` still reports 1 partition | 8 rows, should be 4 | | Parquet `UNION ALL` | 4 rows ✓ | 4 rows ✓ | | Parquet `UNION ALL … ORDER BY id` | fails, same error | 8 rows, should be 4 | On main, the Iceberg queries fail with `cannot restrict … to partitions []` (#2721), and the Parquet ones behave as in the last column, since this PR doesn't change how a Parquet scan is restricted. The `ORDER BY` case is a separate bug in the adaptive planner, for Parquet too: the static planner gives a union branch that collapses its input a stage of its own (`can_stay_inline`, the fix for #2184), and the adaptive planner doesn't. This PR turns its wrong results into an error; I'll open an issue for the fix. # Are there any user-facing changes? - A split stage with a union branch over a custom source now succeeds instead of failing. - A split stage whose union branch collapses its input under the adaptive planner, such as `UNION ALL … ORDER BY`, now fails instead of returning duplicate rows. - Some static plans with a union have one stage fewer. The 56.0.0 upgrade guide and the resolver notes in `extending-components.md` are updated. -- 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]
