ARROW-12290: [Rust][DataFusion] Add input_file_name function [WIP]

rust/datafusion/src/datasource/datasource.rs

@@ -25,18 +25,58 @@ use crate::logical_plan::Expr;
 use crate::physical_plan::ExecutionPlan;
 use crate::{arrow::datatypes::SchemaRef, scalar::ScalarValue};
 
-/// This table statistics are estimates.
+/// These table statistics are estimates.
 /// It can not be used directly in the precise compute
 #[derive(Debug, Clone, Default)]
 pub struct Statistics {
+    /// The partition level statistics
+    pub partition_statistics: Vec<PartitionStatistics>,
+}
+
+impl Statistics {
     /// The number of table rows
+    pub fn num_rows(&self) -> Option<usize> {
+        self.partition_statistics.iter().map(|ps| ps.num_rows).sum()
+    }
+
+    /// total bytes of the table rows
+    pub fn total_byte_size(&self) -> Option<usize> {
+        self.partition_statistics
+            .iter()
+            .map(|ps| ps.total_byte_size)
+            .sum()
+    }
+
+    /// Statistics on a column level for the table
+    pub fn column_statistics(&self) -> Option<Vec<ColumnStatistics>> {
+        self.partition_statistics
+            .iter()
+            .map(|ps| {
+                ps.clone().column_statistics.map(|cs| ColumnStatistics {
+                    null_count: cs.iter().map(|cs| cs.null_count).sum(),
+                    distinct_count: None,
+                    max_value: None,
+                    min_value: None,
+                })
+            })
+            .collect()
+    }
+}
+
+/// These partitions statistics are estimates about the partition
+#[derive(Clone, Default, Debug, PartialEq)]
+pub struct PartitionStatistics {
+    /// The filename for this partition
+    pub filename: Option<String>,
+    /// The number of partition rows
     pub num_rows: Option<usize>,
-    /// total byte of the table rows
+    /// total byte of the partition rows
     pub total_byte_size: Option<usize>,
-    /// Statistics on a column level
+    /// Statistics on a column level for this partition
     pub column_statistics: Option<Vec<ColumnStatistics>>,
 }
-/// This table statistics are estimates about column
+
+/// These table statistics are estimates about column
 #[derive(Clone, Debug, PartialEq)]
 pub struct ColumnStatistics {
     /// Number of null values on column

rust/datafusion/src/datasource/empty.rs

@@ -22,7 +22,7 @@ use std::sync::Arc;
 
 use arrow::datatypes::*;
 
-use crate::datasource::datasource::Statistics;
+use crate::datasource::datasource::{PartitionStatistics, Statistics};
 use crate::datasource::TableProvider;
 use crate::error::Result;
 use crate::logical_plan::Expr;
@@ -72,9 +72,12 @@ impl TableProvider for EmptyTable {
 
     fn statistics(&self) -> Statistics {
         Statistics {
-            num_rows: Some(0),
-            total_byte_size: Some(0),
-            column_statistics: None,
+            partition_statistics: vec![PartitionStatistics {
+                filename: None,
+                num_rows: Some(0),
+                total_byte_size: Some(0),
+                column_statistics: None,
+            }],
         }
     }
 }

rust/datafusion/src/datasource/memory.rs

@@ -34,7 +34,7 @@ use crate::physical_plan::common;
 use crate::physical_plan::memory::MemoryExec;
 use crate::physical_plan::ExecutionPlan;
 use crate::{
-    datasource::datasource::Statistics,
+    datasource::datasource::{PartitionStatistics, Statistics},
     physical_plan::{repartition::RepartitionExec, Partitioning},
 };
 
@@ -48,40 +48,32 @@ pub struct MemTable {
 }
 
 // Calculates statistics based on partitions
-fn calculate_statistics(
-    schema: &SchemaRef,
-    partitions: &[Vec<RecordBatch>],
-) -> Statistics {
-    let num_rows: usize = partitions
+fn calculate_statistics(partitions: &[Vec<RecordBatch>]) -> Statistics {
+    let partition_statistics = partitions
         .iter()
-        .flat_map(|batches| batches.iter().map(RecordBatch::num_rows))
-        .sum();
-
-    let mut null_count: Vec<usize> = vec![0; schema.fields().len()];
-    for partition in partitions.iter() {
-        for batch in partition {
-            for (i, array) in batch.columns().iter().enumerate() {
-                null_count[i] += array.null_count();
-            }
-        }
-    }
-
-    let column_statistics = Some(
-        null_count
-            .iter()
-            .map(|null_count| ColumnStatistics {
-                null_count: Some(*null_count),
-                distinct_count: None,
-                max_value: None,
-                min_value: None,
+        .flat_map(|batches| {
+            batches.iter().map(|batch| PartitionStatistics {
+                filename: None,
+                num_rows: Some(batch.num_rows()),
+                total_byte_size: None,
+                column_statistics: Some(
+                    batch
+                        .columns()
+                        .iter()
+                        .map(|array| ColumnStatistics {
+                            null_count: Some(array.null_count()),
+                            distinct_count: None,
+                            max_value: None,
+                            min_value: None,
+                        })
+                        .collect::<Vec<ColumnStatistics>>(),
+                ),
             })
-            .collect(),
-    );
+        })
+        .collect::<Vec<PartitionStatistics>>();
 
     Statistics {
-        num_rows: Some(num_rows),
-        total_byte_size: None,
-        column_statistics,
+        partition_statistics,
     }
 }
 
@@ -93,7 +85,7 @@ impl MemTable {
             .flatten()
             .all(|batches| schema.contains(&batches.schema()))
         {
-            let statistics = calculate_statistics(&schema, &partitions);
+            let statistics = calculate_statistics(&partitions);
             debug!("MemTable statistics: {:?}", statistics);
 
             Ok(Self {
@@ -247,35 +239,15 @@ mod tests {
 
         let provider = MemTable::try_new(schema, vec![vec![batch]])?;
 
-        assert_eq!(provider.statistics().num_rows, Some(3));
+        assert_eq!(provider.statistics().num_rows(), Some(3));
         assert_eq!(
-            provider.statistics().column_statistics,
-            Some(vec![
-                ColumnStatistics {
-                    null_count: Some(0),
-                    max_value: None,
-                    min_value: None,
-                    distinct_count: None,
-                },
-                ColumnStatistics {
-                    null_count: Some(0),
-                    max_value: None,
-                    min_value: None,
-                    distinct_count: None,
-                },
-                ColumnStatistics {
-                    null_count: Some(0),
-                    max_value: None,
-                    min_value: None,
-                    distinct_count: None,
-                },
-                ColumnStatistics {
-                    null_count: Some(2),
-                    max_value: None,
-                    min_value: None,
-                    distinct_count: None,
-                },
-            ])
+            provider.statistics().column_statistics(),
+            Some(vec![ColumnStatistics {
+                null_count: Some(2),
+                max_value: None,
+                min_value: None,
+                distinct_count: None,
+            },])
         );
 
         // scan with projection
@@ -465,7 +437,7 @@ mod tests {
         let batch1 = it.next().await.unwrap()?;
         assert_eq!(3, batch1.schema().fields().len());
         assert_eq!(3, batch1.num_columns());
-        assert_eq!(provider.statistics().num_rows, Some(6));
+        assert_eq!(provider.statistics().num_rows(), Some(6));
 
         Ok(())
     }

rust/datafusion/src/datasource/parquet.rs

@@ -48,7 +48,9 @@ impl ParquetTable {
         Ok(Self {
             path: path.to_string(),
             schema,
-            statistics: parquet_exec.statistics().to_owned(),
+            statistics: Statistics {
+                partition_statistics: parquet_exec.partitions().to_owned(),
+            },
             max_concurrency,
         })
     }
@@ -130,8 +132,8 @@ mod tests {
             .await;
 
         // test metadata
-        assert_eq!(table.statistics().num_rows, Some(8));
-        assert_eq!(table.statistics().total_byte_size, Some(671));
+        assert_eq!(table.statistics().num_rows(), Some(8));
+        assert_eq!(table.statistics().total_byte_size(), Some(671));
 
         Ok(())
     }

rust/datafusion/src/optimizer/hash_build_probe_order.rs

@@ -13,7 +13,7 @@
 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
-// under the License
+// under the License.
 
 //! Optimizer rule to switch build and probe order of hash join
 //! based on statistics of a `TableProvider`. If the number of
@@ -39,7 +39,7 @@ pub struct HashBuildProbeOrder {}
 // Gets exact number of rows, if known by the statistics of the underlying
 fn get_num_rows(logical_plan: &LogicalPlan) -> Option<usize> {
     match logical_plan {
-        LogicalPlan::TableScan { source, .. } => source.statistics().num_rows,
+        LogicalPlan::TableScan { source, .. } => source.statistics().num_rows(),
         LogicalPlan::EmptyRelation {
             produce_one_row, ..
         } => {
@@ -187,13 +187,14 @@ mod tests {
     use std::sync::Arc;
 
     use crate::{
-        datasource::{datasource::Statistics, TableProvider},
+        datasource::datasource::{PartitionStatistics, Statistics},
+        datasource::TableProvider,
         logical_plan::{DFSchema, Expr},
         test::*,
     };
 
     struct TestTableProvider {
-        num_rows: usize,
+        partition_statistics: Vec<PartitionStatistics>,
     }
 
     impl TableProvider for TestTableProvider {
@@ -215,9 +216,7 @@ mod tests {
         }
         fn statistics(&self) -> crate::datasource::datasource::Statistics {
             Statistics {
-                num_rows: Some(self.num_rows),
-                total_byte_size: None,
-                column_statistics: None,
+                partition_statistics: self.partition_statistics.clone(),
             }
         }
     }
@@ -236,7 +235,14 @@ mod tests {
         let lp_left = LogicalPlan::TableScan {
             table_name: "left".to_string(),
             projection: None,
-            source: Arc::new(TestTableProvider { num_rows: 1000 }),
+            source: Arc::new(TestTableProvider {
+                partition_statistics: vec![PartitionStatistics {
+                    filename: None,
+                    num_rows: Some(1000),
+                    total_byte_size: None,
+                    column_statistics: None,
+                }],
+            }),
             projected_schema: Arc::new(DFSchema::empty()),
             filters: vec![],
             limit: None,
@@ -245,7 +251,14 @@ mod tests {
         let lp_right = LogicalPlan::TableScan {
             table_name: "right".to_string(),
             projection: None,
-            source: Arc::new(TestTableProvider { num_rows: 100 }),
+            source: Arc::new(TestTableProvider {
+                partition_statistics: vec![PartitionStatistics {
+                    filename: None,
+                    num_rows: Some(100),
+                    total_byte_size: None,
+                    column_statistics: None,
+                }],
+            }),
             projected_schema: Arc::new(DFSchema::empty()),
             filters: vec![],
             limit: None,

rust/datafusion/src/physical_optimizer/repartition.rs

@@ -106,18 +106,20 @@ mod tests {
     use arrow::datatypes::Schema;
 
     use super::*;
-    use crate::datasource::datasource::Statistics;
-    use crate::physical_plan::parquet::{ParquetExec, ParquetPartition};
+    use crate::datasource::datasource::PartitionStatistics;
+    use crate::physical_plan::parquet::ParquetExec;
     use crate::physical_plan::projection::ProjectionExec;
 
     #[test]
     fn added_repartition_to_single_partition() -> Result<()> {
         let parquet_project = ProjectionExec::try_new(
             vec![],
             Arc::new(ParquetExec::new(
-                vec![ParquetPartition {
-                    filenames: vec!["x".to_string()],
-                    statistics: Statistics::default(),
+                vec![PartitionStatistics {
+                    filename: Some("x".to_string()),
+                    num_rows: None,
+                    total_byte_size: None,
+                    column_statistics: None,
                 }],
                 Schema::empty(),
                 None,
@@ -151,9 +153,11 @@ mod tests {
             Arc::new(ProjectionExec::try_new(
                 vec![],
                 Arc::new(ParquetExec::new(
-                    vec![ParquetPartition {
-                        filenames: vec!["x".to_string()],
-                        statistics: Statistics::default(),
+                    vec![PartitionStatistics {
+                        filename: Some("x".to_string()),
+                        num_rows: None,
+                        total_byte_size: None,
+                        column_statistics: None,
                     }],
                     Schema::empty(),
                     None,