feat: datafusion_temporary_file()

Author: PokIsemaine

Cargo.lock

@@ -1036,7 +1036,6 @@ impl SessionStateBuilder {
                     .into_iter()
                     .map(|f| (f.name().to_string(), f)),
             );
-
         self
     }
 

datafusion/core/src/execution/session_state.rs

@@ -144,6 +144,42 @@ impl DiskManager {
                 .map_err(DataFusionError::IoError)?,
         })
     }
+
+    pub fn get_temporary_files(&self) -> Vec<(String, u64)> {
+        let mut files = Vec::new();
+
+        let guard = self.local_dirs.lock();
+        if let Some(local_dirs) = guard.as_ref() {
+            local_dirs
+                .iter()
+                .flat_map(|temp_dir| {
+                    std::fs::read_dir(temp_dir.path())
+                        .into_iter()
+                        .flat_map(|entries| {
+                            entries.flatten().filter_map(|entry| {
+                                let path = entry.path();
+                                entry
+                                    .metadata()
+                                    .ok()
+                                    .filter(|metadata| metadata.is_file())
+                                    .map(|metadata| {
+                                        (
+                                            path.to_string_lossy().to_string(),
+                                            metadata.len(),
+                                        )
+                                    })
+                            })
+                        })
+                })
+                .for_each(|file| files.push(file));
+        }
+
+        // test data
+        files.push(("/tmp/file1.tmp".to_string(), 1024));
+        files.push(("/tmp/file2.tmp".to_string(), 2048));
+
+        files
+    }
 }
 
 /// A wrapper around a [`NamedTempFile`] that also contains

datafusion/execution/src/disk_manager.rs

@@ -44,6 +44,8 @@ datafusion-catalog = { workspace = true }
 datafusion-common = { workspace = true }
 datafusion-expr = { workspace = true }
 datafusion-physical-plan = { workspace = true }
+datafusion-execution ={ workspace = true }
+datafusion-datasource = { workspace = true }
 parking_lot = { workspace = true }
 paste = "1.0.14"
 

datafusion/functions-table/Cargo.toml

@@ -25,13 +25,14 @@
 #![cfg_attr(not(test), deny(clippy::clone_on_ref_ptr))]
 
 pub mod generate_series;
+pub mod memory_profile;
 
 use datafusion_catalog::TableFunction;
 use std::sync::Arc;
 
 /// Returns all default table functions
 pub fn all_default_table_functions() -> Vec<Arc<TableFunction>> {
-    vec![generate_series(), range()]
+    vec![generate_series(), range(), datafusion_temporary_files()]
 }
 
 /// Creates a singleton instance of a table function
@@ -59,3 +60,7 @@ macro_rules! create_udtf_function {
 
 create_udtf_function!(generate_series::GenerateSeriesFunc, "generate_series");
 create_udtf_function!(generate_series::RangeFunc, "range");
+create_udtf_function!(
+    memory_profile::TemporaryFilesFunc,
+    "datafusion_temporary_files"
+);

datafusion/functions-table/src/lib.rs

@@ -0,0 +1,100 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{StringArray, UInt64Array};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use arrow::record_batch::RecordBatch;
+use async_trait::async_trait;
+use datafusion_catalog::{Session, TableFunctionImpl, TableProvider};
+use datafusion_common::Result;
+use datafusion_datasource::memory::MemorySourceConfig;
+use datafusion_expr::{Expr, TableType};
+use datafusion_physical_plan::ExecutionPlan;
+use std::sync::Arc;
+
+#[derive(Debug, Clone)]
+struct TemporaryFilesTable {
+    schema: SchemaRef,
+}
+
+#[async_trait]
+impl TableProvider for TemporaryFilesTable {
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn schema(&self) -> SchemaRef {
+        Arc::clone(&self.schema)
+    }
+
+    fn table_type(&self) -> TableType {
+        TableType::Base
+    }
+
+    async fn scan(
+        &self,
+        state: &dyn Session,
+        projection: Option<&Vec<usize>>,
+        _filters: &[Expr],
+        _limit: Option<usize>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let runtime_env = state.runtime_env();
+        let temporary_files = runtime_env.disk_manager.get_temporary_files();
+        let paths = temporary_files
+            .iter()
+            .map(|file| file.0.clone())
+            .collect::<Vec<_>>();
+        let sizes = temporary_files
+            .iter()
+            .map(|file| file.1)
+            .collect::<Vec<_>>();
+        let batches = vec![RecordBatch::try_new(
+            Arc::clone(&self.schema),
+            vec![
+                Arc::new(StringArray::from(paths)),
+                Arc::new(UInt64Array::from(sizes)),
+            ],
+        )?];
+
+        Ok(MemorySourceConfig::try_new_exec(
+            &[batches],
+            TableProvider::schema(self),
+            projection.cloned(),
+        )?)
+    }
+}
+
+/// A table function that returns temporary files with their paths and sizes
+#[derive(Debug)]
+pub struct TemporaryFilesFunc {}
+
+impl TableFunctionImpl for TemporaryFilesFunc {
+    fn call(&self, _exprs: &[Expr]) -> Result<Arc<dyn TableProvider>> {
+        // Create the schema for the table
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("path", DataType::Utf8, false),
+            Field::new("size", DataType::UInt64, false),
+        ]));
+
+        // Create a MemTable plan that returns the RecordBatch
+        let provider = TemporaryFilesTable {
+            schema,
+        };
+
+        Ok(Arc::new(provider))
+    }
+}

datafusion/functions-table/src/memory_profile.rs

@@ -303,3 +303,10 @@ SELECT range(1, t1.end) FROM range(3, 5) as t1(end)
 ----
 [1, 2, 3]
 [1, 2]
+
+# test
+query TI
+SELECT * FROM datafusion_temporary_files()
+----
+/tmp/file1.tmp 1024
+/tmp/file2.tmp 2048