Improve constant folding of UDFs (#45316)

Co-authored-by: Velikhov Pavel <[email protected]>

1 files changed, 6 insertions, 3 deletions

diff --git a/ydb/core/kqp/opt/cbo/solver/kqp_opt_stat.cpp b/ydb/core/kqp/opt/cbo/solver/kqp_opt_stat.cpp
index 1b03b0ecc7b..2c09a8211fb 100644
--- a/ydb/core/kqp/opt/cbo/solver/kqp_opt_stat.cpp
+++ b/ydb/core/kqp/opt/cbo/solver/kqp_opt_stat.cpp
@@ -90,7 +90,8 @@ namespace {
         if (!TCoApply::Match(input.Get())) {
             return false;
         }
-        if (input->ChildrenSize() != 2) {
+
+        if (input->ChildrenSize() != 2 && input->ChildrenSize() != 3) {
             return false;
         }
         if (input->Child(0)->IsCallable("Udf")) {
@@ -102,10 +103,12 @@ namespace {
                 }
             }
             if (withParams) {
-                return IsConstantExprWithParams(input->Child(1));
+                return IsConstantExprWithParams(input->Child(1)) || IsConstantUdf(input->Child(1));
             } else {
-                return IsConstantExpr(input->Child(1));
+                return IsConstantExpr(input->Child(1)) || IsConstantUdf(input->Child(1));
             }
+        } else if (TCoApply::Match(input->Child(1))){
+            return IsConstantUdf(input->ChildPtr(1));
         }
         return false;
     }