From e97ca71ad8673517d84fb9ab8cca5dbb4c537dea Mon Sep 17 00:00:00 2001
From: SteNicholas <programgeek@163.com>
Date: Thu, 30 Jan 2025 17:22:28 +0800
Subject: [PATCH] [BLAZE-805] Support statistics of ExecutionPlan for SortExec

---
 .../datafusion-ext-plans/src/sort_exec.rs          | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/native-engine/datafusion-ext-plans/src/sort_exec.rs b/native-engine/datafusion-ext-plans/src/sort_exec.rs
index 2c397fb7..7711cef9 100644
--- a/native-engine/datafusion-ext-plans/src/sort_exec.rs
+++ b/native-engine/datafusion-ext-plans/src/sort_exec.rs
@@ -185,7 +185,7 @@ impl ExecutionPlan for SortExec {
     }
 
     fn statistics(&self) -> Result<Statistics> {
-        todo!()
+        Statistics::with_fetch(self.input.statistics()?, self.schema(), self.fetch, 0, 1)
     }
 }
 
@@ -1393,9 +1393,9 @@ mod fuzztest {
         record_batch::RecordBatch,
     };
     use datafusion::{
-        common::Result,
+        common::{stats::Precision, Result},
         physical_expr::{expressions::Column, PhysicalSortExpr},
-        physical_plan::memory::MemoryExec,
+        physical_plan::{memory::MemoryExec, ExecutionPlan},
         prelude::{SessionConfig, SessionContext},
     };
 
@@ -1460,8 +1460,9 @@ mod fuzztest {
             None,
         )?);
         let sort = Arc::new(SortExec::new(input, sort_exprs.clone(), None));
-        let output = datafusion::physical_plan::collect(sort, task_ctx.clone()).await?;
+        let output = datafusion::physical_plan::collect(sort.clone(), task_ctx.clone()).await?;
         let a = concat_batches(&schema, &output)?;
+        let a_row_count = sort.clone().statistics()?.num_rows;
 
         let input = Arc::new(MemoryExec::try_new(
             &[batches.clone()],
@@ -1472,10 +1473,13 @@ mod fuzztest {
             sort_exprs.clone(),
             input,
         ));
-        let output = datafusion::physical_plan::collect(sort, task_ctx.clone()).await?;
+        let output = datafusion::physical_plan::collect(sort.clone(), task_ctx.clone()).await?;
         let b = concat_batches(&schema, &output)?;
+        let b_row_count = sort.clone().statistics()?.num_rows;
 
         assert_eq!(a.num_rows(), b.num_rows());
+        assert_eq!(a_row_count, b_row_count);
+        assert_eq!(a_row_count, Precision::Exact(n));
         assert!(a == b);
         Ok(())
     }