chore: update datafusion to 51.0 (#1345)

* cherry pick Marko's work

* update code and test based on new DF apis

* fmt toml and bump rust min version

* update benchmark tests

* fmt

* Using task ctx directly

* update protos

* fmt

* remove session context from few more places

* remove unnecessary clone

---------

Co-authored-by: Marko Milenković <[email protected]>

Author: danielhumanmod

Cargo.lock

@@ -27,27 +27,28 @@ resolver = "2"
 # 
 edition = "2021"
 # we should try to follow datafusion version 
-rust-version = "1.86.0"
+rust-version = "1.88.0"
 
 [workspace.dependencies]
-arrow = { version = "56", features = ["ipc_compression"] }
-arrow-flight = { version = "56", features = ["flight-sql-experimental"] }
+arrow = { version = "57", features = ["ipc_compression"] }
+arrow-flight = { version = "57", features = ["flight-sql-experimental"] }
 clap = { version = "4.5", features = ["derive", "cargo"] }
 configure_me = { version = "0.4.0" }
 configure_me_codegen = { version = "0.4.4" }
-datafusion = "50.3.0"
-datafusion-cli = "50.3.0"
-datafusion-proto = "50.3.0"
-datafusion-proto-common = "50.3.0"
+datafusion = "51.0.0"
+datafusion-cli = "51.0.0"
+datafusion-proto = "51.0.0"
+datafusion-proto-common = "51.0.0"
 object_store = "0.12"
-prost = "0.13"
-prost-types = "0.13"
+prost = "0.14"
+prost-types = "0.14"
 rstest = { version = "0.26" }
-tonic = { version = "0.13" }
-tonic-build = { version = "0.13", default-features = false, features = [
+tonic = { version = "0.14" }
+tonic-build = { version = "0.14", default-features = false, features = [
     "transport",
-    "prost"
 ] }
+tonic-prost = { version = "0.14" }
+tonic-prost-build = { version = "0.14" }
 tracing = "0.1"
 tracing-appender = "0.2.2"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }

Cargo.toml

@@ -15,8 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::env;
 use std::path::Path;
+use std::{env, sync::Arc};
 
 use ballista::{extension::SessionConfigExt, prelude::SessionContextExt};
 use ballista_cli::{
@@ -28,7 +28,9 @@ use datafusion::{
     execution::SessionStateBuilder,
     prelude::{SessionConfig, SessionContext},
 };
-use datafusion_cli::print_options::MaxRows;
+use datafusion_cli::{
+    object_storage::instrumented::InstrumentedObjectStoreRegistry, print_options::MaxRows,
+};
 use mimalloc::MiMalloc;
 
 #[global_allocator]
@@ -151,6 +153,7 @@ pub async fn main() -> Result<()> {
         quiet: args.quiet,
         maxrows: MaxRows::Unlimited,
         color: args.color,
+        instrumented_registry: Arc::new(InstrumentedObjectStoreRegistry::new()),
     };
 
     let files = args.file;

ballista-cli/src/main.rs

@@ -113,12 +113,14 @@ mod basic {
 
         let result = df.unwrap().collect().await.unwrap();
 
-        let expected = ["+---------------+--------------+---------------------+-------------+-----------------------------+-------------+",
-            "| table_catalog | table_schema | table_name          | column_name | data_type                   | is_nullable |",
-            "+---------------+--------------+---------------------+-------------+-----------------------------+-------------+",
-            "| datafusion    | public       | csv_with_timestamps | name        | Utf8                        | YES         |",
-            "| datafusion    | public       | csv_with_timestamps | ts          | Timestamp(Nanosecond, None) | YES         |",
-            "+---------------+--------------+---------------------+-------------+-----------------------------+-------------+"];
+        let expected = [
+            "+---------------+--------------+---------------------+-------------+---------------+-------------+",
+            "| table_catalog | table_schema | table_name          | column_name | data_type     | is_nullable |",
+            "+---------------+--------------+---------------------+-------------+---------------+-------------+",
+            "| datafusion    | public       | csv_with_timestamps | name        | Utf8          | YES         |",
+            "| datafusion    | public       | csv_with_timestamps | ts          | Timestamp(ns) | YES         |",
+            "+---------------+--------------+---------------------+-------------+---------------+-------------+",
+        ];
         datafusion::assert_batches_eq!(expected, &result);
     }
 
@@ -359,7 +361,7 @@ mod basic {
         let res = df.collect().await.unwrap();
         let expected = vec![
             "+-------------------+",
-            "| var(test.id)      |",
+            "| var_samp(test.id) |",
             "+-------------------+",
             "| 6.000000000000001 |",
             "+-------------------+",
@@ -390,11 +392,11 @@ mod basic {
             .unwrap();
         let res = df.collect().await.unwrap();
         let expected = vec![
-            "+--------------------+",
-            "| stddev(test.id)    |",
-            "+--------------------+",
-            "| 2.4494897427831783 |",
-            "+--------------------+",
+            "+----------------------+",
+            "| stddev_samp(test.id) |",
+            "+----------------------+",
+            "| 2.4494897427831783   |",
+            "+----------------------+",
         ];
         assert_result_eq(expected, &res);
     }
@@ -408,11 +410,11 @@ mod basic {
             .unwrap();
         let res = df.collect().await.unwrap();
         let expected = vec![
-            "+--------------------------------------+",
-            "| covar_samp(test.id,test.tinyint_col) |",
-            "+--------------------------------------+",
-            "| 0.28571428571428586                  |",
-            "+--------------------------------------+",
+            "+---------------------------------+",
+            "| covar(test.id,test.tinyint_col) |",
+            "+---------------------------------+",
+            "| 0.28571428571428586             |",
+            "+---------------------------------+",
         ];
         assert_result_eq(expected, &res);
     }

ballista/client/tests/context_basic.rs

@@ -110,7 +110,9 @@ mod standalone {
     use datafusion::{
         assert_batches_eq,
         common::exec_err,
-        execution::{context::QueryPlanner, SessionState, SessionStateBuilder},
+        execution::{
+            context::QueryPlanner, SessionState, SessionStateBuilder, TaskContext,
+        },
         logical_expr::LogicalPlan,
         physical_plan::ExecutionPlan,
         prelude::{SessionConfig, SessionContext},
@@ -275,7 +277,7 @@ mod standalone {
             &self,
             _buf: &[u8],
             _inputs: &[datafusion::logical_expr::LogicalPlan],
-            _ctx: &SessionContext,
+            _ctx: &TaskContext,
         ) -> datafusion::error::Result<datafusion::logical_expr::Extension> {
             self.invoked
                 .store(true, std::sync::atomic::Ordering::Relaxed);
@@ -297,7 +299,7 @@ mod standalone {
             _buf: &[u8],
             _table_ref: &datafusion::sql::TableReference,
             _schema: datafusion::arrow::datatypes::SchemaRef,
-            _ctx: &SessionContext,
+            _ctx: &TaskContext,
         ) -> datafusion::error::Result<
             std::sync::Arc<dyn datafusion::catalog::TableProvider>,
         > {
@@ -320,7 +322,7 @@ mod standalone {
         fn try_decode_file_format(
             &self,
             _buf: &[u8],
-            _ctx: &SessionContext,
+            _ctx: &TaskContext,
         ) -> datafusion::error::Result<
             Arc<dyn datafusion::datasource::file_format::FileFormatFactory>,
         > {
@@ -361,12 +363,12 @@ mod standalone {
             &self,
             buf: &[u8],
             inputs: &[Arc<dyn datafusion::physical_plan::ExecutionPlan>],
-            registry: &dyn datafusion::execution::FunctionRegistry,
+            ctx: &TaskContext,
         ) -> datafusion::error::Result<Arc<dyn datafusion::physical_plan::ExecutionPlan>>
         {
             self.invoked
                 .store(true, std::sync::atomic::Ordering::Relaxed);
-            self.codec.try_decode(buf, inputs, registry)
+            self.codec.try_decode(buf, inputs, ctx)
         }
 
         fn try_encode(

ballista/client/tests/context_setup.rs

@@ -63,6 +63,7 @@ serde = { workspace = true, features = ["derive"] }
 tokio = { workspace = true }
 tokio-stream = { workspace = true, features = ["net"] }
 tonic = { workspace = true }
+tonic-prost = { workspace = true }
 url = { workspace = true }
 uuid = { workspace = true }
 
@@ -72,3 +73,4 @@ tempfile = { workspace = true }
 [build-dependencies]
 rustc_version = "0.4.1"
 tonic-build = { workspace = true }
+tonic-prost-build = { workspace = true }

ballista/core/Cargo.toml

@@ -40,7 +40,7 @@ fn main() -> Result<(), String> {
         println!("cargo:rerun-if-changed=proto/datafusion_common.proto");
         println!("cargo:rerun-if-changed=proto/datafusion.proto");
         println!("cargo:rerun-if-changed=proto/ballista.proto");
-        tonic_build::configure()
+        tonic_prost_build::configure()
             .extern_path(".datafusion_common", "::datafusion_proto_common")
             .extern_path(".datafusion", "::datafusion_proto::protobuf")
             .protoc_arg("--experimental_allow_proto3_optional")

ballista/core/build.rs

@@ -264,6 +264,18 @@ message NamedTime {
   uint64 value = 2;
 }
 
+message NamedPruningMetrics {
+  string name = 1;
+  uint64 pruned = 2;
+  uint64 matched = 3;
+}
+
+message NamedRatio {
+  string name = 1;
+  uint64 part = 2;
+  uint64 total = 3;
+}
+
 message OperatorMetric {
   oneof metric {
     uint64 output_rows = 1;
@@ -277,6 +289,9 @@ message OperatorMetric {
     int64 start_timestamp = 9;
     int64 end_timestamp = 10;
     uint64 spilled_rows = 11;
+    uint64 output_bytes = 12;
+    NamedPruningMetrics pruning_metrics = 13;
+    NamedRatio ratio = 14;
   }
 }
 

ballista/core/proto/ballista.proto

@@ -167,6 +167,7 @@ message CreateExternalTableNode {
   datafusion_common.DfSchema schema = 4;
   repeated string table_partition_cols = 5;
   bool if_not_exists = 6;
+  bool or_replace = 15;
   bool temporary = 14;
   string definition = 7;
   repeated SortExprNodeCollection order_exprs = 10;
@@ -180,6 +181,7 @@ message PrepareNode {
   string name = 1;
   repeated datafusion_common.ArrowType data_types = 2;
   LogicalPlanNode input = 3;
+  repeated datafusion_common.Field fields = 4;
 }
 
 message CreateCatalogSchemaNode {
@@ -412,6 +414,8 @@ message Wildcard {
 message PlaceholderNode {
   string id = 1;
   datafusion_common.ArrowType data_type = 2;
+  optional bool nullable = 3;
+  map<string, string> metadata = 4;
 }
 
 message LogicalExprList {
@@ -517,6 +521,7 @@ message AggregateUDFExprNode {
   LogicalExprNode filter = 3;
   repeated SortExprNode order_by = 4;
   optional bytes fun_definition = 6;
+  optional NullTreatment null_treatment = 7;
 }
 
 message ScalarUDFExprNode {
@@ -537,6 +542,9 @@ message WindowExprNode {
   // repeated LogicalExprNode filter = 7;
   WindowFrame window_frame = 8;
   optional bytes fun_definition = 10;
+  optional NullTreatment null_treatment = 11;
+  bool distinct = 12;
+  LogicalExprNode filter = 13;
 }
 
 message BetweenNode {
@@ -621,6 +629,11 @@ message WindowFrameBound {
   datafusion_common.ScalarValue bound_value = 2;
 }
 
+enum NullTreatment {
+  RESPECT_NULLS = 0;
+  IGNORE_NULLS = 1;
+}
+
 ///////////////////////////////////////////////////////////////////////////////////////////////////
 // Arrow Data Types
 ///////////////////////////////////////////////////////////////////////////////////////////////////
@@ -1032,6 +1045,7 @@ message CsvScanExecNode {
     string comment = 6;
   }
   bool newlines_in_values = 7;
+  bool truncate_rows = 8;
 }
 
 message JsonScanExecNode {

ballista/core/proto/datafusion.proto

@@ -460,6 +460,7 @@ message CsvOptions {
   bytes double_quote = 15; // Indicates if quotes are doubled
   bytes newlines_in_values = 16; // Indicates if newlines are supported in values
   bytes terminator = 17; // Optional terminator character as a byte
+  bytes truncated_rows = 18; // Indicates if truncated rows are allowed
 }
 
 // Options controlling CSV format
@@ -580,6 +581,10 @@ message ParquetOptions {
   oneof coerce_int96_opt {
     string coerce_int96 = 32;
   }
+
+  oneof max_predicate_cache_size_opt {
+    uint64 max_predicate_cache_size = 33;
+  }
 }
 
 enum JoinSide {