Merge branch 'main' into aa/string_builder_char

Author: ajita-asthana

Cargo.lock

@@ -168,7 +168,7 @@ rstest = "0.25.0"
 serde_json = "1"
 sqlparser = { version = "0.55.0", features = ["visitor"] }
 tempfile = "3"
-tokio = { version = "1.44", features = ["macros", "rt", "sync"] }
+tokio = { version = "1.45", features = ["macros", "rt", "sync"] }
 url = "2.5.4"
 
 [profile.release]

Cargo.toml

@@ -109,29 +109,32 @@ impl RunOpt {
         };
 
         let mut benchmark_run = BenchmarkRun::new();
-        for query_id in query_range {
-            benchmark_run.start_new_case(&format!("Query {query_id}"));
-            let query_run = self.benchmark_query(query_id).await?;
-            for iter in query_run {
-                benchmark_run.write_iter(iter.elapsed, iter.row_count);
-            }
-        }
-        benchmark_run.maybe_write_json(self.output_path.as_ref())?;
-        Ok(())
-    }
-
-    async fn benchmark_query(&self, query_id: usize) -> Result<Vec<QueryResult>> {
         let mut config = self
             .common
             .config()?
             .with_collect_statistics(!self.disable_statistics);
         config.options_mut().optimizer.prefer_hash_join = self.prefer_hash_join;
         let rt_builder = self.common.runtime_env_builder()?;
         let ctx = SessionContext::new_with_config_rt(config, rt_builder.build_arc()?);
-
         // register tables
         self.register_tables(&ctx).await?;
 
+        for query_id in query_range {
+            benchmark_run.start_new_case(&format!("Query {query_id}"));
+            let query_run = self.benchmark_query(query_id, &ctx).await?;
+            for iter in query_run {
+                benchmark_run.write_iter(iter.elapsed, iter.row_count);
+            }
+        }
+        benchmark_run.maybe_write_json(self.output_path.as_ref())?;
+        Ok(())
+    }
+
+    async fn benchmark_query(
+        &self,
+        query_id: usize,
+        ctx: &SessionContext,
+    ) -> Result<Vec<QueryResult>> {
         let mut millis = vec![];
         // run benchmark
         let mut query_results = vec![];
@@ -146,14 +149,14 @@ impl RunOpt {
             if query_id == 15 {
                 for (n, query) in sql.iter().enumerate() {
                     if n == 1 {
-                        result = self.execute_query(&ctx, query).await?;
+                        result = self.execute_query(ctx, query).await?;
                     } else {
-                        self.execute_query(&ctx, query).await?;
+                        self.execute_query(ctx, query).await?;
                     }
                 }
             } else {
                 for query in sql {
-                    result = self.execute_query(&ctx, query).await?;
+                    result = self.execute_query(ctx, query).await?;
                 }
             }
 

benchmarks/src/tpch/run.rs

@@ -19,7 +19,7 @@ use std::{num::NonZeroUsize, sync::Arc};
 
 use datafusion::{
     execution::{
-        disk_manager::DiskManagerConfig,
+        disk_manager::DiskManagerBuilder,
         memory_pool::{FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool},
         runtime_env::RuntimeEnvBuilder,
     },
@@ -110,7 +110,7 @@ impl CommonOpt {
             };
             rt_builder = rt_builder
                 .with_memory_pool(pool)
-                .with_disk_manager(DiskManagerConfig::NewOs);
+                .with_disk_manager_builder(DiskManagerBuilder::default());
         }
         Ok(rt_builder)
     }

benchmarks/src/util/options.rs

@@ -39,7 +39,7 @@ arrow = { workspace = true }
 async-trait = { workspace = true }
 aws-config = "1.6.2"
 aws-credential-types = "1.2.0"
-clap = { version = "4.5.37", features = ["derive", "cargo"] }
+clap = { version = "4.5.39", features = ["derive", "cargo"] }
 datafusion = { workspace = true, features = [
     "avro",
     "crypto_expressions",
@@ -60,7 +60,7 @@ object_store = { workspace = true, features = ["aws", "gcp", "http"] }
 parking_lot = { workspace = true }
 parquet = { workspace = true, default-features = false }
 regex = { workspace = true }
-rustyline = "15.0"
+rustyline = "16.0"
 tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread", "sync", "parking_lot", "signal"] }
 url = { workspace = true }
 

datafusion-cli/Cargo.toml

@@ -28,7 +28,6 @@ use datafusion::execution::memory_pool::{
     FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool,
 };
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
-use datafusion::execution::DiskManager;
 use datafusion::prelude::SessionContext;
 use datafusion_cli::catalog::DynamicObjectStoreCatalog;
 use datafusion_cli::functions::ParquetMetadataFunc;
@@ -43,7 +42,7 @@ use datafusion_cli::{
 use clap::Parser;
 use datafusion::common::config_err;
 use datafusion::config::ConfigOptions;
-use datafusion::execution::disk_manager::DiskManagerConfig;
+use datafusion::execution::disk_manager::{DiskManagerBuilder, DiskManagerMode};
 use mimalloc::MiMalloc;
 
 #[global_allocator]
@@ -200,15 +199,10 @@ async fn main_inner() -> Result<()> {
 
     // set disk limit
     if let Some(disk_limit) = args.disk_limit {
-        let mut disk_manager = DiskManager::try_new(DiskManagerConfig::NewOs)?;
-
-        DiskManager::set_arc_max_temp_directory_size(
-            &mut disk_manager,
-            disk_limit.try_into().unwrap(),
-        )?;
-
-        let disk_config = DiskManagerConfig::new_existing(disk_manager);
-        rt_builder = rt_builder.with_disk_manager(disk_config);
+        let builder = DiskManagerBuilder::default()
+            .with_mode(DiskManagerMode::OsTmpDirectory)
+            .with_max_temp_directory_size(disk_limit.try_into().unwrap());
+        rt_builder = rt_builder.with_disk_manager_builder(builder);
     }
 
     let runtime_env = rt_builder.build_arc()?;

datafusion-cli/src/main.rs

@@ -77,7 +77,7 @@ tonic = "0.12.1"
 tracing = { version = "0.1" }
 tracing-subscriber = { version = "0.3" }
 url = { workspace = true }
-uuid = "1.16"
+uuid = "1.17"
 
 [target.'cfg(not(target_os = "windows"))'.dev-dependencies]
 nix = { version = "0.30.1", features = ["fs"] }

datafusion-examples/Cargo.toml

@@ -25,6 +25,7 @@ use arrow::array::{
 };
 use arrow::datatypes::{ArrowNativeTypeOp, ArrowPrimitiveType, Float64Type, UInt32Type};
 use arrow::record_batch::RecordBatch;
+use arrow_schema::FieldRef;
 use datafusion::common::{cast::as_float64_array, ScalarValue};
 use datafusion::error::Result;
 use datafusion::logical_expr::{
@@ -92,10 +93,10 @@ impl AggregateUDFImpl for GeoMeanUdaf {
     }
 
     /// This is the description of the state. accumulator's state() must match the types here.
-    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<Field>> {
+    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
         Ok(vec![
-            Field::new("prod", args.return_type().clone(), true),
-            Field::new("n", DataType::UInt32, true),
+            Field::new("prod", args.return_type().clone(), true).into(),
+            Field::new("n", DataType::UInt32, true).into(),
         ])
     }
 
@@ -401,7 +402,7 @@ impl AggregateUDFImpl for SimplifiedGeoMeanUdaf {
         unimplemented!("should not be invoked")
     }
 
-    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<Field>> {
+    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
         unimplemented!("should not be invoked")
     }
 

datafusion-examples/examples/advanced_udaf.rs

@@ -23,6 +23,7 @@ use arrow::{
     array::{ArrayRef, AsArray, Float64Array},
     datatypes::Float64Type,
 };
+use arrow_schema::FieldRef;
 use datafusion::common::ScalarValue;
 use datafusion::error::Result;
 use datafusion::functions_aggregate::average::avg_udaf;
@@ -87,8 +88,8 @@ impl WindowUDFImpl for SmoothItUdf {
         Ok(Box::new(MyPartitionEvaluator::new()))
     }
 
-    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
-        Ok(Field::new(field_args.name(), DataType::Float64, true))
+    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<FieldRef> {
+        Ok(Field::new(field_args.name(), DataType::Float64, true).into())
     }
 }
 
@@ -205,8 +206,8 @@ impl WindowUDFImpl for SimplifySmoothItUdf {
         Some(Box::new(simplify))
     }
 
-    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
-        Ok(Field::new(field_args.name(), DataType::Float64, true))
+    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<FieldRef> {
+        Ok(Field::new(field_args.name(), DataType::Float64, true).into())
     }
 }