[fix] replace file format providers in ballista

Author: rdettai

ballista-examples/src/bin/ballista-dataframe.rs

@@ -33,7 +33,8 @@ async fn main() -> Result<()> {
 
     // define the query using the DataFrame trait
     let df = ctx
-        .read_parquet(filename)?
+        .read_parquet(filename)
+        .await?
         .select_columns(&["id", "bool_col", "timestamp_col"])?
         .filter(col("id").gt(lit(1)))?;
 

ballista-examples/src/bin/ballista-sql.rs

@@ -34,15 +34,18 @@ async fn main() -> Result<()> {
         "aggregate_test_100",
         &format!("{}/csv/aggregate_test_100.csv", testdata),
         CsvReadOptions::new(),
-    )?;
+    )
+    .await?;
 
     // execute the query
-    let df = ctx.sql(
-        "SELECT c1, MIN(c12), MAX(c12) \
+    let df = ctx
+        .sql(
+            "SELECT c1, MIN(c12), MAX(c12) \
         FROM aggregate_test_100 \
         WHERE c11 > 0.1 AND c11 < 0.9 \
         GROUP BY c1",
-    )?;
+        )
+        .await?;
 
     // print the results
     df.show().await?;

ballista/rust/client/README.md

@@ -100,15 +100,15 @@ async fn main() -> Result<()> {
        "tripdata",
        "/path/to/yellow_tripdata_2020-01.csv",
        CsvReadOptions::new(),
-   )?;
+   ).await?;
 
    // execute the query
    let df = ctx.sql(
        "SELECT passenger_count, MIN(fare_amount), MAX(fare_amount), AVG(fare_amount), SUM(fare_amount)
        FROM tripdata
        GROUP BY passenger_count
        ORDER BY passenger_count",
-   )?;
+   ).await?;
 
    // collect the results and print them to stdout
    let results = df.collect().await?;

ballista/rust/client/src/context.rs

@@ -31,8 +31,7 @@ use datafusion::datasource::TableProvider;
 use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::dataframe_impl::DataFrameImpl;
 use datafusion::logical_plan::LogicalPlan;
-use datafusion::physical_plan::avro::AvroReadOptions;
-use datafusion::physical_plan::csv::CsvReadOptions;
+use datafusion::prelude::{AvroReadOptions, CsvReadOptions};
 use datafusion::sql::parser::FileType;
 
 struct BallistaContextState {
@@ -128,11 +127,11 @@ impl BallistaContext {
     }
 
     /// Create a DataFrame representing an Avro table scan
-
-    pub fn read_avro(
+    /// TODO fetch schema from scheduler instead of resolving locally
+    pub async fn read_avro(
         &self,
         path: &str,
-        options: AvroReadOptions,
+        options: AvroReadOptions<'_>,
     ) -> Result<Arc<dyn DataFrame>> {
         // convert to absolute path because the executor likely has a different working directory
         let path = PathBuf::from(path);
@@ -147,13 +146,13 @@ impl BallistaContext {
                 guard.config(),
             )
         };
-        let df = ctx.read_avro(path.to_str().unwrap(), options)?;
+        let df = ctx.read_avro(path.to_str().unwrap(), options).await?;
         Ok(df)
     }
 
     /// Create a DataFrame representing a Parquet table scan
-
-    pub fn read_parquet(&self, path: &str) -> Result<Arc<dyn DataFrame>> {
+    /// TODO fetch schema from scheduler instead of resolving locally
+    pub async fn read_parquet(&self, path: &str) -> Result<Arc<dyn DataFrame>> {
         // convert to absolute path because the executor likely has a different working directory
         let path = PathBuf::from(path);
         let path = fs::canonicalize(&path)?;
@@ -167,16 +166,16 @@ impl BallistaContext {
                 guard.config(),
             )
         };
-        let df = ctx.read_parquet(path.to_str().unwrap())?;
+        let df = ctx.read_parquet(path.to_str().unwrap()).await?;
         Ok(df)
     }
 
     /// Create a DataFrame representing a CSV table scan
-
-    pub fn read_csv(
+    /// TODO fetch schema from scheduler instead of resolving locally
+    pub async fn read_csv(
         &self,
         path: &str,
-        options: CsvReadOptions,
+        options: CsvReadOptions<'_>,
     ) -> Result<Arc<dyn DataFrame>> {
         // convert to absolute path because the executor likely has a different working directory
         let path = PathBuf::from(path);
@@ -191,7 +190,7 @@ impl BallistaContext {
                 guard.config(),
             )
         };
-        let df = ctx.read_csv(path.to_str().unwrap(), options)?;
+        let df = ctx.read_csv(path.to_str().unwrap(), options).await?;
         Ok(df)
     }
 
@@ -206,39 +205,41 @@ impl BallistaContext {
         Ok(())
     }
 
-    pub fn register_csv(
+    pub async fn register_csv(
         &self,
         name: &str,
         path: &str,
-        options: CsvReadOptions,
+        options: CsvReadOptions<'_>,
     ) -> Result<()> {
-        match self.read_csv(path, options)?.to_logical_plan() {
+        match self.read_csv(path, options).await?.to_logical_plan() {
             LogicalPlan::TableScan { source, .. } => self.register_table(name, source),
             _ => Err(DataFusionError::Internal("Expected tables scan".to_owned())),
         }
     }
 
-    pub fn register_parquet(&self, name: &str, path: &str) -> Result<()> {
-        match self.read_parquet(path)?.to_logical_plan() {
+    pub async fn register_parquet(&self, name: &str, path: &str) -> Result<()> {
+        match self.read_parquet(path).await?.to_logical_plan() {
             LogicalPlan::TableScan { source, .. } => self.register_table(name, source),
             _ => Err(DataFusionError::Internal("Expected tables scan".to_owned())),
         }
     }
 
-    pub fn register_avro(
+    pub async fn register_avro(
         &self,
         name: &str,
         path: &str,
-        options: AvroReadOptions,
+        options: AvroReadOptions<'_>,
     ) -> Result<()> {
-        match self.read_avro(path, options)?.to_logical_plan() {
+        match self.read_avro(path, options).await?.to_logical_plan() {
             LogicalPlan::TableScan { source, .. } => self.register_table(name, source),
             _ => Err(DataFusionError::Internal("Expected tables scan".to_owned())),
         }
     }
 
-    /// Create a DataFrame from a SQL statement
-    pub fn sql(&self, sql: &str) -> Result<Arc<dyn DataFrame>> {
+    /// Create a DataFrame from a SQL statement.
+    ///
+    /// Async because CreateExternalTable might require to resolve the schema
+    pub async fn sql(&self, sql: &str) -> Result<Arc<dyn DataFrame>> {
         let mut ctx = {
             let state = self.state.lock().unwrap();
             create_df_ctx_with_ballista_query_planner(
@@ -275,15 +276,17 @@ impl BallistaContext {
                         CsvReadOptions::new()
                             .schema(&schema.as_ref().to_owned().into())
                             .has_header(*has_header),
-                    )?;
+                    )
+                    .await?;
                     Ok(Arc::new(DataFrameImpl::new(ctx.state, &plan)))
                 }
                 FileType::Parquet => {
-                    self.register_parquet(name, location)?;
+                    self.register_parquet(name, location).await?;
                     Ok(Arc::new(DataFrameImpl::new(ctx.state, &plan)))
                 }
                 FileType::Avro => {
-                    self.register_avro(name, location, AvroReadOptions::default())?;
+                    self.register_avro(name, location, AvroReadOptions::default())
+                        .await?;
                     Ok(Arc::new(DataFrameImpl::new(ctx.state, &plan)))
                 }
                 _ => Err(DataFusionError::NotImplemented(format!(
@@ -292,7 +295,7 @@ impl BallistaContext {
                 ))),
             },
 
-            _ => ctx.sql(sql),
+            _ => ctx.sql(sql).await,
         }
     }
 }
@@ -306,7 +309,7 @@ mod tests {
         let context = BallistaContext::standalone(&BallistaConfig::new().unwrap(), 1)
             .await
             .unwrap();
-        let df = context.sql("SELECT 1;").unwrap();
+        let df = context.sql("SELECT 1;").await.unwrap();
         df.collect().await.unwrap();
     }
 }

ballista/rust/core/Cargo.toml

@@ -41,6 +41,7 @@ sqlparser = "0.11.0"
 tokio = "1.0"
 tonic = "0.5"
 uuid = { version = "0.8", features = ["v4"] }
+chrono = "0.4"
 
 arrow-flight = { version = "^5.3"  }
 

ballista/rust/core/proto/ballista.proto

@@ -20,6 +20,8 @@ syntax = "proto3";
 
 package ballista.protobuf;
 
+import "google/protobuf/timestamp.proto";
+
 option java_multiple_files = true;
 option java_package = "org.ballistacompute.protobuf";
 option java_outer_classname = "BallistaProto";
@@ -239,8 +241,7 @@ message SortExprNode {
 // LogicalPlan is a nested type
 message LogicalPlanNode {
   oneof LogicalPlanType {
-    CsvTableScanNode csv_scan = 1;
-    ParquetTableScanNode parquet_scan = 2;
+    ListingTableScanNode listing_scan = 1;
     ProjectionNode projection = 3;
     SelectionNode selection = 4;
     LimitNode limit = 5;
@@ -254,25 +255,13 @@ message LogicalPlanNode {
     WindowNode window = 13;
     AnalyzeNode analyze = 14;
     CrossJoinNode cross_join = 15;
-    AvroTableScanNode avro_scan = 16;
   }
 }
 
 message ProjectionColumns {
   repeated string columns = 1;
 }
 
-message CsvTableScanNode {
-  string table_name = 1;
-  string path = 2;
-  bool has_header = 3;
-  string delimiter = 4;
-  string file_extension = 5;
-  ProjectionColumns projection = 6;
-  Schema schema = 7;
-  repeated LogicalExprNode filters = 8;
-}
-
 message Statistics {
   int64 num_rows = 1;
   int64 total_byte_size = 2;
@@ -282,30 +271,36 @@ message Statistics {
 
 message PartitionedFile {
   string path = 1;
-  Statistics statistics = 2;
+  uint64 size = 2;
+  uint64 last_modified_ns = 3;
 }
 
-message TableDescriptor {
-  string path = 1;
-  repeated PartitionedFile partition_files = 2;
-  Schema schema = 3;
+message CsvFormat {
+  bool has_header = 1;
+  string delimiter = 2;
 }
 
-message ParquetTableScanNode {
-  string table_name = 1;
-  TableDescriptor table_desc = 2;
-  ProjectionColumns projection = 3;
-  repeated LogicalExprNode filters = 4;
-  uint32 target_partitions = 5;
+message ParquetFormat {
+  bool enable_pruning = 1;
 }
 
-message AvroTableScanNode {
+message AvroFormat {}
+
+message ListingTableScanNode {
   string table_name = 1;
   string path = 2;
   string file_extension = 3;
   ProjectionColumns projection = 4;
   Schema schema = 5;
   repeated LogicalExprNode filters = 6;
+  repeated string partitions = 7;
+  bool collect_stat = 8;
+  uint32 target_partitions = 9;
+  oneof FileFormatType {
+    CsvFormat csv = 10;
+    ParquetFormat parquet = 11;
+    AvroFormat avro = 12;
+  }
 }
 
 message ProjectionNode {
@@ -598,40 +593,42 @@ message FilterExecNode {
   PhysicalExprNode expr = 2;
 }
 
-message ParquetPartition {
-  uint32 index = 1;
-  repeated PartitionedFile files = 2;
+message FilePartition {
+  repeated PartitionedFile files = 1;
+}
+
+message ScanLimit {
+  // wrap into a message to make it optional
+  uint32 limit = 1;
 }
 
 message ParquetScanExecNode {
-  repeated ParquetPartition partitions = 1;
+  repeated FilePartition partitions = 1;
   Schema schema = 2;
-  repeated uint32 projection = 3;
   uint32 batch_size = 4;
+  repeated uint32 projection = 6;
+  ScanLimit limit = 7;
+  Statistics statistics = 8;
 }
 
 message CsvScanExecNode {
-  string path = 1;
-  repeated uint32 projection = 2;
-  Schema schema = 3;
-  string file_extension = 4;
-  bool has_header = 5;
-  uint32 batch_size = 6;
-  string delimiter = 7;
-
-  // partition filenames
-  repeated string filename = 8;
+  repeated PartitionedFile files = 1;
+  Schema schema = 2;
+  bool has_header = 3;
+  uint32 batch_size = 4;
+  string delimiter = 5;
+  repeated uint32 projection = 6;
+  ScanLimit limit = 7;
+  Statistics statistics = 8;
 }
 
 message AvroScanExecNode {
-  string path = 1;
-  repeated uint32 projection = 2;
-  Schema schema = 3;
-  string file_extension = 4;
-  uint32 batch_size = 5;
-
-  // partition filenames
-  repeated string filename = 8;
+  repeated PartitionedFile files = 1;
+  Schema schema = 2;
+  uint32 batch_size = 4;
+  repeated uint32 projection = 6;
+  ScanLimit limit = 7;
+  Statistics statistics = 8;
 }
 
 enum PartitionMode {
@@ -946,7 +943,6 @@ message GetFileMetadataParams {
 
 message GetFileMetadataResult {
   Schema schema = 1;
-  repeated FilePartitionMetadata partitions = 2;
 }
 
 message FilePartitionMetadata {