fixed LambdaArgumentsCount for optional args

3 files changed, 50 insertions, 3 deletions

diff --git a/ydb/docs/ru/core/yql/reference/yql-core/builtins/_includes/types.md b/ydb/docs/ru/core/yql/reference/yql-core/builtins/_includes/types.md
index 360ca4f5a2..6af97c139c 100644
--- a/ydb/docs/ru/core/yql/reference/yql-core/builtins/_includes/types.md
+++ b/ydb/docs/ru/core/yql/reference/yql-core/builtins/_includes/types.md
@@ -795,4 +795,18 @@ SELECT LambdaArgumentsCount(($x, $y)->($x+$y))
 ; -- 2
 ```
 
+## LambdaOptionalArgumentsCount
+
+**Сигнатура**
+```
+LambdaOptionalArgumentsCount(LambdaFunction)->Uint32 
+```
+Получение количества опциональных аргументов в лямбда-функции.
+
+**Примеры:**
+``` yql
+SELECT LambdaOptionalArgumentsCount(($x, $y, $z?)->(if($x,$y,$z)))
+; -- 1
+```
+
 {% endif %}
diff --git a/ydb/library/yql/core/type_ann/type_ann_core.cpp b/ydb/library/yql/core/type_ann/type_ann_core.cpp
index 057c290b60..678a587cec 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -9751,14 +9751,45 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
             return IGraphTransformer::TStatus::Error;
         }
 
-        if (!input->Head().IsLambda()) {
-            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()), TStringBuilder() << "Expected lambda, but got: " << input->Head().Type()));
+        auto lambda = input->HeadPtr();
+        if (input->Head().IsCallable("WithOptionalArgs")) {
+            lambda = input->Head().HeadPtr();
+        }
+
+        if (!lambda->IsLambda()) {
+            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()), TStringBuilder() << "Expected lambda, but got: " << lambda->Type()));
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        output = ctx.Expr.Builder(input->Pos())
+            .Callable("Uint32")
+                .Atom(0, ToString(lambda->Head().ChildrenSize()))
+            .Seal()
+            .Build();
+
+        return IGraphTransformer::TStatus::Repeat;
+    }
+
+    IGraphTransformer::TStatus LambdaOptionalArgumentsCountWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        if (!EnsureArgsCount(*input, 1, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        ui32 optionalArgsCount = 0;
+        auto lambda = input->HeadPtr();
+        if (input->Head().IsCallable("WithOptionalArgs")) {
+            lambda = input->Head().HeadPtr();
+            optionalArgsCount = FromString<ui32>(input->Head().Tail().Content());
+        }
+
+        if (!lambda->IsLambda()) {
+            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()), TStringBuilder() << "Expected lambda, but got: " << lambda->Type()));
             return IGraphTransformer::TStatus::Error;
         }
 
         output = ctx.Expr.Builder(input->Pos())
             .Callable("Uint32")
-                .Atom(0, ToString(input->Head().Head().ChildrenSize()))
+                .Atom(0, ToString(optionalArgsCount))
             .Seal()
             .Build();
 
@@ -11515,6 +11546,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         Functions["CallableArgument"] = &CallableArgumentWrapper;
         Functions["CallableTypeHandle"] = &MakeTypeHandleWrapper<ETypeAnnotationKind::Callable>;
         Functions["LambdaArgumentsCount"] = LambdaArgumentsCountWrapper;
+        Functions["LambdaOptionalArgumentsCount"] = LambdaOptionalArgumentsCountWrapper;
         Functions["FormatCode"] = &FormatCodeWrapper;
         Functions["FormatCodeWithPositions"] = &FormatCodeWrapper;
         Functions["SerializeCode"] = &FormatCodeWrapper;
diff --git a/ydb/library/yql/sql/v1/builtin.cpp b/ydb/library/yql/sql/v1/builtin.cpp
index 9fa9334047..5ee117730c 100644
--- a/ydb/library/yql/sql/v1/builtin.cpp
+++ b/ydb/library/yql/sql/v1/builtin.cpp
@@ -2864,6 +2864,7 @@ struct TBuiltinFuncData {
             {"reprcode", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ReprCode", 1, 1) },
             {"quotecode", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("QuoteCode", 1, 1) },
             {"lambdaargumentscount", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("LambdaArgumentsCount", 1, 1) },
+            {"lambdaoptionalargumentscount", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("LambdaOptionalArgumentsCount", 1, 1) },
             {"subqueryextend", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("SubqueryExtend", 1, -1) },
             {"subqueryunionall", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("SubqueryUnionAll", 1, -1) },
             {"subquerymerge", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("SubqueryMerge", 1, -1) },