jiayuasu opened a new issue, #26862:
URL: https://github.com/apache/datafusion/issues/26862
### Describe the bug
A `DataFrame` sorted with `sort()` loses its order when a later `select()`
computes a column, once the sort's output has more than one batch. `select()`
of plain columns keeps the order.
The physical plan deals the sorted batches round-robin to
`target_partitions` partitions so the projection runs in parallel, and nothing
merges them back by the sort key:
```
ProjectionExec: expr=[a@0 as a, a@0 + 1 as t.a + Int32(1)]
RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1,
maintains_sort_order=true
SortExec: expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false]
DataSourceExec: partitions=1, partition_sizes=[1]
```
Each partition is sorted (`maintains_sort_order=true`), but the plan has 4
output partitions where the sort alone has 1, and `collect()` gets the batches
in whatever order the partitions produce them. So the order is lost by a race,
which 55.2.0 lost in every run below and main in some.
### To Reproduce
```rust
use std::sync::Arc;
use datafusion::arrow::array::{AsArray, Int64Array, RecordBatch};
use datafusion::arrow::datatypes::Int64Type;
use datafusion::error::Result;
use datafusion::physical_plan::ExecutionPlanProperties;
use datafusion::prelude::*;
#[tokio::main]
async fn main() -> Result<()> {
let ctx =
SessionContext::new_with_config(SessionConfig::new().with_target_partitions(4));
let a = Arc::new(Int64Array::from_iter_values(0..100_000));
ctx.register_batch("t", RecordBatch::try_from_iter([("a", a as _)])?)?;
let sorted = ctx
.table("t")
.await?
.sort(vec![col("a").sort(false, false)])?;
let expected: Vec<i64> = (0..100_000).rev().collect();
for (name, df) in [
("sort", sorted.clone()),
("sort, select a", sorted.clone().select(vec![col("a")])?),
(
"sort, select a, a + 1",
sorted.clone().select(vec![col("a"), col("a") + lit(1)])?,
),
] {
let mut in_order = 0;
for _ in 0..10 {
let a: Vec<i64> = (df.clone().collect().await?.iter())
.flat_map(|b|
b.column(0).as_primitive::<Int64Type>().values().to_vec())
.collect();
in_order += (a == expected) as usize;
}
let partitions = df
.clone()
.create_physical_plan()
.await?
.output_partitioning()
.partition_count();
println!("{name}: {partitions} output partition(s), in order in
{in_order} of 10 runs");
}
sorted
.select(vec![col("a"), col("a") + lit(1)])?
.explain(false, false)?
.show()
.await?;
Ok(())
}
```
With `datafusion = { version = "55.2.0", default-features = false }` and
`tokio`:
```
sort: 1 output partition(s), in order in 10 of 10 runs
sort, select a: 1 output partition(s), in order in 10 of 10 runs
sort, select a, a + 1: 4 output partition(s), in order in 0 of 10 runs
```
On main (56eca4a7c) the plan is the same, with 4 output partitions, but the
rows came out of order less often: in three runs of the program, in order in 9,
10 and 10 of 10 runs; with 1,000,000 rows and 8 partitions, in 7, 10 and 7 of
10.
With 8,192 rows or fewer (the default batch size), the sort gives a single
batch, so there is nothing to reorder.
### Expected behavior
The rows of `sort(...).select(...)` come in the sort's order, as they do for
`select(vec![col("a")])` and for the SQL `SELECT a, a + 1 FROM t ORDER BY a
DESC`: a `SortPreservingMergeExec` above the parallel projection, or no
repartition there.
### Additional context
The physical optimizer already treats such a sort as the order of the final
result. `OutputRequirements` (`require_top_ordering_helper` in
`datafusion/physical-optimizer/src/output_requirements.rs`) looks down from the
root through operators that maintain their input's order, which
`ProjectionExec` does:
```rust
// Global ordering defines desired ordering in the final result.
...
// Keep searching for a `SortExec` / `SortPreservingMergeExec` as long as
// ordering is maintained, and on-the-way operators do not themselves
// require an ordering.
```
But it wraps the `SortExec` it finds, below the projection, in the
`OutputRequirementExec`. Above that, the requirement no longer applies, so the
projection, which benefits from input partitioning when its expressions are not
all columns, gets a round-robin repartition, and no merge restores the order at
the root.
In SQL a sort in a subquery promises nothing about the outer query, but the
DataFrame API builds exactly this plan for the common chain of sorting and then
adding or transforming columns. Whether the result is ordered then depends on
whether a selected column is computed and on how many rows there are, so small
tests pass and the same code on larger data does not.
We found this in Apache SedonaDB (apache/sedona-db#1532), where a sorted
frame written to a file after computing its columns came out of order.
### Willingness to contribute
I plan to contribute a fix for this bug independently
--
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]
