apache / apache/datafusion

Sanity check failed when sort and aggregate on a multi-partitioned table

Open
#18,989 2 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Rust
Stars
9.3k
Forks
2.4k
Avg merge
3d 7h
Merged PRs (30d)
344

Description

### Describe the bug

When creating a logical plan that involves a sequence of `Sort -> Aggregate -> Sort -> Aggregate`, and the input table is multi-partitioned, the physical planner fails during the `SanityCheckPlan` phase.

It appears that the optimizer fails to inject a necessary `RepartitionExec` node between the two Aggregate operations, leading to a distribution requirement mismatch panic.

### To Reproduce

Here is a minimal reproduction script (using datafusion and tokio crate):
```rust
use std::sync::Arc;

use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::datasource::MemTable;
use datafusion::functions_aggregate::count::count_udaf;
use datafusion::logical_expr::col;
use datafusion::prelude::*;

#[tokio::main]
async fn main() {
let ctx = SessionContext::default();

let schema = Arc::new(Schema::new(vec![
Field::new("ts", DataType::Int64, false),
Field::new("region", DataType::Utf8, false),
Field::new("value", DataType::Float64, false),
]));

// create an empty but multi-partitioned MemTable
let mem_table = MemTable::try_new(schema.clone(), vec![vec![], vec![]]).unwrap();
ctx.register_table("metrics", Arc::new(mem_table)).unwrap();

// aggregate and sort twice
let data_frame = ctx
.table("metrics")
.await
.unwrap()
.aggregate(
vec![col("region"), col("ts")],
vec![count_udaf().call(vec![col("value")])],
)
.unwrap()
.sort(vec![
col("region").sort(true, true),
col("ts").sort(true, true),
])
.unwrap()
.aggregate(
vec![col("ts")],
vec![count_udaf().call(vec![col("count(metrics.value)")])],
)
.unwrap()
.sort(vec![col("ts").sort(true, true)])
.unwrap();

println!(
"Logical Plan:\n{}",
data_frame.logical_plan().display_indent()
);

data_frame.show().await.unwrap();
}
```

### Expected behavior

The query should execute successfully, but it panics:
```text
Logical Plan:
Sort: metrics.ts ASC NULLS FIRST
Aggregate: groupBy=[[metrics.ts]], aggr=[[count(count(metrics.value))]]
Sort: metrics.region ASC NULLS FIRST, metrics.ts ASC NULLS FIRST
Aggregate: groupBy=[[metrics.region, metrics.ts]], aggr=[[count(metrics.value)]]
TableScan: metrics

thread 'main' panicked at src/main.rs:51:29:
called `Result::unwrap()` on an `Err` value: Context("SanityCheckPlan", Plan("Plan: [\"AggregateExec: mode=SinglePartitioned, gby=[ts@0 as ts], aggr=[count(count(metrics.value))]\", \" ProjectionExec: expr=[ts@1 as ts, count(metrics.value)@2 as count(metrics.value)]\", \" AggregateExec: mode=FinalPartitioned, gby=[region@0 as region, ts@1 as ts], aggr=[count(metrics.value)]\", \" CoalesceBatchesExec: target_batch_size=8192\", \" RepartitionExec: partitioning=Hash([region@0, ts@1], 64), input_partitions=2\", \" AggregateExec: mode=Partial, gby=[region@1 as region, ts@0 as ts], aggr=[count(metrics.value)]\", \" DataSourceExec: partitions=2, partition_sizes=[0, 0]\"] does not satisfy distribution requirements: HashPartitioned[[ts@0]]). Child-0 output partitioning: Hash([region@0, ts@0], 64)"))
```

### Additional context

If the input table is single-partitioned, or any sort operation is missing, the query executes successfully.

Contributor guide

Open the contributing guide

Research direction

Start by running the provided Rust reproduction with the multi-partitioned MemTable and inspect the physical plan at the SanityCheckPlan failure. Then trace how the physical planner enforces distribution requirements between the two Aggregate operations. Done means this query executes successfully with the sort operations and multiple input partitions.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
databases
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.