DiskManager: include the description in file names

datafusion/execution/src/disk_manager.rs

@@ -373,6 +373,7 @@ impl DiskManager {
             parent_temp_dir: Arc::clone(&local_dirs[dir_index]),
             tempfile: Arc::new(
                 Builder::new()
+                    .prefix(&escape_for_temp_file_name_prefix(request_description))
                     .tempfile_in(local_dirs[dir_index].as_ref())
                     .map_err(DataFusionError::IoError)?,
             ),
@@ -585,6 +586,37 @@ impl SpillFile for RefCountedTempFile {
         }))
     }
 }
+
+/// Escape a string to be safe to use in a temp file name prefix.
+///
+/// All non-alphanumeric characters are replaced by `_` and characters are put in lowercase.
+/// A `.tmp.` is added at the beginning and a last `.` is set at the end.
+///
+/// For example "This kind.of /input" becomes ".tmp.this_kind_of_input."
+///
+/// The output is also truncated to fit in 128 bytes.
+fn escape_for_temp_file_name_prefix(text: &str) -> String {
+    let mut output = String::with_capacity(text.len() + 6);
+    output.push_str(".tmp.");
+    let mut add_separator = false;
+    for c in text.chars() {
+        if c.is_alphanumeric() {
+            if add_separator && !output.is_empty() {
+                output.push('_');
+                add_separator = false;
+            }
+            output.extend(c.to_lowercase());
+        } else {
+            add_separator = true;
+        }
+        if output.len() >= 123 {
+            break; // We limit the output at 128 bytes and keep the space for an extra UTF-8 char and the final dot
+        }
+    }
+    output.push('.');
+    output
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -622,7 +654,7 @@ mod tests {
     }
 
     #[test]
-    fn file_in_right_dir() -> Result<()> {
+    fn file_in_right_dir_and_with_description() -> Result<()> {
         let local_dir1 = TempDir::new()?;
         let local_dir2 = TempDir::new()?;
         let local_dir3 = TempDir::new()?;
@@ -636,11 +668,23 @@ mod tests {
         );
 
         assert!(dm.tmp_files_enabled());
-        let actual = dm.create_tmp_file("Testing")?;
+        let actual = dm.create_tmp_file("Testing with./spécial")?;
 
         // the file should be in one of the specified local directories
         assert_path_in_dirs(actual.path().unwrap(), local_dirs.into_iter());
 
+        let file_name = actual
+            .path()
+            .unwrap()
+            .file_name()
+            .unwrap()
+            .to_str()
+            .unwrap();
+        assert!(
+            file_name.starts_with(".tmp.testing_with_spécial."),
+            "{file_name}"
+        );
+
         Ok(())
     }
 
@@ -1175,4 +1219,14 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_escape_for_temp_file_name_prefix() {
+        assert_eq!(escape_for_temp_file_name_prefix("foo"), ".tmp.foo.");
+        assert_eq!(escape_for_temp_file_name_prefix("fooé \nt"), ".tmp.fooé_t.");
+        assert_eq!(
+            escape_for_temp_file_name_prefix(&"a".repeat(200)),
+            format!(".tmp.{}.", "a".repeat(118))
+        );
+    }
 }