Add ability to set statistics in the SQLExecutor

Author: nuno-faria

datafusion-federation/src/schema_cast/mod.rs

@@ -1,5 +1,6 @@
 use async_stream::stream;
 use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::common::Statistics;
 use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
@@ -113,4 +114,8 @@ impl ExecutionPlan for SchemaCastScanExec {
             },
         )))
     }
+
+    fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
+        self.input.partition_statistics(partition)
+    }
 }

datafusion-federation/src/sql/executor.rs

@@ -2,6 +2,7 @@ use async_trait::async_trait;
 use core::fmt;
 use datafusion::{
     arrow::datatypes::SchemaRef,
+    common::Statistics,
     error::Result,
     logical_expr::LogicalPlan,
     physical_plan::SendableRecordBatchStream,
@@ -43,6 +44,18 @@ pub trait SQLExecutor: Sync + Send {
     /// Execute a SQL query
     fn execute(&self, query: &str, schema: SchemaRef) -> Result<SendableRecordBatchStream>;
 
+    /// Returns statistics for this `SQLExecutor` node. If statistics are not available, it should
+    /// return [`Statistics::new_unknown`] (the default), not an error. See the `ExecutionPlan`
+    /// trait.
+    fn partition_statistics(
+        &self,
+        _partition: Option<usize>,
+        _query: &str,
+        schema: SchemaRef,
+    ) -> Result<Statistics> {
+        Ok(Statistics::new_unknown(&schema))
+    }
+
     /// Returns the tables provided by the remote
     async fn table_names(&self) -> Result<Vec<String>>;
 

datafusion-federation/src/sql/mod.rs

@@ -12,6 +12,7 @@ use async_trait::async_trait;
 use datafusion::{
     arrow::datatypes::{Schema, SchemaRef},
     common::tree_node::{Transformed, TreeNode},
+    common::Statistics,
     error::{DataFusionError, Result},
     execution::{context::SessionState, TaskContext},
     logical_expr::{Extension, LogicalPlan},
@@ -343,6 +344,11 @@ impl ExecutionPlan for VirtualExecutionPlan {
     fn properties(&self) -> &PlanProperties {
         &self.props
     }
+
+    fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
+        self.executor
+            .partition_statistics(partition, &self.final_sql()?, self.schema())
+    }
 }
 
 #[cfg(test)]