Upgrade to datafusion 16

ballista-cli/Cargo.toml

@@ -33,8 +33,8 @@ ballista = { path = "../ballista/client", version = "0.10.0", features = [
     "standalone",
 ] }
 clap = { version = "3", features = ["derive", "cargo"] }
-datafusion = "15.0.0"
-datafusion-cli = "15.0.0"
+datafusion = "16.0.0"
+datafusion-cli = "16.0.0"
 dirs = "4.0.0"
 env_logger = "0.10"
 mimalloc = { version = "0.1", default-features = false }

ballista/client/Cargo.toml

@@ -31,12 +31,12 @@ rust-version = "1.63"
 ballista-core = { path = "../core", version = "0.10.0" }
 ballista-executor = { path = "../executor", version = "0.10.0", optional = true }
 ballista-scheduler = { path = "../scheduler", version = "0.10.0", optional = true }
-datafusion = "15.0.0"
-datafusion-proto = "15.0.0"
+datafusion = "16.0.0"
+datafusion-proto = "16.0.0"
 futures = "0.3"
 log = "0.4"
 parking_lot = "0.12"
-sqlparser = "0.27"
+sqlparser = "0.30"
 tempfile = "3"
 tokio = "1.0"
 

ballista/client/src/context.rs

@@ -33,6 +33,7 @@ use ballista_core::utils::{
 use datafusion_proto::protobuf::LogicalPlanNode;
 
 use datafusion::catalog::TableReference;
+use datafusion::common::OwnedTableReference;
 use datafusion::dataframe::DataFrame;
 use datafusion::datasource::{source_as_provider, TableProvider};
 use datafusion::error::{DataFusionError, Result};
@@ -212,7 +213,7 @@ impl BallistaContext {
         &self,
         path: &str,
         options: AvroReadOptions<'_>,
-    ) -> Result<Arc<DataFrame>> {
+    ) -> Result<DataFrame> {
         let df = self.context.read_avro(path, options).await?;
         Ok(df)
     }
@@ -223,7 +224,7 @@ impl BallistaContext {
         &self,
         path: &str,
         options: ParquetReadOptions<'_>,
-    ) -> Result<Arc<DataFrame>> {
+    ) -> Result<DataFrame> {
         let df = self.context.read_parquet(path, options).await?;
         Ok(df)
     }
@@ -234,7 +235,7 @@ impl BallistaContext {
         &self,
         path: &str,
         options: CsvReadOptions<'_>,
-    ) -> Result<Arc<DataFrame>> {
+    ) -> Result<DataFrame> {
         let df = self.context.read_csv(path, options).await?;
         Ok(df)
     }
@@ -256,13 +257,10 @@ impl BallistaContext {
         path: &str,
         options: CsvReadOptions<'_>,
     ) -> Result<()> {
-        let plan = self
-            .read_csv(path, options)
-            .await
-            .map_err(|e| {
-                DataFusionError::Context(format!("Can't read CSV: {}", path), Box::new(e))
-            })?
-            .to_logical_plan()?;
+        let df = self.read_csv(path, options).await.map_err(|e| {
+            DataFusionError::Context(format!("Can't read CSV: {}", path), Box::new(e))
+        })?;
+        let plan = df.into_optimized_plan()?;
         match plan {
             LogicalPlan::TableScan(TableScan { source, .. }) => {
                 self.register_table(name, source_as_provider(&source)?)
@@ -277,7 +275,11 @@ impl BallistaContext {
         path: &str,
         options: ParquetReadOptions<'_>,
     ) -> Result<()> {
-        match self.read_parquet(path, options).await?.to_logical_plan()? {
+        match self
+            .read_parquet(path, options)
+            .await?
+            .into_optimized_plan()?
+        {
             LogicalPlan::TableScan(TableScan { source, .. }) => {
                 self.register_table(name, source_as_provider(&source)?)
             }
@@ -291,7 +293,7 @@ impl BallistaContext {
         path: &str,
         options: AvroReadOptions<'_>,
     ) -> Result<()> {
-        match self.read_avro(path, options).await?.to_logical_plan()? {
+        match self.read_avro(path, options).await?.into_optimized_plan()? {
             LogicalPlan::TableScan(TableScan { source, .. }) => {
                 self.register_table(name, source_as_provider(&source)?)
             }
@@ -302,17 +304,16 @@ impl BallistaContext {
     /// is a 'show *' sql
     pub async fn is_show_statement(&self, sql: &str) -> Result<bool> {
         let mut is_show_variable: bool = false;
-        let statements = DFParser::parse_sql(sql)?;
+        let mut statements = DFParser::parse_sql(sql)?;
 
         if statements.len() != 1 {
             return Err(DataFusionError::NotImplemented(
                 "The context currently only supports a single SQL statement".to_string(),
             ));
         }
 
-        if let DFStatement::Statement(s) = &statements[0] {
-            let st: &Statement = s;
-            match st {
+        if let Some(DFStatement::Statement(s)) = statements.remove(0) {
+            match *s {
                 Statement::ShowVariable { .. } => {
                     is_show_variable = true;
                 }
@@ -332,7 +333,7 @@ impl BallistaContext {
     ///
     /// This method is `async` because queries of type `CREATE EXTERNAL TABLE`
     /// might require the schema to be inferred.
-    pub async fn sql(&self, sql: &str) -> Result<Arc<DataFrame>> {
+    pub async fn sql(&self, sql: &str) -> Result<DataFrame> {
         let mut ctx = self.context.clone();
 
         let is_show = self.is_show_statement(sql).await?;
@@ -361,7 +362,7 @@ impl BallistaContext {
             }
         }
 
-        let plan = ctx.create_logical_plan(sql)?;
+        let plan = ctx.state().create_logical_plan(sql).await?;
 
         match plan {
             LogicalPlan::CreateExternalTable(CreateExternalTable {
@@ -375,6 +376,11 @@ impl BallistaContext {
                 ref if_not_exists,
                 ..
             }) => {
+                let name = match name {
+                    OwnedTableReference::Bare { table, .. } => table,
+                    OwnedTableReference::Partial { table, .. } => table,
+                    OwnedTableReference::Full { table, .. } => table,
+                };
                 let table_exists = ctx.table_exist(name.as_str())?;
                 let schema: SchemaRef = Arc::new(schema.as_ref().to_owned().into());
                 let table_partition_cols = table_partition_cols
@@ -398,7 +404,7 @@ impl BallistaContext {
                                 options = options.schema(&schema);
                             }
                             self.register_csv(name, location, options).await?;
-                            Ok(Arc::new(DataFrame::new(ctx.state.clone(), &plan)))
+                            Ok(DataFrame::new(ctx.state(), plan))
                         }
                         "parquet" => {
                             self.register_parquet(
@@ -408,7 +414,7 @@ impl BallistaContext {
                                     .table_partition_cols(table_partition_cols),
                             )
                             .await?;
-                            Ok(Arc::new(DataFrame::new(ctx.state.clone(), &plan)))
+                            Ok(DataFrame::new(ctx.state(), plan))
                         }
                         "avro" => {
                             self.register_avro(
@@ -418,16 +424,14 @@ impl BallistaContext {
                                     .table_partition_cols(table_partition_cols),
                             )
                             .await?;
-                            Ok(Arc::new(DataFrame::new(ctx.state.clone(), &plan)))
+                            Ok(DataFrame::new(ctx.state(), plan))
                         }
                         _ => Err(DataFusionError::NotImplemented(format!(
                             "Unsupported file type {:?}.",
                             file_type
                         ))),
                     },
-                    (true, true) => {
-                        Ok(Arc::new(DataFrame::new(ctx.state.clone(), &plan)))
-                    }
+                    (true, true) => Ok(DataFrame::new(ctx.state(), plan)),
                     (false, true) => Err(DataFusionError::Execution(format!(
                         "Table '{:?}' already exists",
                         name

ballista/core/Cargo.toml

@@ -46,13 +46,13 @@ simd = ["datafusion/simd"]
 [dependencies]
 ahash = { version = "0.8", default-features = false }
 
-arrow-flight = { version = "28.0.0", features = ["flight-sql-experimental"] }
+arrow-flight = { version = "29.0.0", features = ["flight-sql-experimental"] }
 async-trait = "0.1.41"
 chrono = { version = "0.4", default-features = false }
 clap = { version = "3", features = ["derive", "cargo"] }
-datafusion = "15.0.0"
+datafusion = "16.0.0"
 datafusion-objectstore-hdfs = { version = "0.1.1", default-features = false, optional = true }
-datafusion-proto = "15.0.0"
+datafusion-proto = "16.0.0"
 futures = "0.3"
 hashbrown = "0.13"
 
@@ -68,7 +68,7 @@ prost = "0.11"
 prost-types = "0.11"
 rand = "0.8"
 serde = { version = "1", features = ["derive"] }
-sqlparser = "0.27"
+sqlparser = "0.30"
 sys-info = "0.9.0"
 tokio = "1.0"
 tokio-stream = { version = "0.1", features = ["net"] }

ballista/core/src/error.rs

@@ -158,14 +158,14 @@ impl From<tokio::task::JoinError> for BallistaError {
     }
 }
 
-impl From<datafusion_proto::from_proto::Error> for BallistaError {
-    fn from(e: datafusion_proto::from_proto::Error) -> Self {
+impl From<datafusion_proto::logical_plan::to_proto::Error> for BallistaError {
+    fn from(e: datafusion_proto::logical_plan::to_proto::Error) -> Self {
         BallistaError::General(e.to_string())
     }
 }
 
-impl From<datafusion_proto::to_proto::Error> for BallistaError {
-    fn from(e: datafusion_proto::to_proto::Error) -> Self {
+impl From<datafusion_proto::logical_plan::from_proto::Error> for BallistaError {
+    fn from(e: datafusion_proto::logical_plan::from_proto::Error) -> Self {
         BallistaError::General(e.to_string())
     }
 }

ballista/core/src/execution_plans/distributed_query.rs

@@ -136,10 +136,6 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
         None
     }
 
-    fn relies_on_input_order(&self) -> bool {
-        false
-    }
-
     fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
         vec![]
     }

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -98,10 +98,6 @@ impl ExecutionPlan for ShuffleReaderExec {
         None
     }
 
-    fn relies_on_input_order(&self) -> bool {
-        false
-    }
-
     fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
         vec![]
     }

ballista/core/src/execution_plans/shuffle_writer.rs

@@ -313,10 +313,6 @@ impl ExecutionPlan for ShuffleWriterExec {
         None
     }
 
-    fn relies_on_input_order(&self) -> bool {
-        false
-    }
-
     fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
         vec![self.plan.clone()]
     }

ballista/core/src/execution_plans/unresolved_shuffle.rs

@@ -81,10 +81,6 @@ impl ExecutionPlan for UnresolvedShuffleExec {
         None
     }
 
-    fn relies_on_input_order(&self) -> bool {
-        false
-    }
-
     fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
         vec![]
     }

ballista/core/src/serde/generated/ballista.rs

@@ -1985,7 +1985,7 @@ pub mod executor_grpc_client {
 pub mod scheduler_grpc_server {
     #![allow(unused_variables, dead_code, missing_docs, clippy::let_unit_value)]
     use tonic::codegen::*;
-    ///Generated trait containing gRPC methods that should be implemented for use with SchedulerGrpcServer.
+    /// Generated trait containing gRPC methods that should be implemented for use with SchedulerGrpcServer.
     #[async_trait]
     pub trait SchedulerGrpc: Send + Sync + 'static {
         /// Executors must poll the scheduler for heartbeat and to receive tasks
@@ -2531,7 +2531,7 @@ pub mod scheduler_grpc_server {
 pub mod executor_grpc_server {
     #![allow(unused_variables, dead_code, missing_docs, clippy::let_unit_value)]
     use tonic::codegen::*;
-    ///Generated trait containing gRPC methods that should be implemented for use with ExecutorGrpcServer.
+    /// Generated trait containing gRPC methods that should be implemented for use with ExecutorGrpcServer.
     #[async_trait]
     pub trait ExecutorGrpc: Send + Sync + 'static {
         async fn launch_task(

ballista/core/src/serde/mod.rs

@@ -262,10 +262,9 @@ mod tests {
         SendableRecordBatchStream, Statistics,
     };
     use datafusion::prelude::{CsvReadOptions, SessionConfig, SessionContext};
+    use datafusion_proto::logical_plan::from_proto::parse_expr;
     use prost::Message;
     use std::any::Any;
-
-    use datafusion_proto::from_proto::parse_expr;
     use std::convert::TryInto;
     use std::fmt;
     use std::fmt::{Debug, Formatter};
@@ -628,13 +627,13 @@ mod tests {
         let scan = ctx
             .read_csv("tests/customer.csv", CsvReadOptions::default())
             .await?
-            .to_logical_plan()?;
+            .into_optimized_plan()?;
 
         let topk_plan = LogicalPlan::Extension(Extension {
             node: Arc::new(TopKPlanNode::new(3, scan, col("revenue"))),
         });
 
-        let topk_exec = ctx.create_physical_plan(&topk_plan).await?;
+        let topk_exec = ctx.state().create_physical_plan(&topk_plan).await?;
 
         let extension_codec = TopKExtensionCodec {};
 

ballista/core/src/serde/physical_plan/from_proto.rs

@@ -23,7 +23,6 @@ use std::sync::Arc;
 
 use chrono::{TimeZone, Utc};
 use datafusion::arrow::datatypes::Schema;
-use datafusion::config::ConfigOptions;
 use datafusion::datasource::listing::{FileRange, PartitionedFile};
 use datafusion::datasource::object_store::ObjectStoreUrl;
 use datafusion::execution::context::ExecutionProps;
@@ -40,10 +39,9 @@ use datafusion::physical_plan::{
     functions, Partitioning,
 };
 use datafusion::physical_plan::{ColumnStatistics, PhysicalExpr, Statistics};
-use datafusion_proto::from_proto::from_proto_binary_op;
+use datafusion_proto::logical_plan::from_proto::from_proto_binary_op;
 use object_store::path::Path;
 use object_store::ObjectMeta;
-use parking_lot::RwLock;
 
 use crate::serde::protobuf::physical_expr_node::ExprType;
 
@@ -399,7 +397,6 @@ impl TryInto<FileScanConfig> for &protobuf::FileScanExecConf {
         let statistics = convert_required!(self.statistics)?;
 
         Ok(FileScanConfig {
-            config_options: Arc::new(RwLock::new(ConfigOptions::new())), // TODO add serde
             object_store_url: ObjectStoreUrl::parse(&self.object_store_url)?,
             file_schema: schema,
             file_groups: self
@@ -413,6 +410,7 @@ impl TryInto<FileScanConfig> for &protobuf::FileScanExecConf {
             table_partition_cols: vec![],
             // TODO add ordering info to the ballista proto file
             output_ordering: None,
+            infinite_source: false,
         })
     }
 }