feat: support copy into table with transform.

youngsofun · youngsofun · commit d0f94823e78f · 2023-03-15T00:38:20.000+08:00
diff --git a/src/query/service/src/interpreters/interpreter_copy.rs b/src/query/service/src/interpreters/interpreter_copy.rs
@@ -41,6 +41,7 @@ use tracing::info;
 use crate::interpreters::common::append2table;
 use crate::interpreters::Interpreter;
 use crate::interpreters::SelectInterpreter;
+use crate::pipelines::processors::TransformCastSchema;
 use crate::pipelines::processors::TransformLimit;
 use crate::pipelines::PipelineBuildResult;
 use crate::sessions::QueryContext;
@@ -152,6 +153,76 @@ impl CopyInterpreter {
         }
     }
 
+    #[allow(clippy::too_many_arguments)]
+    async fn build_copy_into_table_with_transform_pipeline(
+        &self,
+        catalog_name: &str,
+        database_name: &str,
+        table_name: &str,
+        query: &Plan,
+        stage_info: StageInfo,
+        all_source_file_infos: Vec<StageFileInfo>,
+        need_copy_file_infos: Vec<StageFileInfo>,
+    ) -> Result<PipelineBuildResult> {
+        let start = Instant::now();
+        let ctx = self.ctx.clone();
+        let (mut build_res, source_schema) = self.build_query(query).await?;
+        let to_table = ctx
+            .get_table(catalog_name, database_name, table_name)
+            .await?;
+
+        let dst_schema = Arc::new(to_table.schema().into());
+        if source_schema != dst_schema {
+            let func_ctx = ctx.get_function_context()?;
+            build_res.main_pipeline.add_transform(
+                |transform_input_port, transform_output_port| {
+                    TransformCastSchema::try_create(
+                        transform_input_port,
+                        transform_output_port,
+                        source_schema.clone(),
+                        dst_schema.clone(),
+                        func_ctx,
+                    )
+                },
+            )?;
+        }
+
+        // Build append data pipeline.
+        to_table.append_data(
+            ctx.clone(),
+            &mut build_res.main_pipeline,
+            AppendMode::Copy,
+            false,
+        )?;
+
+        let database_name = database_name.to_string();
+        let catalog_name = catalog_name.to_string();
+        let table_name = table_name.to_string();
+        build_res.main_pipeline.set_on_finished(move |may_error| {
+            if may_error.is_none() {
+                CopyInterpreter::commit_copy_into_table(
+                    ctx.clone(),
+                    to_table,
+                    stage_info,
+                    all_source_file_infos,
+                    need_copy_file_infos,
+                    catalog_name,
+                    database_name,
+                    table_name,
+                )?;
+                // Status.
+                {
+                    info!("all copy finished, elapsed:{}", start.elapsed().as_secs());
+                }
+                Ok(())
+            } else {
+                Err(may_error.as_ref().unwrap().clone())
+            }
+        });
+
+        Ok(build_res)
+    }
+
     #[allow(clippy::too_many_arguments)]
     async fn build_copy_into_table_pipeline(
         &self,
@@ -422,6 +493,27 @@ impl Interpreter for CopyInterpreter {
                     other
                 ))),
             },
+            CopyPlan::IntoTableWithTransform {
+                catalog_name,
+                database_name,
+                table_name,
+                stage_info,
+                from,
+                all_source_file_infos,
+                need_copy_file_infos,
+                ..
+            } => {
+                self.build_copy_into_table_with_transform_pipeline(
+                    catalog_name,
+                    database_name,
+                    table_name,
+                    from,
+                    *stage_info.clone(),
+                    all_source_file_infos.clone(),
+                    need_copy_file_infos.clone(),
+                )
+                .await
+            }
             CopyPlan::IntoStage {
                 stage, from, path, ..
             } => self.build_copy_into_stage_pipeline(stage, path, from).await,
diff --git a/src/query/sql/src/planner/binder/copy.rs b/src/query/sql/src/planner/binder/copy.rs
@@ -15,6 +15,7 @@
 use std::collections::BTreeMap;
 use std::str::FromStr;
 use std::sync::Arc;
+use std::time::Instant;
 
 use common_ast::ast::CopyStmt;
 use common_ast::ast::CopyUnit;
@@ -40,9 +41,13 @@ use common_exception::ErrorCode;
 use common_exception::Result;
 use common_meta_app::principal::OnErrorMode;
 use common_meta_app::principal::StageInfo;
+use common_storage::init_stage_operator;
+use common_storage::StageFileInfo;
+use common_storage::StageFileStatus;
 use common_storage::StageFilesInfo;
 use common_users::UserApiProvider;
 use tracing::debug;
+use tracing::info;
 
 use crate::binder::location::parse_uri_location;
 use crate::binder::Binder;
@@ -181,6 +186,26 @@ impl<'a> Binder {
                 self.bind_copy_from_query_into_uri(bind_context, stmt, query, &mut ul)
                     .await
             }
+            (
+                CopyUnit::Query(query),
+                CopyUnit::Table {
+                    catalog,
+                    database,
+                    table,
+                },
+            ) => {
+                let (catalog_name, database_name, table_name) =
+                    self.normalize_object_identifier_triple(catalog, database, table);
+                self.bind_copy_from_query_into_table(
+                    bind_context,
+                    stmt,
+                    query,
+                    &catalog_name,
+                    &database_name,
+                    &table_name,
+                )
+                .await
+            }
             (src, dst) => Err(ErrorCode::SyntaxException(format!(
                 "COPY INTO <{}> FROM <{}> is invalid",
                 dst.target(),
@@ -470,6 +495,148 @@ impl<'a> Binder {
         })))
     }
 
+    /// Bind COPY INFO <table> FROM <query>
+    async fn bind_copy_from_query_into_table(
+        &mut self,
+        bind_context: &BindContext,
+        stmt: &CopyStmt,
+        src_query: &Query,
+        dst_catalog_name: &str,
+        dst_database_name: &str,
+        dst_table_name: &str,
+    ) -> Result<Plan> {
+        // Validation mode.
+        let validation_mode = ValidationMode::from_str(stmt.validation_mode.as_str())
+            .map_err(ErrorCode::SyntaxException)?;
+
+        // dst
+        let dst_table = self
+            .ctx
+            .get_table(dst_catalog_name, dst_database_name, dst_table_name)
+            .await?;
+
+        // src
+        let (select_list, location, alias) = check_transform_query(src_query)?;
+        if matches!(location, FileLocation::Uri(_)) {
+            // todo!(youngsofun): need to refactor parser
+            return Err(ErrorCode::SyntaxException(
+                "copy into table from uri with transform not supported yet",
+            ));
+        }
+
+        let (mut stage_info, path) =
+            parse_file_location(&self.ctx, location, BTreeMap::new()).await?;
+        self.apply_stage_options(stmt, &mut stage_info).await?;
+        let files_info = StageFilesInfo {
+            path,
+            pattern: stmt.pattern.clone(),
+            files: stmt.files.clone(),
+        };
+
+        let start = Instant::now();
+        {
+            let status = "begin to list files";
+            self.ctx.set_status_info(status);
+            info!(status);
+        }
+
+        let operator = init_stage_operator(&stage_info)?;
+        let files = if operator.info().can_blocking() {
+            files_info.blocking_list(&operator, false)
+        } else {
+            files_info.list(&operator, false).await
+        }?;
+
+        let mut all_source_file_infos = files
+            .into_iter()
+            .map(|file_with_meta| StageFileInfo::new(file_with_meta.path, &file_with_meta.metadata))
+            .collect::<Vec<_>>();
+
+        info!("end to list files: {}", all_source_file_infos.len());
+
+        if !stmt.force {
+            // Status.
+            {
+                let status = "begin to color copied files";
+                self.ctx.set_status_info(status);
+                info!(status);
+            }
+
+            all_source_file_infos = self
+                .ctx
+                .color_copied_files(
+                    dst_catalog_name,
+                    dst_database_name,
+                    dst_table_name,
+                    all_source_file_infos,
+                )
+                .await?;
+
+            info!("end to color copied files: {}", all_source_file_infos.len());
+        }
+
+        let mut need_copy_file_infos = vec![];
+        for file in &all_source_file_infos {
+            if file.status == StageFileStatus::NeedCopy {
+                need_copy_file_infos.push(file.clone());
+            }
+        }
+
+        info!(
+            "copy: read all files finished, all:{}, need copy:{}, elapsed:{}",
+            all_source_file_infos.len(),
+            need_copy_file_infos.len(),
+            start.elapsed().as_secs()
+        );
+
+        if need_copy_file_infos.is_empty() {
+            return Err(ErrorCode::EmptyData("no file need to copy"));
+        }
+
+        let (s_expr, mut from_context) = self
+            .bind_stage_table(
+                bind_context,
+                stage_info.clone(),
+                files_info,
+                alias,
+                Some(need_copy_file_infos.clone()),
+            )
+            .await?;
+
+        // Generate a analyzed select list with from context
+        let select_list = self
+            .normalize_select_list(&from_context, select_list)
+            .await?;
+        let (scalar_items, projections) = self.analyze_projection(&select_list)?;
+        let s_expr =
+            self.bind_projection(&mut from_context, &projections, &scalar_items, s_expr)?;
+        let mut output_context = BindContext::new();
+        output_context.parent = from_context.parent;
+        output_context.columns = from_context.columns;
+
+        let query_plan = Plan::Query {
+            s_expr: Box::new(s_expr),
+            metadata: self.metadata.clone(),
+            bind_context: Box::new(output_context),
+            rewrite_kind: None,
+            ignore_result: false,
+            formatted_ast: None,
+        };
+
+        Ok(Plan::Copy(Box::new(CopyPlan::IntoTableWithTransform {
+            catalog_name: dst_catalog_name.to_string(),
+            database_name: dst_database_name.to_string(),
+            table_name: dst_table_name.to_string(),
+            table_id: dst_table.get_id(),
+            schema: dst_table.schema(),
+            from: Box::new(query_plan),
+            stage_info: Box::new(stage_info),
+            all_source_file_infos,
+            need_copy_file_infos,
+            validation_mode,
+        })))
+    }
+
     async fn apply_stage_options(&mut self, stmt: &CopyStmt, stage: &mut StageInfo) -> Result<()> {
         if !stmt.file_format.is_empty() {
             stage.file_format_options = self.try_resolve_file_format(&stmt.file_format).await?;
diff --git a/src/query/sql/src/planner/plans/copy.rs b/src/query/sql/src/planner/plans/copy.rs
@@ -20,6 +20,7 @@ use common_catalog::plan::DataSourcePlan;
 use common_expression::TableSchemaRef;
 use common_meta_app::principal::StageInfo;
 use common_meta_types::MetaId;
+use common_storage::StageFileInfo;
 
 use crate::plans::Plan;
 
@@ -65,6 +66,18 @@ pub enum CopyPlan {
         from: Box<DataSourcePlan>,
         force: bool,
     },
+    IntoTableWithTransform {
+        catalog_name: String,
+        database_name: String,
+        table_name: String,
+        table_id: MetaId,
+        schema: TableSchemaRef,
+        stage_info: Box<StageInfo>,
+        validation_mode: ValidationMode,
+        from: Box<Plan>,
+        all_source_file_infos: Vec<StageFileInfo>,
+        need_copy_file_infos: Vec<StageFileInfo>,
+    },
     IntoStage {
         stage: Box<StageInfo>,
         path: String,
@@ -90,6 +103,17 @@ impl Debug for CopyPlan {
                 write!(f, ", from: {from:?}")?;
                 write!(f, " force: {force}")?;
             }
+            CopyPlan::IntoTableWithTransform {
+                database_name,
+                table_name,
+                from,
+                validation_mode,
+                ..
+            } => {
+                write!(f, "Copy into {database_name:}.{table_name:}")?;
+                write!(f, ", validation_mode: {validation_mode:?}")?;
+                write!(f, ", from: {from:?}")?;
+            }
             CopyPlan::IntoStage {
                 stage,
                 path,
