fix(query): fix subquery writer in udf

src/query/sql/src/evaluator/cse.rs

@@ -15,6 +15,7 @@
 use std::collections::HashMap;
 
 use databend_common_expression::Expr;
+use databend_common_functions::BUILTIN_FUNCTIONS;
 use log::debug;
 
 use super::BlockOperator;
@@ -117,6 +118,9 @@ pub fn apply_cse(
 /// `count_expressions` recursively counts the occurrences of expressions in an expression tree
 /// and stores the count in a HashMap.
 fn count_expressions(expr: &Expr, counter: &mut HashMap<Expr, usize>) {
+    if !expr.is_deterministic(&BUILTIN_FUNCTIONS) {
+        return;
+    }
     match expr {
         Expr::FunctionCall { function, .. } if function.signature.name == "if" => {}
         Expr::FunctionCall { function, .. } if function.signature.name == "is_not_error" => {}

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

@@ -168,17 +168,20 @@ impl SubqueryRewriter {
                 ))
             }
 
-            RelOperator::Limit(_) | RelOperator::Sort(_) => Ok(SExpr::create_unary(
+            RelOperator::Limit(_)
+            | RelOperator::Sort(_)
+            | RelOperator::Udf(_)
+            | RelOperator::AsyncFunction(_) => Ok(SExpr::create_unary(
                 Arc::new(s_expr.plan().clone()),
                 Arc::new(self.rewrite(s_expr.child(0)?)?),
             )),
 
             RelOperator::DummyTableScan(_)
             | RelOperator::Scan(_)
             | RelOperator::CteScan(_)
-            | RelOperator::ConstantTableScan(_) => Ok(s_expr.clone()),
-
-            _ => Err(ErrorCode::Internal("Invalid plan type")),
+            | RelOperator::ConstantTableScan(_)
+            | RelOperator::AddRowNumber(_)
+            | RelOperator::Exchange(_) => Ok(s_expr.clone()),
         }
     }
 

tests/sqllogictests/suites/udf_server/udf_server_test.test

@@ -170,6 +170,12 @@ NULL
 6
 NULL
 
+query I
+SELECT gcd(a,b) d from (select number + 1 a, a * 2  b from numbers(3) where number in (select * from numbers(300))) order by d;
+----
+1
+2
+3
 
 
 statement ok