feat: information.temporary_files

Author: PokIsemaine

datafusion/catalog/src/information_schema.rs

@@ -31,7 +31,7 @@ use async_trait::async_trait;
 use datafusion_common::config::{ConfigEntry, ConfigOptions};
 use datafusion_common::error::Result;
 use datafusion_common::DataFusionError;
-use datafusion_execution::TaskContext;
+use datafusion_execution::{disk_manager, TaskContext};
 use datafusion_expr::{AggregateUDF, ScalarUDF, Signature, TypeSignature, WindowUDF};
 use datafusion_expr::{TableType, Volatility};
 use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
@@ -49,6 +49,7 @@ pub(crate) const DF_SETTINGS: &str = "df_settings";
 pub(crate) const SCHEMATA: &str = "schemata";
 pub(crate) const ROUTINES: &str = "routines";
 pub(crate) const PARAMETERS: &str = "parameters";
+pub(crate) const TEMPORARY_FILES: &str = "temporary_files";
 
 /// All information schema tables
 pub const INFORMATION_SCHEMA_TABLES: &[&str] = &[
@@ -59,6 +60,7 @@ pub const INFORMATION_SCHEMA_TABLES: &[&str] = &[
     SCHEMATA,
     ROUTINES,
     PARAMETERS,
+    TEMPORARY_FILES,
 ];
 
 /// Implements the `information_schema` virtual schema and tables
@@ -397,6 +399,19 @@ impl InformationSchemaConfig {
             TypeSignature::Variadic(_) | TypeSignature::VariadicAny
         )
     }
+
+    fn make_temporary_files(
+        &self,
+        disk_manager: &disk_manager::DiskManager,
+        builder: &mut InformationSchemaTemporaryFilesBuilder,
+    ) -> Result<(), DataFusionError> {
+        let temporary_files = disk_manager.get_temporary_files();
+
+        for (path, size) in temporary_files {
+            builder.add_file(&path, size);
+        }
+        Ok(())
+    }
 }
 
 /// get the arguments and return types of a UDF
@@ -496,6 +511,7 @@ impl SchemaProvider for InformationSchemaProvider {
             SCHEMATA => Arc::new(InformationSchemata::new(config)),
             ROUTINES => Arc::new(InformationSchemaRoutines::new(config)),
             PARAMETERS => Arc::new(InformationSchemaParameters::new(config)),
+            TEMPORARY_FILES => Arc::new(InformationSchemaTemporaryFiles::new(config)),
             _ => return Ok(None),
         };
 
@@ -1348,3 +1364,70 @@ impl PartitionStream for InformationSchemaParameters {
         ))
     }
 }
+
+#[derive(Debug)]
+struct InformationSchemaTemporaryFiles {
+    schema: SchemaRef,
+    config: InformationSchemaConfig,
+}
+
+impl InformationSchemaTemporaryFiles {
+    fn new(config: InformationSchemaConfig) -> Self {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("path", DataType::Utf8, false),
+            Field::new("size", DataType::UInt64, false),
+        ]));
+
+        Self { schema, config }
+    }
+
+    fn builder(&self) -> InformationSchemaTemporaryFilesBuilder {
+        InformationSchemaTemporaryFilesBuilder {
+            paths: StringBuilder::new(),
+            sizes: UInt64Builder::new(),
+            schema: Arc::clone(&self.schema),
+        }
+    }
+}
+
+impl PartitionStream for InformationSchemaTemporaryFiles {
+    fn schema(&self) -> &SchemaRef {
+        &self.schema
+    }
+
+    fn execute(&self, ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
+        let mut builder = self.builder();
+        let config = self.config.clone();
+        Box::pin(RecordBatchStreamAdapter::new(
+            Arc::clone(&self.schema),
+            futures::stream::once(async move {
+                config.make_temporary_files(&ctx.runtime_env().disk_manager, &mut builder)?;
+                Ok(builder.finish())
+            }),
+        ))
+    }
+}
+
+struct InformationSchemaTemporaryFilesBuilder {
+    schema: SchemaRef,
+    paths: StringBuilder,
+    sizes: UInt64Builder,
+}
+
+impl InformationSchemaTemporaryFilesBuilder {
+    fn add_file(&mut self, path: &str, size: u64) {
+        self.paths.append_value(path);
+        self.sizes.append_value(size);
+    }
+
+    fn finish(&mut self) -> RecordBatch {
+        RecordBatch::try_new(
+            Arc::clone(&self.schema),
+            vec![
+                Arc::new(self.paths.finish()),
+                Arc::new(self.sizes.finish()),
+            ],
+        )
+        .unwrap()
+    }
+}

datafusion/execution/src/disk_manager.rs

@@ -144,6 +144,31 @@ 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/sqllogictest/test_files/information_schema_tmp.slt

@@ -0,0 +1,41 @@
+# 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.
+
+
+# Verify the information schema does not exit by default
+statement error DataFusion error: Error during planning: table 'datafusion.information_schema.tables' not found
+SELECT * from information_schema.tables
+
+# Verify the information schema does not exit by default
+statement error DataFusion error: Error during planning: table 'datafusion.information_schema.schemata' not found
+SELECT * from information_schema.schemata
+
+statement error DataFusion error: Error during planning: SHOW \[VARIABLE\] is not supported unless information_schema is enabled
+show all
+
+# Turn it on
+
+# expect that the queries now work
+statement ok
+set datafusion.catalog.information_schema = true;
+
+# Check information_schema.temporary_files
+query TI
+SELECT * FROM information_schema.temporary_files;
+----
+/tmp/file1.tmp 1024
+/tmp/file2.tmp 2048