From 92a5b09237ed9e81ee352a076cc8887e70762f03 Mon Sep 17 00:00:00 2001 From: Sarah Yurick Date: Mon, 4 Dec 2023 13:45:01 -0800 Subject: [PATCH] handle single pq file --- src/sql/optimizer/dynamic_partition_pruning.rs | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/src/sql/optimizer/dynamic_partition_pruning.rs b/src/sql/optimizer/dynamic_partition_pruning.rs index d94359f3a..021eeff7c 100644 --- a/src/sql/optimizer/dynamic_partition_pruning.rs +++ b/src/sql/optimizer/dynamic_partition_pruning.rs @@ -484,11 +484,18 @@ fn read_table( field_string: String, tables: HashMap, ) -> Option> { - // Obtain filepaths to all relevant Parquet files, e.g., in a directory of Parquet files - let paths = fs::read_dir(tables.get(&table_string).unwrap().filepath.clone()).unwrap(); + let file_path = tables.get(&table_string).unwrap().filepath.clone(); + let paths: fs::ReadDir; let mut files = vec![]; - for path in paths { - files.push(path.unwrap().path().display().to_string()) + if fs::metadata(&file_path).map(|metadata| metadata.is_dir()).unwrap_or(false) { + // Obtain filepaths to all relevant Parquet files, e.g., in a directory of Parquet files + paths = fs::read_dir(&file_path).unwrap(); + for path in paths { + files.push(path.unwrap().path().display().to_string()) + } + } else { + // Obtain single Parquet file + files.push(file_path); } // Using the filepaths to the Parquet tables, obtain the schemas of the relevant tables