YQL-19163 Fix StaticFold behavior with optional args lambdas

commit_hash:bbea0c78d15d9ba627f15182d80dd2406a494b93

4 files changed, 36 insertions, 10 deletions

diff --git a/yql/essentials/core/type_ann/type_ann_core.cpp b/yql/essentials/core/type_ann/type_ann_core.cpp
index 4b8f9c63b9..ec3061a1b1 100644
--- a/yql/essentials/core/type_ann/type_ann_core.cpp
+++ b/yql/essentials/core/type_ann/type_ann_core.cpp
@@ -6245,11 +6245,11 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
 
         TExprNode::TPtr mergeLambda = nullptr;
         if (input->ChildrenSize() == 3) {
-            mergeLambda = input->ChildPtr(2);
-            auto status = ConvertToLambda(mergeLambda, ctx.Expr, 3);
+            auto status = ConvertToLambda(input->ChildRef(2), ctx.Expr, 3);
             if (status.Level != IGraphTransformer::TStatus::Ok) {
                 return status;
             }
+            mergeLambda = input->ChildPtr(2);
         } else {
             mergeLambda = ctx.Expr.Builder(input->Pos())
                 .Lambda()
@@ -6540,17 +6540,16 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         TExprNode::TPtr result = nullptr;
         TExprNode::TPtr initFunc = nullptr;
         if (input->Content() == "StaticFold1") {
-            initFunc = input->ChildPtr(1);
-
-            auto status = ConvertToLambda(initFunc, ctx.Expr, 1);
+            auto status = ConvertToLambda(input->ChildRef(1), ctx.Expr, 1);
             if (status.Level != IGraphTransformer::TStatus::Ok) {
                 return status;
             }
+            initFunc = input->ChildPtr(1);
         } else {
             result = input->ChildPtr(1);
         }
 
-        auto reduceFunc = input->ChildPtr(2);
+        auto& reduceFunc = input->ChildRef(2);
         auto status = ConvertToLambda(reduceFunc, ctx.Expr, 2);
         if (status.Level != IGraphTransformer::TStatus::Ok) {
             return status;
diff --git a/yql/essentials/tests/sql/sql2yql/canondata/result.json b/yql/essentials/tests/sql/sql2yql/canondata/result.json
index ab345774f4..8e339438a1 100644
--- a/yql/essentials/tests/sql/sql2yql/canondata/result.json
+++ b/yql/essentials/tests/sql/sql2yql/canondata/result.json
@@ -6693,9 +6693,9 @@
     ],
     "test_sql2yql.test[expr-static_fold]": [
         {
-            "checksum": "90e676aca21d7b7a84997aa261cba619",
-            "size": 3647,
-            "uri": "https://{canondata_backend}/1936273/4a1b39013e1bae40e722cff8ccef8829784964e2/resource.tar.gz#test_sql2yql.test_expr-static_fold_/sql.yql"
+            "checksum": "397ed54ebb027b5d4ea51a7c577bfad0",
+            "size": 4863,
+            "uri": "https://{canondata_backend}/1900335/87071896eae3051f9f7c20c1c894d1187145cde1/resource.tar.gz#test_sql2yql.test_expr-static_fold_/sql.yql"
         }
     ],
     "test_sql2yql.test[expr-static_zip]": [
diff --git a/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_expr-static_fold_/formatted.sql b/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_expr-static_fold_/formatted.sql
index edc564c1f5..428d75f9ee 100644
--- a/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_expr-static_fold_/formatted.sql
+++ b/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_expr-static_fold_/formatted.sql
@@ -25,3 +25,18 @@ SELECT
     CAST(StaticFold1($st_empty, $start, $concat) AS Optional<String>),
     StaticFold1($tup, $start, $concat)
 ;
+
+--WithOptionalArgs lambda test
+$puk = ($row_struct) -> {
+    RETURN StaticFold(
+        $row_struct,
+        0,
+        ($item, $sum?) -> {
+            RETURN $sum + IF($item IS NULL, 0, 1)
+        }
+    )
+};
+
+SELECT
+    $puk(Unwrap(CAST(<|one: "8912", two: 42|> AS Struct<one: Int64, two: Int64>)))
+;
diff --git a/yql/essentials/tests/sql/suites/expr/static_fold.sql b/yql/essentials/tests/sql/suites/expr/static_fold.sql
index 9118da0e06..81444ea9b3 100644
--- a/yql/essentials/tests/sql/suites/expr/static_fold.sql
+++ b/yql/essentials/tests/sql/suites/expr/static_fold.sql
@@ -18,4 +18,16 @@ SELECT
     StaticFold1($st_many, $start, $concat),
     StaticFold1($st_single, $start, $concat),
     CAST(StaticFold1($st_empty, $start, $concat) AS Optional<String>),
-    StaticFold1($tup, $start, $concat);
\ No newline at end of file
+    StaticFold1($tup, $start, $concat);
+
+--WithOptionalArgs lambda test
+$puk = ($row_struct) -> {
+    RETURN StaticFold(
+        $row_struct, 
+        0, 
+        ($item, $sum?) -> {
+            RETURN $sum + IF($item IS NULL, 0, 1)
+        }
+    )
+};
+SELECT $puk(Unwrap(CAST(<|one:"8912", two:42|> AS Struct<one: Int64, two:Int64>)));
\ No newline at end of file