Merge remote-tracking branch 'apache/master' into gauge

yjshen · yjshen · commit 9df9b38bc91e · 2022-01-26T22:01:46.000+08:00
diff --git a/README.md b/README.md
@@ -67,6 +67,7 @@ Here are some of the projects known to use DataFusion:
 - [ROAPI](https://github.com/roapi/roapi)
 - [Tensorbase](https://github.com/tensorbase/tensorbase)
 - [Squirtle](https://github.com/DSLAM-UMD/Squirtle)
+- [VegaFusion](https://vegafusion.io/) Server-side acceleration for the [Vega](https://vega.github.io/) visualization grammar
 
 (if you know of another project, please submit a PR to add a link!)
 
diff --git a/ballista/rust/core/src/execution_plans/shuffle_writer.rs b/ballista/rust/core/src/execution_plans/shuffle_writer.rs
@@ -265,11 +265,10 @@ impl ShuffleWriterExec {
                                 std::fs::create_dir_all(&path)?;
 
                                 path.push(format!("data-{}.arrow", input_partition));
-                                let path = path.to_str().unwrap();
-                                info!("Writing results to {}", path);
+                                info!("Writing results to {:?}", path);
 
                                 let mut writer =
-                                    IPCWriter::new(path, stream.schema().as_ref())?;
+                                    IPCWriter::new(&path, stream.schema().as_ref())?;
 
                                 writer.write(&output_batch)?;
                                 writers[output_partition] = Some(writer);
@@ -287,7 +286,7 @@ impl ShuffleWriterExec {
                         Some(w) => {
                             w.finish()?;
                             info!(
-                                "Finished writing shuffle partition {} at {}. Batches: {}. Rows: {}. Bytes: {}.",
+                                "Finished writing shuffle partition {} at {:?}. Batches: {}. Rows: {}. Bytes: {}.",
                                 i,
                                 w.path(),
                                 w.num_batches,
@@ -297,7 +296,7 @@ impl ShuffleWriterExec {
 
                             part_locs.push(ShuffleWritePartition {
                                 partition_id: i as u64,
-                                path: w.path().to_owned(),
+                                path: w.path().to_string_lossy().to_string(),
                                 num_batches: w.num_batches,
                                 num_rows: w.num_rows,
                                 num_bytes: w.num_bytes,
diff --git a/datafusion/src/execution/disk_manager.rs b/datafusion/src/execution/disk_manager.rs
@@ -20,14 +20,10 @@
 
 use crate::error::{DataFusionError, Result};
 use log::{debug, info};
-use rand::distributions::Alphanumeric;
 use rand::{thread_rng, Rng};
-use std::collections::hash_map::DefaultHasher;
-use std::fs::File;
-use std::hash::{Hash, Hasher};
-use std::path::{Path, PathBuf};
+use std::path::PathBuf;
 use std::sync::Arc;
-use tempfile::{Builder, TempDir};
+use tempfile::{Builder, NamedTempFile, TempDir};
 
 /// Configuration for temporary disk access
 #[derive(Debug, Clone)]
@@ -101,8 +97,8 @@ impl DiskManager {
         }
     }
 
-    /// Create a file in conf dirs in randomized manner and return the file path
-    pub fn create_tmp_file(&self) -> Result<String> {
+    /// Return a temporary file from a randomized choice in the configured locations
+    pub fn create_tmp_file(&self) -> Result<NamedTempFile> {
         create_tmp_file(&self.local_dirs)
     }
 }
@@ -120,34 +116,15 @@ fn create_local_dirs(local_dirs: Vec<PathBuf>) -> Result<Vec<TempDir>> {
         .collect()
 }
 
-fn get_file(file_name: &str, local_dirs: &[TempDir]) -> String {
-    let mut hasher = DefaultHasher::new();
-    file_name.hash(&mut hasher);
-    let hash = hasher.finish();
-    let dir = &local_dirs[hash.rem_euclid(local_dirs.len() as u64) as usize];
-    let mut path = PathBuf::new();
-    path.push(dir);
-    path.push(file_name);
-    path.to_str().unwrap().to_string()
-}
+fn create_tmp_file(local_dirs: &[TempDir]) -> Result<NamedTempFile> {
+    let dir_index = thread_rng().gen_range(0..local_dirs.len());
+    let dir = local_dirs.get(dir_index).ok_or_else(|| {
+        DataFusionError::Internal("No directories available to DiskManager".into())
+    })?;
 
-fn create_tmp_file(local_dirs: &[TempDir]) -> Result<String> {
-    let name = rand_name();
-    let mut path = get_file(&*name, local_dirs);
-    while Path::new(path.as_str()).exists() {
-        path = get_file(&rand_name(), local_dirs);
-    }
-    File::create(&path)?;
-    Ok(path)
-}
-
-/// Return a random string suitable for use as a database name
-fn rand_name() -> String {
-    thread_rng()
-        .sample_iter(&Alphanumeric)
-        .take(10)
-        .map(char::from)
-        .collect()
+    Builder::new()
+        .tempfile_in(dir)
+        .map_err(DataFusionError::IoError)
 }
 
 #[cfg(test)]
@@ -161,19 +138,28 @@ mod tests {
         let local_dir1 = TempDir::new()?;
         let local_dir2 = TempDir::new()?;
         let local_dir3 = TempDir::new()?;
-        let config = DiskManagerConfig::new_specified(vec![
-            local_dir1.path().into(),
-            local_dir2.path().into(),
-            local_dir3.path().into(),
-        ]);
+        let local_dirs = vec![local_dir1.path(), local_dir2.path(), local_dir3.path()];
+        let config = DiskManagerConfig::new_specified(
+            local_dirs.iter().map(|p| p.into()).collect(),
+        );
 
         let dm = DiskManager::try_new(config)?;
         let actual = dm.create_tmp_file()?;
-        let name = actual.rsplit_once(std::path::MAIN_SEPARATOR).unwrap().1;
 
-        let expected = get_file(name, &dm.local_dirs);
-        // file should be located in dir by it's name hash
-        assert_eq!(actual, expected);
+        // the file should be in one of the specified local directories
+        let found = local_dirs.iter().any(|p| {
+            actual
+                .path()
+                .ancestors()
+                .any(|candidate_path| *p == candidate_path)
+        });
+
+        assert!(
+            found,
+            "Can't find {:?} in specified local dirs: {:?}",
+            actual, local_dirs
+        );
+
         Ok(())
     }
 }
diff --git a/datafusion/src/physical_plan/common.rs b/datafusion/src/physical_plan/common.rs
@@ -33,6 +33,7 @@ use futures::{Future, SinkExt, Stream, StreamExt, TryStreamExt};
 use pin_project_lite::pin_project;
 use std::fs;
 use std::fs::{metadata, File};
+use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use std::task::{Context, Poll};
 use tokio::task::JoinHandle;
@@ -387,7 +388,7 @@ mod tests {
 /// Write in Arrow IPC format.
 pub struct IPCWriter {
     /// path
-    pub path: String,
+    pub path: PathBuf,
     /// Inner writer
     pub writer: FileWriter<File>,
     /// bathes written
@@ -400,18 +401,18 @@ pub struct IPCWriter {
 
 impl IPCWriter {
     /// Create new writer
-    pub fn new(path: &str, schema: &Schema) -> Result<Self> {
+    pub fn new(path: &Path, schema: &Schema) -> Result<Self> {
         let file = File::create(path).map_err(|e| {
             DataFusionError::Execution(format!(
-                "Failed to create partition file at {}: {:?}",
+                "Failed to create partition file at {:?}: {:?}",
                 path, e
             ))
         })?;
         Ok(Self {
             num_batches: 0,
             num_rows: 0,
             num_bytes: 0,
-            path: path.to_owned(),
+            path: path.into(),
             writer: FileWriter::try_new(file, schema)?,
         })
     }
@@ -432,7 +433,7 @@ impl IPCWriter {
     }
 
     /// Path write to
-    pub fn path(&self) -> &str {
+    pub fn path(&self) -> &Path {
         &self.path
     }
 }
diff --git a/datafusion/src/physical_plan/sorts/sort.rs b/datafusion/src/physical_plan/sorts/sort.rs
@@ -50,8 +50,10 @@ use std::fmt;
 use std::fmt::{Debug, Formatter};
 use std::fs::File;
 use std::io::BufReader;
+use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use tokio::sync::mpsc::{Receiver, Sender};
+use tempfile::NamedTempFile;
 use tokio::task;
 
 /// Sort arbitrary size of data to get a total order (may spill several times during sorting based on free memory available).
@@ -67,7 +69,7 @@ struct ExternalSorter {
     id: MemoryConsumerId,
     schema: SchemaRef,
     in_mem_batches: Mutex<Vec<RecordBatch>>,
-    spills: Mutex<Vec<String>>,
+    spills: Mutex<Vec<NamedTempFile>>,
     /// Sort expressions
     expr: Vec<PhysicalSortExpr>,
     runtime: Arc<RuntimeEnv>,
@@ -222,20 +224,19 @@ impl MemoryConsumer for ExternalSorter {
 
         let baseline_metrics = self.metrics.new_intermediate_baseline(partition);
 
-        let path = self.runtime.disk_manager.create_tmp_file()?;
+        let spillfile = self.runtime.disk_manager.create_tmp_file()?;
         let stream = in_mem_partial_sort(
             &mut *in_mem_batches,
             self.schema.clone(),
             &*self.expr,
             baseline_metrics,
         );
 
-        spill_partial_sorted_stream(&mut stream?, path.clone(), self.schema.clone())
-            .await?;
+        spill_partial_sorted_stream(&mut stream?, spillfile.path(), self.schema.clone()).await?;
         let mut spills = self.spills.lock().await;
         let used = self.inner_metrics.mem_used().set(0);
         self.inner_metrics.record_spill(used);
-        spills.push(path);
+        spills.push(spillfile);
         Ok(used)
     }
 
@@ -280,12 +281,12 @@ fn in_mem_partial_sort(
 
 async fn spill_partial_sorted_stream(
     in_mem_stream: &mut SendableRecordBatchStream,
-    path: String,
+    path: &Path,
     schema: SchemaRef,
 ) -> Result<()> {
     let (sender, receiver) = tokio::sync::mpsc::channel(2);
-    let path_clone = path.clone();
-    let handle = task::spawn_blocking(move || write_sorted(receiver, path_clone, schema));
+    let path: PathBuf = path.into();
+    let handle = task::spawn_blocking(move || write_sorted(receiver, path, schema));
     while let Some(item) = in_mem_stream.next().await {
         sender.send(item).await.ok();
     }
@@ -300,17 +301,16 @@ async fn spill_partial_sorted_stream(
 }
 
 fn read_spill_as_stream(
-    path: String,
+    path: NamedTempFile,
     schema: SchemaRef,
 ) -> Result<SendableRecordBatchStream> {
     let (sender, receiver): (
         Sender<ArrowResult<RecordBatch>>,
         Receiver<ArrowResult<RecordBatch>>,
     ) = tokio::sync::mpsc::channel(2);
-    let path_clone = path.clone();
     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);
+        if let Err(e) = read_spill(sender, path.path()) {
+            error!("Failure while reading spill file: {:?}. Error: {}", path, e);
         }
     });
     Ok(RecordBatchReceiverStream::create(
@@ -322,7 +322,7 @@ fn read_spill_as_stream(
 
 fn write_sorted(
     mut receiver: Receiver<ArrowResult<RecordBatch>>,
-    path: String,
+    path: PathBuf,
     schema: SchemaRef,
 ) -> Result<()> {
     let mut writer = IPCWriter::new(path.as_ref(), schema.as_ref())?;
@@ -337,7 +337,7 @@ fn write_sorted(
     Ok(())
 }
 
-fn read_spill(sender: Sender<ArrowResult<RecordBatch>>, path: String) -> Result<()> {
+fn read_spill(sender: Sender<ArrowResult<RecordBatch>>, path: &Path) -> Result<()> {
     let file = BufReader::new(File::open(&path)?);
     let reader = FileReader::try_new(file)?;
     for batch in reader {
