fix(query): fix update set with udf and filter (#18487)

Author: b41sh

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

@@ -25,6 +25,8 @@ use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_expression::udf_client::error_kind;
 use databend_common_expression::udf_client::UDFFlightClient;
+use databend_common_expression::BlockEntry;
+use databend_common_expression::ColumnBuilder;
 use databend_common_expression::DataBlock;
 use databend_common_metrics::external_server::record_error_external;
 use databend_common_metrics::external_server::record_retry_external;
@@ -168,6 +170,14 @@ impl AsyncTransform for TransformUdfServer {
 
     #[async_backtrace::framed]
     async fn transform(&mut self, mut data_block: DataBlock) -> Result<DataBlock> {
+        if data_block.is_empty() {
+            for func in self.funcs.iter() {
+                let builder = ColumnBuilder::repeat_default(&func.data_type, 0);
+                let entry = BlockEntry::Column(builder.build());
+                data_block.add_entry(entry);
+            }
+            return Ok(data_block);
+        }
         let rows = data_block.num_rows();
         let batch_rows = self.request_batch_rows;
         let mut batch_blocks: Vec<DataBlock> = (0..rows)

src/query/sql/src/planner/optimizer/optimizer.rs

@@ -47,6 +47,7 @@ use crate::plans::Join;
 use crate::plans::JoinType;
 use crate::plans::MatchedEvaluator;
 use crate::plans::Mutation;
+use crate::plans::MutationSource;
 use crate::plans::Operator;
 use crate::plans::Plan;
 use crate::plans::RelOp;
@@ -398,7 +399,22 @@ async fn optimize_mutation(opt_ctx: Arc<OptimizerContext>, s_expr: SExpr) -> Res
             }
         }
         MutationType::Update | MutationType::Delete => {
-            if let RelOperator::MutationSource(rel) = input_s_expr.plan() {
+            // Helper function to find MutationSource in a chain of operators
+            fn find_mutation_source(s_expr: &SExpr) -> Option<&MutationSource> {
+                match s_expr.plan() {
+                    RelOperator::MutationSource(rel) => Some(rel),
+                    RelOperator::Udf(_) | RelOperator::EvalScalar(_) => {
+                        if s_expr.arity() == 1 {
+                            find_mutation_source(s_expr.unary_child())
+                        } else {
+                            None
+                        }
+                    }
+                    _ => None,
+                }
+            }
+
+            if let Some(rel) = find_mutation_source(&input_s_expr) {
                 if rel.mutation_type == MutationType::Delete && rel.predicates.is_empty() {
                     mutation.truncate_table = true;
                 }

tests/sqllogictests/suites/udf_server/udf_server_test.test

@@ -579,6 +579,19 @@ SELECT * FROM test_update_udf;
 databend.com 12
 databend.cn 11
 
+statement ok
+INSERT INTO test_update_udf (url) VALUES('databend.com'),('databend.cn');
+
+statement ok
+UPDATE test_update_udf SET length = url_len(url) WHERE length IS NULL;
+
+query TI
+SELECT * FROM test_update_udf ORDER BY url;
+----
+databend.cn 11
+databend.cn 11
+databend.com 12
+databend.com 12
 
 statement ok
 CREATE OR REPLACE TABLE test_update_udf_1(url STRING, a INT64,b INT64,c INT64);
@@ -632,19 +645,13 @@ SELECT embedding_4('databend.com')::vector(4);
 ----
 [1.1,1.2,1.3,1.4]
 
-statement ok
-CREATE OR REPLACE TABLE test(url STRING, length INT64);
-
-statement ok
-INSERT INTO test (url) VALUES('databend.com'),('databend.cn');
-
 query T
-SELECT embedding_4('databend.com')::vector(4) fro
+SELECT embedding_4('databend.com')::vector(4)
 ----
 [1.1,1.2,1.3,1.4]
 
 query T
-SELECT embedding_4(url)::vector(4) FROM test;
+SELECT embedding_4(url)::vector(4) FROM test_update_udf_1;
 ----
 [1.1,1.2,1.3,1.4]
 [1.1,1.2,1.3,1.4]
@@ -673,3 +680,4 @@ select sum((value::Int)::Int) from system.metrics where metric = 'external_runni
 ----
 0
 
+