Refine the statistics estimation for the limit and aggregate operator

Author: yangzhong

datafusion/core/src/physical_plan/aggregates/mod.rs

@@ -505,7 +505,11 @@ impl ExecutionPlan for AggregateExec {
                     ..Default::default()
                 }
             }
-            _ => Statistics::default(),
+            _ => Statistics {
+                num_rows: self.input.statistics().num_rows,
+                is_exact: false,
+                ..Default::default()
+            },
         }
     }
 }

datafusion/core/src/physical_plan/limit.rs

@@ -194,6 +194,16 @@ impl ExecutionPlan for GlobalLimitExec {
     fn statistics(&self) -> Statistics {
         let input_stats = self.input.statistics();
         let skip = self.skip;
+        let max_row_num = self
+            .fetch
+            .map(|fetch| {
+                if fetch >= usize::MAX - skip {
+                    usize::MAX
+                } else {
+                    fetch + skip
+                }
+            })
+            .unwrap_or(usize::MAX);
         match input_stats {
             Statistics {
                 num_rows: Some(nr), ..
@@ -205,22 +215,25 @@ impl ExecutionPlan for GlobalLimitExec {
                         is_exact: input_stats.is_exact,
                         ..Default::default()
                     }
-                } else if nr - skip <= self.fetch.unwrap_or(usize::MAX) {
+                } else if nr <= max_row_num {
                     // if the input does not reach the "fetch" globally, return input stats
                     input_stats
-                } else if nr - skip > self.fetch.unwrap_or(usize::MAX) {
+                } else {
                     // if the input is greater than the "fetch", the num_row will be the "fetch",
                     // but we won't be able to predict the other statistics
                     Statistics {
-                        num_rows: self.fetch,
+                        num_rows: Some(max_row_num),
                         is_exact: input_stats.is_exact,
                         ..Default::default()
                     }
-                } else {
-                    Statistics::default()
                 }
             }
-            _ => Statistics::default(),
+            _ => Statistics {
+                // the result output row number will always be no greater than the limit number
+                num_rows: Some(max_row_num),
+                is_exact: false,
+                ..Default::default()
+            },
         }
     }
 }
@@ -358,8 +371,12 @@ impl ExecutionPlan for LocalLimitExec {
                 is_exact: input_stats.is_exact,
                 ..Default::default()
             },
-            // if we don't know the input size, we can't predict the limit's behaviour
-            _ => Statistics::default(),
+            _ => Statistics {
+                // the result output row number will always be no greater than the limit number
+                num_rows: Some(self.fetch * self.output_partitioning().partition_count()),
+                is_exact: false,
+                ..Default::default()
+            },
         }
     }
 }