DataFrame sort() order is lost by a later select() of a computed column #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

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 (56eca4a) 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:

// 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