feat: ParquetTable support injecting file list.

Author: youngsofun

src/query/catalog/src/plan/datasource/datasource_info/parquet.rs

@@ -18,6 +18,7 @@ use common_arrow::arrow::datatypes::Schema as ArrowSchema;
 use common_expression::TableSchema;
 use common_meta_app::principal::StageInfo;
 use common_meta_app::schema::TableInfo;
+use common_storage::StageFileInfo;
 use common_storage::StageFilesInfo;
 
 use crate::plan::datasource::datasource_info::parquet_read_options::ParquetReadOptions;
@@ -30,6 +31,7 @@ pub struct ParquetTableInfo {
 
     pub table_info: TableInfo,
     pub arrow_schema: ArrowSchema,
+    pub files_to_read: Option<Vec<StageFileInfo>>,
 }
 
 impl ParquetTableInfo {

src/query/sql/src/planner/binder/table.rs

@@ -46,6 +46,7 @@ use common_functions::scalars::BUILTIN_FUNCTIONS;
 use common_meta_app::principal::StageFileFormatType;
 use common_meta_app::principal::StageInfo;
 use common_storage::DataOperator;
+use common_storage::StageFileInfo;
 use common_storage::StageFilesInfo;
 use common_storages_parquet::ParquetTable;
 use common_storages_result_cache::ResultCacheMetaManager;
@@ -380,7 +381,7 @@ impl Binder {
                     pattern: options.pattern.clone(),
                     files: options.files.clone(),
                 };
-                self.bind_stage_table(bind_context, stage_info, files_info, alias)
+                self.bind_stage_table(bind_context, stage_info, files_info, alias, None)
                     .await
             }
         }
@@ -392,14 +393,17 @@ impl Binder {
         stage_info: StageInfo,
         files_info: StageFilesInfo,
         alias: &Option<TableAlias>,
+        files_to_copy: Option<Vec<StageFileInfo>>,
     ) -> Result<(SExpr, BindContext)> {
         if matches!(
             stage_info.file_format_options.format,
             StageFileFormatType::Parquet
         ) {
             let read_options = ParquetReadOptions::default();
 
-            let table = ParquetTable::create(stage_info.clone(), files_info, read_options).await?;
+            let table =
+                ParquetTable::create(stage_info.clone(), files_info, read_options, files_to_copy)
+                    .await?;
 
             let table_alias_name = if let Some(table_alias) = alias {
                 Some(normalize_identifier(&table_alias.name, &self.name_resolution_ctx).name)

src/query/storages/parquet/src/parquet_table/blocking.rs

@@ -21,6 +21,7 @@ use common_catalog::table::Table;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_meta_app::principal::StageInfo;
+use common_storage::StageFileInfo;
 use common_storage::StageFilesInfo;
 use opendal::Operator;
 
@@ -33,9 +34,14 @@ impl ParquetTable {
         read_options: ParquetReadOptions,
         stage_info: StageInfo,
         files_info: StageFilesInfo,
+        files_to_read: Option<Vec<StageFileInfo>>,
     ) -> Result<Arc<dyn Table>> {
-        let first_file = files_info.blocking_first_file(&operator)?;
-        let arrow_schema = Self::blocking_prepare_metas(&first_file.path, operator.clone())?;
+        let first_file = match &files_to_read {
+            Some(files) => files[0].path.clone(),
+            None => files_info.blocking_first_file(&operator)?.path,
+        };
+
+        let arrow_schema = Self::blocking_prepare_metas(&first_file, operator.clone())?;
 
         let table_info = create_parquet_table_info(arrow_schema.clone());
 
@@ -46,6 +52,7 @@ impl ParquetTable {
             read_options,
             stage_info,
             files_info,
+            files_to_read,
         }))
     }
 

src/query/storages/parquet/src/parquet_table/non_blocking.rs

@@ -22,6 +22,7 @@ use common_exception::ErrorCode;
 use common_exception::Result;
 use common_meta_app::principal::StageInfo;
 use common_storage::init_stage_operator;
+use common_storage::StageFileInfo;
 use common_storage::StageFilesInfo;
 use opendal::Operator;
 
@@ -33,14 +34,24 @@ impl ParquetTable {
         stage_info: StageInfo,
         files_info: StageFilesInfo,
         read_options: ParquetReadOptions,
+        files_to_read: Option<Vec<StageFileInfo>>,
     ) -> Result<Arc<dyn Table>> {
         let operator = init_stage_operator(&stage_info)?;
         if operator.info().can_blocking() {
-            return Self::blocking_create(operator, read_options, stage_info, files_info);
+            return Self::blocking_create(
+                operator,
+                read_options,
+                stage_info,
+                files_info,
+                files_to_read,
+            );
         }
-        let first_file = files_info.first_file(&operator).await?;
+        let first_file = match &files_to_read {
+            Some(files) => files[0].path.clone(),
+            None => files_info.first_file(&operator).await?.path.clone(),
+        };
 
-        let arrow_schema = Self::prepare_metas(&first_file.path, operator.clone()).await?;
+        let arrow_schema = Self::prepare_metas(&first_file, operator.clone()).await?;
 
         let table_info = create_parquet_table_info(arrow_schema.clone());
 
@@ -51,6 +62,7 @@ impl ParquetTable {
             read_options,
             stage_info,
             files_info,
+            files_to_read,
         }))
     }
 

src/query/storages/parquet/src/parquet_table/partition.rs

@@ -107,14 +107,20 @@ impl ParquetTable {
             None
         };
 
-        let file_locations = if self.operator.info().can_blocking() {
-            self.files_info.blocking_list(&self.operator, false)
-        } else {
-            self.files_info.list(&self.operator, false).await
-        }?
-        .into_iter()
-        .map(|f| (f.path, f.metadata.content_length()))
-        .collect::<Vec<_>>();
+        let file_locations = match &self.files_to_read {
+            Some(files) => files
+                .iter()
+                .map(|f| (f.path.clone(), f.size))
+                .collect::<Vec<_>>(),
+            None => if self.operator.info().can_blocking() {
+                self.files_info.blocking_list(&self.operator, false)
+            } else {
+                self.files_info.list(&self.operator, false).await
+            }?
+            .into_iter()
+            .map(|f| (f.path, f.metadata.content_length()))
+            .collect::<Vec<_>>(),
+        };
 
         let pruner = PartitionPruner {
             schema,

src/query/storages/parquet/src/parquet_table/table.rs

@@ -38,6 +38,7 @@ use common_meta_app::schema::TableInfo;
 use common_meta_app::schema::TableMeta;
 use common_pipeline_core::Pipeline;
 use common_storage::init_stage_operator;
+use common_storage::StageFileInfo;
 use common_storage::StageFilesInfo;
 use opendal::Operator;
 
@@ -50,6 +51,7 @@ pub struct ParquetTable {
 
     pub(super) table_info: TableInfo,
     pub(super) arrow_schema: ArrowSchema,
+    pub(super) files_to_read: Option<Vec<StageFileInfo>>,
 }
 
 impl ParquetTable {
@@ -63,6 +65,7 @@ impl ParquetTable {
             read_options: info.read_options,
             stage_info: info.stage_info.clone(),
             files_info: info.files_info.clone(),
+            files_to_read: info.files_to_read.clone(),
         }))
     }
 }
@@ -100,6 +103,7 @@ impl Table for ParquetTable {
             read_options: self.read_options,
             stage_info: self.stage_info.clone(),
             files_info: self.files_info.clone(),
+            files_to_read: self.files_to_read.clone(),
         })
     }