feat: support specifying columns to include in the "Attach table" statement (#17442)

* feat: support specifying columns to include in the "Attach table" statement

* update ut golden file of case parser::test_statement

* logic test

* clean up

* tweak test case 02_0004_attach_table

* adjust code comments

* show create table

* tweak case 02_0004_attach_table

Author: dantengsky

src/query/ast/src/ast/statements/table.rs

@@ -249,6 +249,7 @@ pub struct AttachTableStmt {
     pub catalog: Option<Identifier>,
     pub database: Option<Identifier>,
     pub table: Identifier,
+    pub columns_opt: Option<Vec<Identifier>>,
     pub uri_location: UriLocation,
 }
 
@@ -263,6 +264,12 @@ impl Display for AttachTableStmt {
                 .chain(Some(&self.table)),
         )?;
 
+        if let Some(cols) = &self.columns_opt {
+            write!(f, " (")?;
+            write_comma_separated_list(f, cols.iter())?;
+            write!(f, ")")?;
+        }
+
         write!(f, " {}", self.uri_location)?;
 
         Ok(())

src/query/ast/src/parser/statement.rs

@@ -869,13 +869,15 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
 
     let attach_table = map(
         rule! {
-            ATTACH ~ TABLE ~ #dot_separated_idents_1_to_3 ~ #uri_location
+            ATTACH ~ TABLE ~ #dot_separated_idents_1_to_3 ~ ("(" ~ #comma_separated_list1(ident) ~ ")")? ~ #uri_location
         },
-        |(_, _, (catalog, database, table), uri_location)| {
+        |(_, _, (catalog, database, table), columns_opt, uri_location)| {
+            let columns_opt = columns_opt.map(|(_, v, _)| v);
             Statement::AttachTable(AttachTableStmt {
                 catalog,
                 database,
                 table,
+                columns_opt,
                 uri_location,
             })
         },

src/query/ast/tests/it/testdata/stmt.txt

@@ -23705,6 +23705,7 @@ AttachTable(
             quote: None,
             ident_type: None,
         },
+        columns_opt: None,
         uri_location: UriLocation {
             protocol: "s3",
             name: "a",

src/query/ee/src/attach_table/handler.rs

@@ -15,6 +15,9 @@
 use std::sync::Arc;
 
 use databend_common_base::base::GlobalInstance;
+use databend_common_exception::ErrorCode;
+use databend_common_exception::Result;
+use databend_common_expression::TableSchema;
 use databend_common_meta_app::schema::CreateTableReq;
 use databend_common_meta_app::schema::TableMeta;
 use databend_common_meta_app::schema::TableNameIdent;
@@ -23,6 +26,7 @@ use databend_common_sql::plans::CreateTablePlan;
 use databend_common_storage::check_operator;
 use databend_common_storage::init_operator;
 use databend_common_storages_fuse::io::MetaReaders;
+use databend_common_storages_fuse::FUSE_OPT_KEY_ATTACH_COLUMN_IDS;
 use databend_common_storages_fuse::FUSE_TBL_LAST_SNAPSHOT_HINT;
 use databend_enterprise_attach_table::AttachTableHandler;
 use databend_enterprise_attach_table::AttachTableHandlerWrapper;
@@ -71,9 +75,11 @@ impl AttachTableHandler for RealAttachTableHandler {
             number_of_blocks: Some(snapshot.summary.block_count),
         };
 
+        let attach_table_schema = Self::gen_schema(&plan, &snapshot)?;
+
         let field_comments = vec!["".to_string(); snapshot.schema.num_fields()];
         let table_meta = TableMeta {
-            schema: Arc::new(snapshot.schema.clone()),
+            schema: Arc::new(attach_table_schema),
             engine: plan.engine.to_string(),
             storage_params: plan.storage_params.clone(),
             options,
@@ -105,4 +111,63 @@ impl RealAttachTableHandler {
         GlobalInstance::set(Arc::new(wrapper));
         Ok(())
     }
+
+    fn gen_schema(
+        plan: &&CreateTablePlan,
+        base_table_snapshot: &Arc<TableSnapshot>,
+    ) -> Result<TableSchema> {
+        let schema = if let Some(attached_columns) = &plan.attached_columns {
+            // Columns to include are specified, let's check them
+            let base_table_schema = &base_table_snapshot.schema;
+            let mut fields_to_attach = Vec::with_capacity(attached_columns.len());
+
+            // The ids of columns being included
+            let mut field_ids_to_include = Vec::with_capacity(attached_columns.len());
+
+            // Columns that do not exist in the table being attached to, if any
+            let mut invalid_cols = vec![];
+            for field in attached_columns {
+                match base_table_schema.field_with_name(&field.name) {
+                    Ok(f) => {
+                        field_ids_to_include.push(f.column_id);
+                        fields_to_attach.push(f.clone())
+                    }
+                    Err(_) => invalid_cols.push(field.name.as_str()),
+                }
+            }
+            if !invalid_cols.is_empty() {
+                return Err(ErrorCode::InvalidArgument(format!(
+                    "Columns [{}] do not exist in the table being attached to",
+                    invalid_cols.join(",")
+                )));
+            }
+
+            let new_table_schema_metadata = if !field_ids_to_include.is_empty() {
+                // If columns to include are specified explicitly, their ids should
+                // be kept in the metadata of TableSchema.
+                let ids = field_ids_to_include
+                    .iter()
+                    .map(|id| format!("{id}"))
+                    .collect::<Vec<_>>()
+                    .join(",");
+                let mut v = base_table_schema.metadata.clone();
+                v.insert(FUSE_OPT_KEY_ATTACH_COLUMN_IDS.to_owned(), ids);
+                v
+            } else {
+                base_table_schema.metadata.clone()
+            };
+
+            TableSchema {
+                fields: fields_to_attach,
+                metadata: new_table_schema_metadata,
+                next_column_id: base_table_schema.next_column_id,
+            }
+        } else {
+            // If columns are not specified, use all the fields of table being attached to,
+            // in this case, no schema meta of key FUSE_OPT_KEY_ATTACH_COLUMN_IDS will be kept.
+            base_table_snapshot.schema.clone()
+        };
+
+        Ok(schema)
+    }
 }

src/query/ee/tests/it/inverted_index/pruning.rs

@@ -122,6 +122,7 @@ async fn test_block_pruner() -> Result<()> {
         as_select: None,
         cluster_key: None,
         inverted_indexes: None,
+        attached_columns: None,
     };
 
     let interpreter = CreateTableInterpreter::try_create(ctx.clone(), create_table_plan)?;

src/query/service/src/interpreters/interpreter_table_show_create.rs

@@ -27,6 +27,7 @@ use databend_common_expression::DataBlock;
 use databend_common_expression::Scalar;
 use databend_common_expression::Value;
 use databend_common_sql::plans::ShowCreateTablePlan;
+use databend_common_storages_fuse::FUSE_OPT_KEY_ATTACH_COLUMN_IDS;
 use databend_common_storages_stream::stream_table::StreamTable;
 use databend_common_storages_stream::stream_table::STREAM_ENGINE;
 use databend_common_storages_view::view_table::QUERY;
@@ -37,6 +38,7 @@ use databend_storages_common_table_meta::table::OPT_KEY_CLUSTER_TYPE;
 use databend_storages_common_table_meta::table::OPT_KEY_STORAGE_PREFIX;
 use databend_storages_common_table_meta::table::OPT_KEY_TABLE_ATTACHED_DATA_URI;
 use databend_storages_common_table_meta::table::OPT_KEY_TEMP_PREFIX;
+use itertools::Itertools;
 
 use crate::interpreters::Interpreter;
 use crate::pipelines::PipelineBuildResult;
@@ -321,15 +323,35 @@ impl ShowCreateTableInterpreter {
     }
 
     fn show_attach_table_query(table: &dyn Table, database: &str) -> String {
-        // TODO table that attached before this PR, could not show location properly
+        // Note: Tables that attached before this PR #13403, could not show location properly
         let location_not_available = "N/A".to_string();
         let table_data_location = table
             .options()
             .get(OPT_KEY_TABLE_ATTACHED_DATA_URI)
             .unwrap_or(&location_not_available);
 
+        let table_info = table.get_table_info();
+
+        let mut include_cols = "".to_string();
+        if table_info
+            .meta
+            .schema
+            .metadata
+            .contains_key(FUSE_OPT_KEY_ATTACH_COLUMN_IDS)
+        {
+            let cols = table_info
+                .meta
+                .schema
+                .fields
+                .iter()
+                .map(|f| &f.name)
+                .join(",");
+            include_cols = format!("({cols}) ");
+        }
+
         format!(
-            "ATTACH TABLE `{}`.`{}` {}",
+            "ATTACH TABLE {}`{}`.`{}` {}",
+            include_cols,
             database,
             table.name(),
             table_data_location,

src/query/service/src/pipelines/processors/transforms/transform_recursive_cte_source.rs

@@ -306,6 +306,7 @@ async fn create_memory_table_for_cte_scan(
                 cluster_key: None,
                 as_select: None,
                 inverted_indexes: None,
+                attached_columns: None,
             };
             let create_table_interpreter =
                 CreateTableInterpreter::try_create(ctx.clone(), create_table_plan)?;

src/query/service/src/test_kits/fixture.rs

@@ -348,6 +348,7 @@ impl TestFixture {
             as_select: None,
             cluster_key: Some("(id)".to_string()),
             inverted_indexes: None,
+            attached_columns: None,
         }
     }
 
@@ -372,6 +373,7 @@ impl TestFixture {
             as_select: None,
             cluster_key: None,
             inverted_indexes: None,
+            attached_columns: None,
         }
     }
 
@@ -407,6 +409,7 @@ impl TestFixture {
             as_select: None,
             cluster_key: None,
             inverted_indexes: None,
+            attached_columns: None,
         }
     }
 
@@ -442,6 +445,7 @@ impl TestFixture {
             as_select: None,
             cluster_key: None,
             inverted_indexes: None,
+            attached_columns: None,
         }
     }
 
@@ -486,6 +490,7 @@ impl TestFixture {
             as_select: None,
             cluster_key: None,
             inverted_indexes: None,
+            attached_columns: None,
         }
     }
 

src/query/service/tests/it/sql/planner/optimizer/agg_index_query_rewrite.rs

@@ -84,6 +84,7 @@ fn create_table_plan(fixture: &TestFixture, format: &str) -> CreateTablePlan {
         as_select: None,
         cluster_key: None,
         inverted_indexes: None,
+        attached_columns: None,
     }
 }
 

src/query/service/tests/it/storages/fuse/operations/clustering.rs

@@ -53,6 +53,7 @@ async fn test_fuse_alter_table_cluster_key() -> databend_common_exception::Resul
         as_select: None,
         cluster_key: None,
         inverted_indexes: None,
+        attached_columns: None,
     };
 
     // create test table