Merge df-6.0-arrow2-0.8 and make it compile

Author: yjshen

datafusion/Cargo.toml

@@ -40,11 +40,7 @@ path = "src/lib.rs"
 [features]
 default = ["crypto_expressions", "regex_expressions", "unicode_expressions"]
 simd = ["arrow/simd"]
-<<<<<<< HEAD
 crypto_expressions = ["md-5", "sha2", "blake2", "blake3"]
-=======
-crypto_expressions = ["md-5", "sha2"]
->>>>>>> ExternalSortExec v1
 regex_expressions = ["regex"]
 unicode_expressions = ["unicode-segmentation"]
 # FIXME: add pyarrow support to arrow2 pyarrow = ["pyo3", "arrow/pyarrow"]

datafusion/src/error.rs

@@ -63,7 +63,7 @@ pub enum DataFusionError {
     Execution(String),
     /// This error is thrown when a consumer cannot acquire memory from the Memory Manager
     /// we can just cancel the execution of the partition.
-    OutOfMemory(String),
+    ResourcesExhausted(String),
 }
 
 impl DataFusionError {
@@ -132,8 +132,8 @@ impl Display for DataFusionError {
             DataFusionError::Execution(ref desc) => {
                 write!(f, "Execution error: {}", desc)
             }
-            DataFusionError::OutOfMemory(ref desc) => {
-                write!(f, "Out of memory error: {}", desc)
+            DataFusionError::ResourcesExhausted(ref desc) => {
+                write!(f, "Resources exhausted: {}", desc)
             }
         }
     }

datafusion/src/execution/context.rs

@@ -31,8 +31,6 @@ use crate::{
         },
         MemTable,
     },
-    execution::disk_manager::DiskManager,
-    execution::memory_management::MemoryManager,
     logical_plan::{PlanType, ToStringifiedPlan},
     optimizer::eliminate_limit::EliminateLimit,
     physical_optimizer::{

datafusion/src/execution/disk_manager.rs

@@ -34,7 +34,7 @@ pub struct DiskManager {
 
 impl DiskManager {
     /// Create local dirs inside user provided dirs through conf
-    pub fn new(conf_dirs: &Vec<String>) -> Result<Self> {
+    pub fn new(conf_dirs: &[String]) -> Result<Self> {
         Ok(Self {
             local_dirs: create_local_dirs(conf_dirs)?,
         })
@@ -55,9 +55,9 @@ impl DiskManager {
 }
 
 /// Setup local dirs by creating one new dir in each of the given dirs
-fn create_local_dirs(local_dir: &Vec<String>) -> Result<Vec<String>> {
+fn create_local_dirs(local_dir: &[String]) -> Result<Vec<String>> {
     local_dir
-        .into_iter()
+        .iter()
         .map(|root| create_directory(root, "datafusion"))
         .collect()
 }
@@ -82,7 +82,7 @@ fn create_directory(root: &str, prefix: &str) -> Result<String> {
     )))
 }
 
-fn get_file(file_name: &str, local_dirs: &Vec<String>) -> String {
+fn get_file(file_name: &str, local_dirs: &[String]) -> String {
     let mut hasher = DefaultHasher::new();
     file_name.hash(&mut hasher);
     let hash = hasher.finish();
@@ -93,7 +93,7 @@ fn get_file(file_name: &str, local_dirs: &Vec<String>) -> String {
     path.to_str().unwrap().to_string()
 }
 
-fn create_tmp_file(local_dirs: &Vec<String>) -> Result<String> {
+fn create_tmp_file(local_dirs: &[String]) -> Result<String> {
     let name = Uuid::new_v4().to_string();
     let mut path = get_file(&*name, local_dirs);
     while Path::new(path.as_str()).exists() {

datafusion/src/execution/memory_management/mod.rs

@@ -19,7 +19,8 @@
 
 pub mod allocation_strategist;
 
-use crate::error::DataFusionError::OutOfMemory;
+use std::cmp::Reverse;
+use crate::error::DataFusionError::ResourcesExhausted;
 use crate::error::{DataFusionError, Result};
 use crate::execution::memory_management::allocation_strategist::{
     DummyAllocationStrategist, FairStrategist, MemoryAllocationStrategist,
@@ -169,7 +170,7 @@ impl PartitionMemoryManager {
     }
 
     /// Try to acquire `required` of execution memory for the consumer and return the number of bytes
-    /// obtained, or return OutOfMemoryError if no enough memory avaiable even after possible spills.
+    /// obtained, or return ResourcesExhausted if no enough memory available even after possible spills.
     pub async fn acquire_exec_memory(
         &self,
         required: usize,
@@ -192,7 +193,7 @@ impl PartitionMemoryManager {
             for c in consumers.iter() {
                 all_consumers.push(c.1.clone());
             }
-            all_consumers.sort_by(|a, b| b.get_used().cmp(&a.get_used()));
+            all_consumers.sort_by_key(|b| Reverse(b.get_used()));
 
             for c in all_consumers.iter_mut() {
                 if c.id() == consumer_id {
@@ -235,7 +236,7 @@ impl PartitionMemoryManager {
         }
 
         if got < required {
-            return Err(OutOfMemory(format!(
+            return Err(ResourcesExhausted(format!(
                 "Unable to acquire {} bytes of memory, got {}",
                 required, got
             )));

datafusion/src/execution/mod.rs

@@ -19,7 +19,7 @@
 
 pub mod context;
 pub mod dataframe_impl;
-pub mod options;
 pub mod disk_manager;
 pub mod memory_management;
+pub mod options;
 pub mod runtime_env;

datafusion/src/execution/runtime_env.rs

@@ -21,6 +21,7 @@
 use crate::error::Result;
 use crate::execution::disk_manager::DiskManager;
 use crate::execution::memory_management::{MemoryConsumer, MemoryManager};
+use lazy_static::lazy_static;
 use std::sync::Arc;
 
 lazy_static! {
@@ -100,7 +101,7 @@ impl RuntimeConfig {
 
     /// Customize exec size
     pub fn with_local_dirs(mut self, local_dirs: Vec<String>) -> Self {
-        assert!(local_dirs.len() > 0);
+        assert!(!local_dirs.is_empty());
         self.local_dirs = local_dirs;
         self
     }

datafusion/src/physical_plan/common.rs

@@ -294,7 +294,7 @@ pub struct IPCWriterWrapper {
 impl IPCWriterWrapper {
     /// Create new writer
     pub fn new(path: &str, schema: &Schema) -> Result<Self> {
-        let file = File::create(path).map_err(|e| DataFusionError::IoError(e))?;
+        let file = File::create(path).map_err(DataFusionError::IoError)?;
         let buffer_writer = std::io::BufWriter::new(file);
         Ok(Self {
             num_batches: 0,

datafusion/src/physical_plan/mod.rs

@@ -624,11 +624,7 @@ pub mod distinct_expressions;
 pub mod empty;
 pub mod explain;
 pub mod expressions;
-<<<<<<< HEAD
 pub mod file_format;
-pub mod external_sort;
-=======
->>>>>>> move sorts together into submodule
 pub mod filter;
 pub mod functions;
 pub mod hash_aggregate;

datafusion/src/physical_plan/sorts/external_sort.rs

@@ -307,12 +307,16 @@ async fn read_spill_as_stream(
         TKReceiver<ArrowResult<RecordBatch>>,
     ) = tokio::sync::mpsc::channel(2);
     let path_clone = path.clone();
-    task::spawn_blocking(move || {
+    let join_handle = task::spawn_blocking(move || {
         if let Err(e) = read_spill(sender, path_clone) {
             error!("Failure while reading spill file: {}. Error: {}", path, e);
         }
     });
-    Ok(RecordBatchReceiverStream::create(&schema, receiver))
+    Ok(RecordBatchReceiverStream::create(
+        &schema,
+        receiver,
+        join_handle,
+    ))
 }
 
 pub(crate) async fn convert_stream_disk_based(
@@ -521,30 +525,41 @@ pub async fn external_sort(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::datasource::object_store::local::LocalFileSystem;
     use crate::physical_plan::coalesce_partitions::CoalescePartitionsExec;
     use crate::physical_plan::expressions::col;
     use crate::physical_plan::memory::MemoryExec;
-    use crate::physical_plan::sorts::SortOptions;
     use crate::physical_plan::{
         collect,
-        csv::{CsvExec, CsvReadOptions},
+        file_format::{CsvExec, PhysicalPlanConfig},
     };
     use crate::test;
+    use crate::test_util;
     use arrow::array::*;
+    use arrow::compute::sort::SortOptions;
     use arrow::datatypes::*;
 
     #[tokio::test]
     async fn test_sort() -> Result<()> {
-        let schema = test::aggr_test_schema();
+        let schema = test_util::aggr_test_schema();
         let partitions = 4;
-        let path = test::create_partitioned_csv("aggregate_test_100.csv", partitions)?;
-        let csv = CsvExec::try_new(
-            &path,
-            CsvReadOptions::new().schema(&schema),
-            None,
-            1024,
-            None,
-        )?;
+        let (_, files) =
+            test::create_partitioned_csv("aggregate_test_100.csv", partitions)?;
+
+        let csv = CsvExec::new(
+            PhysicalPlanConfig {
+                object_store: Arc::new(LocalFileSystem {}),
+                file_schema: Arc::clone(&schema),
+                file_groups: files,
+                statistics: Statistics::default(),
+                projection: None,
+                batch_size: 1024,
+                limit: None,
+                table_partition_cols: vec![],
+            },
+            true,
+            b',',
+        );
 
         let sort_exec = Arc::new(ExternalSortExec::try_new(
             vec![