andygrove opened a new issue, #2817:
URL: https://github.com/apache/datafusion-ballista/issues/2817
### Describe the bug
The Flight SQL frontend reports a query's result schema from the job's final
stage. `result_schema` in `ballista/scheduler/src/flight_sql.rs` finds that
stage as the first stage with no `output_links`:
```rust
graph
.stages()
.values()
.find(|stage| stage.output_links().is_empty())
.map(|stage| stage.plan().schema())
```
The adaptive graph creates every stage with `output_links` empty and never
fills them in (see #1996). So under the adaptive planner, which is the default,
every stage matches. `stages()` is a `HashMap`, so which stage's schema comes
back depends on iteration order. For a multi-stage query it can be an
intermediate stage's schema, such as the partial aggregate of a `GROUP BY`.
A Flight SQL client (ADBC, JDBC, BI tools) can then get a schema in
`FlightInfo` that doesn't match the batches it reads from the result locations.
### To Reproduce
Run a multi-stage query through the Flight SQL frontend with the adaptive
planner on, for example `SELECT c, min(a) FROM t GROUP BY c`. Compare the
schema in the returned `FlightInfo` with the schema of the fetched batches. The
result depends on `HashMap` order, so it may take a few runs or a deeper plan
to see a mismatch.
### Expected behavior
The schema comes from the stage that wrote the job's output.
### Additional context
The caller already has the output locations: `partition_location` for staged
jobs and `pipelined_output_locations` for pipelined ones. Each
`PartitionLocation` carries the id of the stage that wrote it, so the final
stage could be found from those. #2753 does the same for
`intermediate_stage_ids()`. Another option is an `ExecutionGraph` method that
returns the final stage ids, which `intermediate_stage_ids()` and
`result_schema` could share. This only affects the reported schema, not the
data that is returned.
--
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]
