compiles

andygrove · Sep 12, 2024 · 90dda5a · 90dda5a
1 parent 49d3a58
commit 90dda5a
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 11 deletions.
diff --git a/ballista-cli/src/command.rs b/ballista-cli/src/command.rs
@@ -56,21 +56,25 @@ impl Command {
     ) -> Result<()> {
         let now = Instant::now();
         match self {
-            Self::Help => print_options
-                .print_batches(&[all_commands_info()], now)
+            Self::Help =>
+                // TODO need to provide valid schema
+                print_options
+                .print_batches(Arc::new(Schema::empty()), &[all_commands_info()], now)
                 .map_err(BallistaError::DataFusionError),
             Self::ListTables => {
                 let df = ctx.sql("SHOW TABLES").await?;
+                let schema = Arc::new(df.schema().as_arrow().clone());
                 let batches = df.collect().await?;
                 print_options
-                    .print_batches(&batches, now)
+                    .print_batches(schema, &batches, now)
                     .map_err(BallistaError::DataFusionError)
             }
             Self::DescribeTable(name) => {
                 let df = ctx.sql(&format!("SHOW COLUMNS FROM {name}")).await?;
+                let schema = Arc::new(df.schema().as_arrow().clone());
                 let batches = df.collect().await?;
                 print_options
-                    .print_batches(&batches, now)
+                    .print_batches(schema, &batches, now)
                     .map_err(BallistaError::DataFusionError)
             }
             Self::QuietMode(quiet) => {

diff --git a/ballista-cli/src/exec.rs b/ballista-cli/src/exec.rs
@@ -20,6 +20,7 @@
 use std::fs::File;
 use std::io::prelude::*;
 use std::io::BufReader;
+use std::sync::Arc;
 use std::time::Instant;
 
 use ballista::prelude::{BallistaContext, Result};
@@ -176,8 +177,9 @@ async fn exec_and_print(
 ) -> Result<()> {
     let now = Instant::now();
     let df = ctx.sql(&sql).await?;
+    let schema = Arc::new(df.schema().as_arrow().clone());
     let results = df.collect().await?;
-    print_options.print_batches(&results, now)?;
+    print_options.print_batches(schema, &results, now)?;
 
     Ok(())
 }
diff --git a/ballista/client/src/context.rs b/ballista/client/src/context.rs
@@ -22,7 +22,6 @@ use datafusion::execution::context::DataFilePaths;
 use log::info;
 use parking_lot::Mutex;
 use sqlparser::ast::Statement;
-use std::borrow::Cow;
 use std::collections::HashMap;
 use std::sync::Arc;
 
@@ -373,12 +372,12 @@ impl BallistaContext {
             for (name, prov) in &state.tables {
                 // ctx is shared between queries, check table exists or not before register
                 let table_ref = TableReference::Bare {
-                    table: Cow::Borrowed(name),
+                    table: name.as_str().into(),
                 };
                 if !ctx.table_exist(table_ref)? {
                     ctx.register_table(
                         TableReference::Bare {
-                            table: Cow::Borrowed(name),
+                            table: name.as_str().into(),
                         },
                         Arc::clone(prov),
                     )?;
@@ -402,7 +401,7 @@ impl BallistaContext {
                     ..
                 },
             )) => {
-                let table_exists = ctx.table_exist(name)?;
+                let table_exists = ctx.table_exist(name.to_owned())?;
                 let schema: SchemaRef = Arc::new(schema.as_ref().to_owned().into());
                 let table_partition_cols = table_partition_cols
                     .iter()

diff --git a/ballista/core/src/plugin/udf.rs b/ballista/core/src/plugin/udf.rs
@@ -22,7 +22,7 @@ use std::any::Any;
 use std::collections::HashMap;
 use std::io;
 use std::sync::Arc;
-use datafusion::logical_expr::{AggregateUDF, AggregateUDFImpl, ScalarUDF};
+use datafusion::logical_expr::{AggregateUDF, ScalarUDF};
 
 /// UDF plugin trait
 pub trait UDFPlugin: Plugin {

diff --git a/ballista/executor/src/executor.rs b/ballista/executor/src/executor.rs
@@ -28,7 +28,7 @@ use ballista_core::serde::scheduler::PartitionId;
 use dashmap::DashMap;
 use datafusion::execution::context::TaskContext;
 use datafusion::execution::runtime_env::RuntimeEnv;
-use datafusion::logical_expr::{AggregateUDF, AggregateUDFImpl, ScalarUDF, ScalarUDFImpl, WindowUDF};
+use datafusion::logical_expr::{AggregateUDF, ScalarUDF, WindowUDF};
 use futures::future::AbortHandle;
 use std::collections::HashMap;
 use std::future::Future;