Add support for partitioned parquet data sources

rust/datafusion/src/datasource/parquet.rs

@@ -37,23 +37,27 @@ use parquet::file::reader::*;
 
 use crate::datasource::{ScanResult, TableProvider};
 use crate::error::{ExecutionError, Result};
+use crate::execution::physical_plan::common;
 use crate::execution::physical_plan::BatchIterator;
 
 /// Table-based representation of a `ParquetFile`
 pub struct ParquetTable {
-    filename: String,
+    filenames: Vec<String>,
     schema: Arc<Schema>,
 }
 
 impl ParquetTable {
     /// Attempt to initialize a new `ParquetTable` from a file path
-    pub fn try_new(filename: &str) -> Result<Self> {
-        let parquet_file = ParquetFile::open(filename, None, 0)?;
-        let schema = parquet_file.projection_schema.clone();
-        Ok(Self {
-            filename: filename.to_string(),
-            schema,
-        })
+    pub fn try_new(path: &str) -> Result<Self> {
+        let mut filenames: Vec<String> = vec![];
+        common::build_file_list(path, &mut filenames, ".parquet")?;
+        if filenames.is_empty() {
+            panic!()
+        } else {
+            let parquet_file = ParquetFile::open(&filenames[0], None, 0)?;
+            let schema = parquet_file.projection_schema.clone();
+            Ok(Self { filenames, schema })
+        }
     }
 }
 
@@ -70,17 +74,21 @@ impl TableProvider for ParquetTable {
         projection: &Option<Vec<usize>>,
         batch_size: usize,
     ) -> Result<Vec<ScanResult>> {
-        // note that this code currently assumes the filename is a file rather than a directory
-        // and therefore only returns a single partition
-        let parquet_file = match projection {
-            Some(p) => ParquetScanPartition::try_new(
-                &self.filename,
-                Some(p.clone()),
-                batch_size,
-            )?,
-            None => ParquetScanPartition::try_new(&self.filename, None, batch_size)?,
-        };
-        Ok(vec![Arc::new(Mutex::new(parquet_file))])
+        //TODO remove the unwrap
+        Ok(self
+            .filenames
+            .iter()
+            .map(|filename| {
+                Arc::new(Mutex::new(
+                    ParquetScanPartition::try_new(
+                        filename,
+                        projection.clone(),
+                        batch_size,
+                    )
+                    .unwrap(),
+                )) as Arc<Mutex<dyn BatchIterator>>
+            })
+            .collect())
     }
 }
 

rust/datafusion/src/execution/physical_plan/common.rs

@@ -19,7 +19,6 @@
 
 use std::fs;
 use std::fs::metadata;
-use std::fs::File;
 use std::sync::{Arc, Mutex};
 
 use crate::error::{ExecutionError, Result};

rust/datafusion/src/execution/physical_plan/csv.rs

@@ -17,12 +17,10 @@
 
 //! Execution plan for reading CSV files
 
-use std::fs;
-use std::fs::metadata;
 use std::fs::File;
 use std::sync::{Arc, Mutex};
 
-use crate::error::{ExecutionError, Result};
+use crate::error::Result;
 use crate::execution::physical_plan::common;
 use crate::execution::physical_plan::{BatchIterator, ExecutionPlan, Partition};
 use arrow::csv;