feat: support async function for merge not match clause (#17941)

* add logic test

* rm unused code

* update

* fix

* skip empty

* modify test

* seperate cast schema

* modify test

* modify test

Author: SkyFan2002

src/query/pipeline/transforms/src/processors/transforms/transform_pipeline_helper.rs

@@ -145,6 +145,15 @@ pub trait TransformPipelineHelper {
     where
         F: Fn() -> Result<R>,
         R: Transform + 'static;
+
+    fn try_create_async_transform_pipeline_builder_with_len<F, R>(
+        &mut self,
+        f: F,
+        transform_len: usize,
+    ) -> Result<TransformPipeBuilder>
+    where
+        F: Fn() -> Result<R>,
+        R: AsyncTransform + 'static;
 }
 
 impl TransformPipelineHelper for Pipeline {
@@ -179,4 +188,25 @@ impl TransformPipelineHelper for Pipeline {
             transform_len,
         )
     }
+
+    fn try_create_async_transform_pipeline_builder_with_len<F, R>(
+        &mut self,
+        f: F,
+        transform_len: usize,
+    ) -> Result<TransformPipeBuilder>
+    where
+        F: Fn() -> Result<R>,
+        R: AsyncTransform + 'static,
+    {
+        self.add_transform_with_specified_len(
+            |input, output| {
+                Ok(ProcessorPtr::create(AsyncTransformer::create(
+                    input,
+                    output,
+                    f()?,
+                )))
+            },
+            transform_len,
+        )
+    }
 }

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

@@ -12,10 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use databend_common_base::base::tokio::sync::Semaphore;
 use databend_common_catalog::table::Table;
+use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_expression::BlockThresholds;
 use databend_common_expression::DataSchema;
@@ -34,11 +36,18 @@ use databend_common_pipeline_transforms::processors::TransformPipelineHelper;
 use databend_common_sql::binder::MutationStrategy;
 use databend_common_sql::executor::physical_plans::Mutation;
 use databend_common_sql::executor::physical_plans::MutationKind;
+use databend_common_sql::DefaultExprBinder;
 use databend_common_storages_fuse::operations::TransformSerializeBlock;
 use databend_common_storages_fuse::operations::UnMatchedExprs;
 use databend_common_storages_fuse::FuseTable;
 
+use crate::pipelines::processors::transforms::build_cast_exprs;
+use crate::pipelines::processors::transforms::build_expression_transform;
+use crate::pipelines::processors::transforms::AsyncFunctionBranch;
+use crate::pipelines::processors::transforms::CastSchemaBranch;
 use crate::pipelines::processors::transforms::TransformAddComputedColumns;
+use crate::pipelines::processors::transforms::TransformBranchedAsyncFunction;
+use crate::pipelines::processors::transforms::TransformBranchedCastSchema;
 use crate::pipelines::processors::transforms::TransformResortAddOnWithoutSourceSchema;
 use crate::pipelines::PipelineBuilder;
 
@@ -136,16 +145,110 @@ impl PipelineBuilder {
 
         // fill default columns
         let table_default_schema = &table.schema_with_stream().remove_computed_fields();
+        let default_schema: DataSchemaRef = Arc::new(table_default_schema.into());
+
+        let mut expression_transforms = Vec::with_capacity(unmatched.len());
+        let mut data_schemas = HashMap::with_capacity(unmatched.len());
+        let mut trigger_non_null_errors = Vec::with_capacity(unmatched.len());
+        let mut async_function_branches = HashMap::with_capacity(unmatched.len());
+        let mut cast_schema_branches = HashMap::with_capacity(unmatched.len());
+        for (idx, item) in unmatched.iter().enumerate() {
+            let mut input_schema = item.0.clone();
+            let mut default_expr_binder = DefaultExprBinder::try_new(self.ctx.clone())?;
+            if let Some((async_funcs, new_default_schema, new_default_schema_no_cast)) =
+                default_expr_binder
+                    .split_async_default_exprs(input_schema.clone(), default_schema.clone())?
+            {
+                async_function_branches.insert(idx, AsyncFunctionBranch {
+                    async_func_descs: async_funcs,
+                });
+
+                if new_default_schema != new_default_schema_no_cast {
+                    cast_schema_branches.insert(idx, CastSchemaBranch {
+                        to_schema: new_default_schema.clone(),
+                        from_schema: new_default_schema_no_cast.clone(),
+                        exprs: build_cast_exprs(
+                            new_default_schema_no_cast.clone(),
+                            new_default_schema.clone(),
+                        )?,
+                    });
+                }
+                // update input_schema, which is used in `TransformResortAddOnWithoutSourceSchema`
+                input_schema = new_default_schema;
+            }
+
+            data_schemas.insert(idx, input_schema.clone());
+            match build_expression_transform(
+                input_schema,
+                default_schema.clone(),
+                tbl.clone(),
+                self.ctx.clone(),
+            ) {
+                Ok(expression_transform) => {
+                    expression_transforms.push(Some(expression_transform));
+                    trigger_non_null_errors.push(None);
+                }
+                Err(err) => {
+                    if err.code() != ErrorCode::BAD_ARGUMENTS {
+                        return Err(err);
+                    }
+
+                    expression_transforms.push(None);
+                    trigger_non_null_errors.push(Some(err));
+                }
+            };
+        }
+
+        if !async_function_branches.is_empty() {
+            let branches = Arc::new(async_function_branches);
+            let mut builder = self
+                .main_pipeline
+                .try_create_async_transform_pipeline_builder_with_len(
+                    || {
+                        Ok(TransformBranchedAsyncFunction {
+                            ctx: self.ctx.clone(),
+                            branches: branches.clone(),
+                        })
+                    },
+                    transform_len,
+                )?;
+            if need_match {
+                builder.add_items_prepend(vec![create_dummy_item()]);
+            }
+            self.main_pipeline.add_pipe(builder.finalize());
+        }
+
+        if !cast_schema_branches.is_empty() {
+            let branches = Arc::new(cast_schema_branches);
+            let mut builder = self
+                .main_pipeline
+                .try_create_transform_pipeline_builder_with_len(
+                    || {
+                        Ok(TransformBranchedCastSchema {
+                            ctx: self.ctx.clone(),
+                            branches: branches.clone(),
+                        })
+                    },
+                    transform_len,
+                )?;
+            if need_match {
+                builder.add_items_prepend(vec![create_dummy_item()]);
+            }
+            self.main_pipeline.add_pipe(builder.finalize());
+        }
+
         let mut builder = self
             .main_pipeline
             .try_create_transform_pipeline_builder_with_len(
                 || {
                     TransformResortAddOnWithoutSourceSchema::try_new(
                         self.ctx.clone(),
                         Arc::new(DataSchema::from(table_default_schema)),
-                        unmatched.clone(),
                         tbl.clone(),
                         Arc::new(DataSchema::from(table.schema_with_stream())),
+                        data_schemas.clone(),
+                        expression_transforms.clone(),
+                        trigger_non_null_errors.clone(),
                     )
                 },
                 transform_len,

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

@@ -23,6 +23,8 @@ mod transform_add_const_columns;
 mod transform_add_internal_columns;
 mod transform_add_stream_columns;
 mod transform_async_function;
+mod transform_branched_async_function;
+mod transform_branched_cast_schema;
 mod transform_cache_scan;
 mod transform_cast_schema;
 mod transform_create_sets;
@@ -50,9 +52,14 @@ pub use transform_add_const_columns::TransformAddConstColumns;
 pub use transform_add_internal_columns::TransformAddInternalColumns;
 pub use transform_add_stream_columns::TransformAddStreamColumns;
 pub use transform_async_function::TransformAsyncFunction;
+pub use transform_branched_async_function::AsyncFunctionBranch;
+pub use transform_branched_async_function::TransformBranchedAsyncFunction;
+pub use transform_branched_cast_schema::CastSchemaBranch;
+pub use transform_branched_cast_schema::TransformBranchedCastSchema;
 pub use transform_cache_scan::CacheSourceState;
 pub use transform_cache_scan::HashJoinCacheState;
 pub use transform_cache_scan::TransformCacheScan;
+pub use transform_cast_schema::build_cast_exprs;
 pub use transform_cast_schema::TransformCastSchema;
 pub use transform_create_sets::TransformCreateSets;
 pub use transform_expression_scan::TransformExpressionScan;
@@ -64,6 +71,7 @@ pub use transform_null_if::TransformNullIf;
 pub use transform_recursive_cte_scan::TransformRecursiveCteScan;
 pub use transform_recursive_cte_source::TransformRecursiveCteSource;
 pub use transform_resort_addon::TransformResortAddOn;
+pub use transform_resort_addon_without_source_schema::build_expression_transform;
 pub use transform_resort_addon_without_source_schema::TransformResortAddOnWithoutSourceSchema;
 pub use transform_srf::TransformSRF;
 pub use transform_udf_script::TransformUdfScript;

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

@@ -59,27 +59,7 @@ impl TransformAsyncFunction {
         sequence_name: &String,
         data_type: &DataType,
     ) -> Result<()> {
-        let count = data_block.num_rows() as u64;
-        let value = if count == 0 {
-            UInt64Type::from_data(vec![])
-        } else {
-            let tenant = self.ctx.get_tenant();
-            let catalog = self.ctx.get_default_catalog()?;
-            let req = GetSequenceNextValueReq {
-                ident: SequenceIdent::new(&tenant, sequence_name),
-                count,
-            };
-            let resp = catalog.get_sequence_next_value(req).await?;
-            let range = resp.start..resp.start + count;
-            UInt64Type::from_data(range.collect::<Vec<u64>>())
-        };
-        let entry = BlockEntry {
-            data_type: data_type.clone(),
-            value: Value::Column(value),
-        };
-        data_block.add_column(entry);
-
-        Ok(())
+        transform_sequence(&self.ctx, data_block, sequence_name, data_type).await
     }
 }
 
@@ -114,3 +94,32 @@ impl AsyncTransform for TransformAsyncFunction {
         Ok(data_block)
     }
 }
+
+pub async fn transform_sequence(
+    ctx: &Arc<QueryContext>,
+    data_block: &mut DataBlock,
+    sequence_name: &String,
+    data_type: &DataType,
+) -> Result<()> {
+    let count = data_block.num_rows() as u64;
+    let value = if count == 0 {
+        UInt64Type::from_data(vec![])
+    } else {
+        let tenant = ctx.get_tenant();
+        let catalog = ctx.get_default_catalog()?;
+        let req = GetSequenceNextValueReq {
+            ident: SequenceIdent::new(&tenant, sequence_name),
+            count,
+        };
+        let resp = catalog.get_sequence_next_value(req).await?;
+        let range = resp.start..resp.start + count;
+        UInt64Type::from_data(range.collect::<Vec<u64>>())
+    };
+    let entry = BlockEntry {
+        data_type: data_type.clone(),
+        value: Value::Column(value),
+    };
+    data_block.add_column(entry);
+
+    Ok(())
+}

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

@@ -0,0 +1,74 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use databend_common_exception::Result;
+use databend_common_expression::BlockMetaInfoDowncast;
+use databend_common_expression::DataBlock;
+use databend_common_expression::SourceSchemaIndex;
+use databend_common_pipeline_transforms::processors::AsyncTransform;
+
+use crate::pipelines::processors::transforms::transform_async_function::transform_sequence;
+use crate::sessions::QueryContext;
+use crate::sql::executor::physical_plans::AsyncFunctionDesc;
+use crate::sql::plans::AsyncFunctionArgument;
+
+/// The key of branches is `SourceSchemaIndex`, see `TransformResortAddOnWithoutSourceSchema`.
+pub struct TransformBranchedAsyncFunction {
+    pub ctx: Arc<QueryContext>,
+    pub branches: Arc<HashMap<SourceSchemaIndex, AsyncFunctionBranch>>,
+}
+
+pub struct AsyncFunctionBranch {
+    pub async_func_descs: Vec<AsyncFunctionDesc>,
+}
+
+#[async_trait::async_trait]
+impl AsyncTransform for TransformBranchedAsyncFunction {
+    const NAME: &'static str = "BranchedAsyncFunction";
+
+    #[async_backtrace::framed]
+    async fn transform(&mut self, mut block: DataBlock) -> Result<DataBlock> {
+        // see the comment details of `TransformResortAddOnWithoutSourceSchema`.
+        if block.get_meta().is_none() {
+            return Ok(block);
+        }
+        let input_schema_idx =
+            SourceSchemaIndex::downcast_from(block.clone().get_owned_meta().unwrap()).unwrap();
+        let Some(branch) = self.branches.get(&input_schema_idx) else {
+            // no async function to execute in this branch, just return the original block
+            return Ok(block);
+        };
+
+        let AsyncFunctionBranch { async_func_descs } = branch;
+
+        for async_func_desc in async_func_descs.iter() {
+            match &async_func_desc.func_arg {
+                AsyncFunctionArgument::SequenceFunction(sequence_name) => {
+                    transform_sequence(
+                        &self.ctx,
+                        &mut block,
+                        sequence_name,
+                        &async_func_desc.data_type,
+                    )
+                    .await?;
+                }
+                AsyncFunctionArgument::DictGetFunction(_) => unreachable!(),
+            }
+        }
+        Ok(block)
+    }
+}