Skip to content

Commit e97ca71

Browse files
committed
[BLAZE-805] Support statistics of ExecutionPlan for SortExec
1 parent f6452a7 commit e97ca71

File tree

1 file changed

+9
-5
lines changed

1 file changed

+9
-5
lines changed

native-engine/datafusion-ext-plans/src/sort_exec.rs

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -185,7 +185,7 @@ impl ExecutionPlan for SortExec {
185185
}
186186

187187
fn statistics(&self) -> Result<Statistics> {
188-
todo!()
188+
Statistics::with_fetch(self.input.statistics()?, self.schema(), self.fetch, 0, 1)
189189
}
190190
}
191191

@@ -1393,9 +1393,9 @@ mod fuzztest {
13931393
record_batch::RecordBatch,
13941394
};
13951395
use datafusion::{
1396-
common::Result,
1396+
common::{stats::Precision, Result},
13971397
physical_expr::{expressions::Column, PhysicalSortExpr},
1398-
physical_plan::memory::MemoryExec,
1398+
physical_plan::{memory::MemoryExec, ExecutionPlan},
13991399
prelude::{SessionConfig, SessionContext},
14001400
};
14011401

@@ -1460,8 +1460,9 @@ mod fuzztest {
14601460
None,
14611461
)?);
14621462
let sort = Arc::new(SortExec::new(input, sort_exprs.clone(), None));
1463-
let output = datafusion::physical_plan::collect(sort, task_ctx.clone()).await?;
1463+
let output = datafusion::physical_plan::collect(sort.clone(), task_ctx.clone()).await?;
14641464
let a = concat_batches(&schema, &output)?;
1465+
let a_row_count = sort.clone().statistics()?.num_rows;
14651466

14661467
let input = Arc::new(MemoryExec::try_new(
14671468
&[batches.clone()],
@@ -1472,10 +1473,13 @@ mod fuzztest {
14721473
sort_exprs.clone(),
14731474
input,
14741475
));
1475-
let output = datafusion::physical_plan::collect(sort, task_ctx.clone()).await?;
1476+
let output = datafusion::physical_plan::collect(sort.clone(), task_ctx.clone()).await?;
14761477
let b = concat_batches(&schema, &output)?;
1478+
let b_row_count = sort.clone().statistics()?.num_rows;
14771479

14781480
assert_eq!(a.num_rows(), b.num_rows());
1481+
assert_eq!(a_row_count, b_row_count);
1482+
assert_eq!(a_row_count, Precision::Exact(n));
14791483
assert!(a == b);
14801484
Ok(())
14811485
}

0 commit comments

Comments
 (0)