feat: csv/tsv/ndjson support querying file meta data. (#17512)

* feat: csv/tsv/ndjson support querying file meta data.

Author: youngsofun

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

@@ -44,6 +44,7 @@ pub struct StageTableInfo {
     pub is_select: bool,
     pub copy_into_location_options: CopyIntoLocationOptions,
     pub copy_into_table_options: CopyIntoTableOptions,
+    pub stage_root: String,
 }
 
 impl StageTableInfo {

src/query/catalog/src/plan/datasource/datasource_plan.rs

@@ -20,6 +20,7 @@ use databend_common_expression::Scalar;
 use databend_common_expression::TableSchemaRef;
 
 use crate::plan::datasource::datasource_info::DataSourceInfo;
+use crate::plan::InternalColumn;
 use crate::plan::PartStatistics;
 use crate::plan::Partitions;
 use crate::plan::PushDownInfo;
@@ -38,7 +39,7 @@ pub struct DataSourcePlan {
 
     pub tbl_args: Option<TableArgs>,
     pub push_downs: Option<PushDownInfo>,
-    pub query_internal_columns: bool,
+    pub internal_columns: Option<BTreeMap<FieldIndex, InternalColumn>>,
     pub base_block_ids: Option<Scalar>,
     // used for recluster to update stream columns
     pub update_stream_columns: bool,

src/query/catalog/src/plan/internal_column.rs

@@ -37,6 +37,8 @@ use databend_common_expression::Value;
 use databend_common_expression::BASE_BLOCK_IDS_COLUMN_ID;
 use databend_common_expression::BASE_ROW_ID_COLUMN_ID;
 use databend_common_expression::BLOCK_NAME_COLUMN_ID;
+use databend_common_expression::FILENAME_COLUMN_ID;
+use databend_common_expression::FILE_ROW_NUMBER_COLUMN_ID;
 use databend_common_expression::ROW_ID_COLUMN_ID;
 use databend_common_expression::SEARCH_MATCHED_COLUMN_ID;
 use databend_common_expression::SEARCH_SCORE_COLUMN_ID;
@@ -141,6 +143,9 @@ pub enum InternalColumnType {
     // search columns
     SearchMatched,
     SearchScore,
+
+    FileName,
+    FileRowNumber,
 }
 
 #[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq, Eq)]
@@ -176,6 +181,8 @@ impl InternalColumn {
             )),
             InternalColumnType::SearchMatched => TableDataType::Boolean,
             InternalColumnType::SearchScore => TableDataType::Number(NumberDataType::Float32),
+            InternalColumnType::FileName => TableDataType::String,
+            InternalColumnType::FileRowNumber => TableDataType::Number(NumberDataType::UInt64),
         }
     }
 
@@ -198,6 +205,8 @@ impl InternalColumn {
             InternalColumnType::BaseBlockIds => BASE_BLOCK_IDS_COLUMN_ID,
             InternalColumnType::SearchMatched => SEARCH_MATCHED_COLUMN_ID,
             InternalColumnType::SearchScore => SEARCH_SCORE_COLUMN_ID,
+            InternalColumnType::FileName => FILENAME_COLUMN_ID,
+            InternalColumnType::FileRowNumber => FILE_ROW_NUMBER_COLUMN_ID,
         }
     }
 
@@ -315,6 +324,9 @@ impl InternalColumn {
                     Value::Column(Float32Type::from_data(scores)),
                 )
             }
+            InternalColumnType::FileName | InternalColumnType::FileRowNumber => {
+                todo!("generate_column_values not support for file related")
+            }
         }
     }
 }

src/query/expression/src/schema.rs

@@ -67,11 +67,16 @@ pub const CHANGE_ROW_ID_COL_NAME: &str = "change$row_id";
 
 pub const PREDICATE_COLUMN_NAME: &str = "_predicate";
 
+pub const FILENAME_COLUMN_NAME: &str = "metadata$filename";
+pub const FILE_ROW_NUMBER_COLUMN_NAME: &str = "metadata$file_row_number";
+
 // stream column id.
 pub const ORIGIN_BLOCK_ROW_NUM_COLUMN_ID: u32 = u32::MAX - 10;
 pub const ORIGIN_BLOCK_ID_COLUMN_ID: u32 = u32::MAX - 11;
 pub const ORIGIN_VERSION_COLUMN_ID: u32 = u32::MAX - 12;
 pub const ROW_VERSION_COLUMN_ID: u32 = u32::MAX - 13;
+pub const FILENAME_COLUMN_ID: u32 = u32::MAX - 14;
+pub const FILE_ROW_NUMBER_COLUMN_ID: u32 = u32::MAX - 15;
 // stream column name.
 pub const ORIGIN_VERSION_COL_NAME: &str = "_origin_version";
 pub const ORIGIN_BLOCK_ID_COL_NAME: &str = "_origin_block_id";
@@ -98,12 +103,15 @@ pub static INTERNAL_COLUMNS: LazyLock<HashSet<&'static str>> = LazyLock::new(||
         ORIGIN_BLOCK_ID_COL_NAME,
         ORIGIN_BLOCK_ROW_NUM_COL_NAME,
         ROW_VERSION_COL_NAME,
+        FILENAME_COLUMN_NAME,
+        FILE_ROW_NUMBER_COLUMN_NAME,
     ])
 });
 
 #[inline]
 pub fn is_internal_column_id(column_id: ColumnId) -> bool {
     column_id >= SEARCH_SCORE_COLUMN_ID
+        || (FILE_ROW_NUMBER_COLUMN_ID..=FILENAME_COLUMN_ID).contains(&column_id)
 }
 
 #[inline]

src/query/service/src/interpreters/interpreter_copy_into_location.rs

@@ -113,6 +113,7 @@ impl CopyIntoLocationInterpreter {
                 default_values: None,
                 copy_into_location_options: options.clone(),
                 copy_into_table_options: Default::default(),
+                stage_root: "".to_string(),
             },
         }));
 

src/query/service/src/pipelines/builders/builder_recluster.rs

@@ -83,7 +83,7 @@ impl PipelineBuilder {
                     description,
                     tbl_args: table.table_args(),
                     push_downs: None,
-                    query_internal_columns: false,
+                    internal_columns: None,
                     base_block_ids: None,
                     update_stream_columns: table.change_tracking_enabled(),
                     data_mask_policy: None,

src/query/service/src/sessions/query_ctx.rs

@@ -91,6 +91,7 @@ use databend_common_pipeline_core::InputError;
 use databend_common_pipeline_core::LockGuard;
 use databend_common_settings::Settings;
 use databend_common_sql::IndexType;
+use databend_common_storage::init_stage_operator;
 use databend_common_storage::CopyStatus;
 use databend_common_storage::DataOperator;
 use databend_common_storage::FileStatus;
@@ -1514,6 +1515,14 @@ impl TableContext for QueryContext {
         max_column_position: usize,
         case_sensitive: bool,
     ) -> Result<Arc<dyn Table>> {
+        let operator = init_stage_operator(&stage_info)?;
+        let info = operator.info();
+        let stage_root = format!("{}{}", info.name(), info.root());
+        let stage_root = if stage_root.ends_with('/') {
+            stage_root
+        } else {
+            format!("{}/", stage_root)
+        };
         match stage_info.file_format_params {
             FileFormatParams::Parquet(..) => {
                 let mut read_options = ParquetReadOptions::default();
@@ -1553,6 +1562,7 @@ impl TableContext for QueryContext {
                     default_values: None,
                     copy_into_location_options: Default::default(),
                     copy_into_table_options: Default::default(),
+                    stage_root,
                 };
                 OrcTable::try_create(info).await
             }
@@ -1571,6 +1581,7 @@ impl TableContext for QueryContext {
                     default_values: None,
                     copy_into_location_options: Default::default(),
                     copy_into_table_options: Default::default(),
+                    stage_root,
                 };
                 StageTable::try_create(info)
             }
@@ -1607,6 +1618,7 @@ impl TableContext for QueryContext {
                     default_values: None,
                     copy_into_location_options: Default::default(),
                     copy_into_table_options: Default::default(),
+                    stage_root,
                 };
                 StageTable::try_create(info)
             }

src/query/sql/src/executor/table_read_plan.rs

@@ -272,7 +272,7 @@ impl ToReadDataSourcePlan for dyn Table {
             description,
             tbl_args: self.table_args(),
             push_downs,
-            query_internal_columns: internal_columns.is_some(),
+            internal_columns,
             base_block_ids,
             update_stream_columns,
             data_mask_policy,

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

@@ -225,6 +225,7 @@ impl Binder {
                 default_values,
                 copy_into_location_options: Default::default(),
                 copy_into_table_options: stmt.options.clone(),
+                stage_root: "".to_string(),
             },
             values_consts: vec![],
             required_source_schema: required_values_schema.clone(),
@@ -405,6 +406,7 @@ impl Binder {
                 default_values: Some(default_values),
                 copy_into_location_options: Default::default(),
                 copy_into_table_options: options,
+                stage_root: "".to_string(),
             },
             write_mode,
             query: None,
@@ -482,6 +484,7 @@ impl Binder {
 
         // rewrite async function and udf
         s_expr = self.rewrite_udf(&mut from_context, s_expr)?;
+        s_expr = self.add_internal_column_into_expr(&mut from_context, s_expr)?;
 
         let mut output_context = BindContext::new();
         output_context.parent = from_context.parent;

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

@@ -20,6 +20,8 @@ use databend_common_catalog::plan::InternalColumnType;
 use databend_common_expression::BASE_BLOCK_IDS_COL_NAME;
 use databend_common_expression::BASE_ROW_ID_COL_NAME;
 use databend_common_expression::BLOCK_NAME_COL_NAME;
+use databend_common_expression::FILENAME_COLUMN_NAME;
+use databend_common_expression::FILE_ROW_NUMBER_COLUMN_NAME;
 use databend_common_expression::ROW_ID_COL_NAME;
 use databend_common_expression::SEARCH_MATCHED_COL_NAME;
 use databend_common_expression::SEARCH_SCORE_COL_NAME;
@@ -77,6 +79,19 @@ impl InternalColumnFactory {
             InternalColumn::new(SEARCH_SCORE_COL_NAME, InternalColumnType::SearchScore),
         );
 
+        internal_columns.insert(
+            FILENAME_COLUMN_NAME.to_string(),
+            InternalColumn::new(FILENAME_COLUMN_NAME, InternalColumnType::FileName),
+        );
+
+        internal_columns.insert(
+            FILE_ROW_NUMBER_COLUMN_NAME.to_string(),
+            InternalColumn::new(
+                FILE_ROW_NUMBER_COLUMN_NAME,
+                InternalColumnType::FileRowNumber,
+            ),
+        );
+
         InternalColumnFactory { internal_columns }
     }