niebayes commented on PR #21107:
URL: https://github.com/apache/datafusion/pull/21107#issuecomment-5016393397
@xudong963 @alamb I found this PR also fixed a bug for DataFusion <54.0.0
version.
For a SQL `SELECT DISTINCT host FROM t ORDER BY host LIMIT k`, if the Scan
operator declares the output data is sorted by host already, the SortExec
(TopK) operator would be pushed down under the AggregateExec(Final) operator,
makes the result incorrect.
I have rewritten a bug reproduction:
```rust
//! Reproduces a correctness bug where `SortExec: TopK(fetch=N)` is
//! inserted between the Partial and Final stages of a grouped sorted
//! aggregate, after a hash repartition destroys the sort order.
//!
//! The intermediate TopK counts ROWS, but its input stream still holds
//! duplicate GROUP keys (one partial group row per round-robin partition
//! that saw the key). Limiting ROWS is not limiting GROUPS — duplicates
//! of small keys can evict keys that belong in the global top-N.
//!
//! This test constructs the broken plan shape that Datalayers' optimizer
//! chain produces and verifies that DISTINCT + ORDER BY + LIMIT drops
//! distinct values.
use std::sync::Arc;
use arrow::array::{AsArray, StringArray, UInt64Array};
use arrow::compute::concat_batches;
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use datafusion::catalog::MemTable;
use datafusion::common::Result;
use datafusion::execution::context::SessionContext;
use datafusion::prelude::{SessionConfig, col};
// ── data-loss reproducer ──────────────────────────────────────────
//
// When the plan contains a SortExec:TopK(fetch=N) between
// AggregateExec:FinalPartitioned and RepartitionExec:Hash, the TopK
// counts partial-group **rows** instead of distinct **groups**.
// Because the stream after Hash repartition holds one partial-group
// row per source partition that saw the key, the duplicates of small
// keys can consume the entire fetch budget before larger keys are
// reached. A key that belongs in the global top-N may then be
// entirely absent from the output.
//
// The test below creates many distinct keys and a fetch small enough
// that the intermediate TopK necessarily drops group keys.
//
// Parameters:
// source/hash partitions = 2 (TARGET_PARTITIONS)
// distinct_keys = 8 (h00 -- h07)
// fetch = 2 (LIMIT)
//
// PartialAgg(lim=2) outputs the two smallest keys per source
// partition: h00 and h01. After hash repartition those 2 keys land
// in 2 partitions. If they collide on one partition:
// 2 keys × 2 partial rows each = 4 rows
// TopK(fetch=2) keeps h00×2 → h01 is entirely evicted.
// The query result then contains only h00 instead of the expected
// h00, h01.
const N_PARTITIONS: usize = 2;
const N_KEYS: usize = 8;
const FETCH: usize = 2;
fn bug_data() -> (Arc<Schema>, Vec<Vec<RecordBatch>>) {
let schema = Arc::new(Schema::new(vec![
Field::new("host", DataType::Utf8, false),
Field::new("row", DataType::UInt64, false),
]));
let mut partitions: Vec<Vec<RecordBatch>> = vec![vec![]; N_PARTITIONS];
let mut row = 0u64;
for key in 0..N_KEYS {
for _ in 0..N_PARTITIONS {
let host =
Arc::new(StringArray::from(vec![format!("h{key:02}")]));
let row_arr = Arc::new(UInt64Array::from(vec![row]));
let batch =
RecordBatch::try_new(schema.clone(), vec![host,
row_arr]).unwrap();
partitions[row as usize % N_PARTITIONS].push(batch);
row += 1;
}
}
for (i, p) in partitions.iter().enumerate() {
let merged = concat_batches(&schema, p).unwrap();
let values: Vec<_> = merged
.column_by_name("host")
.unwrap()
.as_string::<i32>()
.iter()
.map(|v| v.unwrap())
.collect();
assert!(values.is_sorted(), "partition {i} not sorted: {values:?}");
}
(schema, partitions)
}
#[tokio::test]
async fn distinct_topk_data_loss() -> Result<()> {
let config = SessionConfig::new().with_target_partitions(N_PARTITIONS);
let ctx = SessionContext::new_with_config(config);
let (schema, partitions) = bug_data();
let table = MemTable::try_new(schema, partitions)?.with_sort_order(vec![
vec![
col("host").sort(true, false)
];
N_PARTITIONS
]);
ctx.register_table("t", Arc::new(table))?;
let sql = format!("SELECT DISTINCT host FROM t ORDER BY host LIMIT
{FETCH}");
let df = ctx.sql(&sql).await?;
// ── Assert on data ───────────────────────────────────────────
let results = df.collect().await?;
let hosts: Vec<String> = results
.iter()
.flat_map(|b| {
b.column(0)
.as_string::<i32>()
.iter()
.map(|o| o.unwrap().to_string())
.collect::<Vec<_>>()
})
.collect();
let expected: Vec<String> = (0..FETCH).map(|i|
format!("h{i:02}")).collect();
assert_eq!(
hosts, expected,
"\nTopK counted rows instead of groups.\n\
got {hosts:?}\n\
want {expected:?}\n\
(partitions={N_PARTITIONS}, keys={N_KEYS}, fetch={FETCH})"
);
Ok(())
}
```
The test fails on v53.1.0 and passes on v54.0.0. Users using earlier
versions should upgrade.
--
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]
