Sanity check failed when sort and aggregate on a multi-partitioned table
- 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
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