[YQL-14347] Implement StaticZip

ref:ba68516ee7d540558b824351fc8854b4ee4ce880

2 files changed, 100 insertions, 0 deletions

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 f627039201c..5846e6cb108 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -5754,6 +5754,104 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         return IGraphTransformer::TStatus::Repeat;
     }
 
+    IGraphTransformer::TStatus StaticZipWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        if (!EnsureMinArgsCount(*input, 1, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        const TStructExprType* argStruct = nullptr;
+        TVector<TStringBuf> argNames;
+        const TTupleExprType* argTuple = nullptr;
+
+        auto getMemberNames = [](const TStructExprType& type) {
+            TVector<TStringBuf> result;
+            for (auto& item : type.GetItems()) {
+                result.push_back(item->GetName());
+            }
+            Sort(result);
+            return result;
+        };
+
+        for (ui32 i = 0; i < input->ChildrenSize(); ++i) {
+            auto child = input->Child(i);
+            const TTypeAnnotationNode* childType = child->GetTypeAnn();
+            if (HasError(childType, ctx.Expr)) {
+                return IGraphTransformer::TStatus::Error;
+            }
+
+            if (!childType) {
+                YQL_ENSURE(child->Type() == TExprNode::Lambda);
+                ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(child->Pos()), TStringBuilder() << "Expected either struct or tuple, but got lambda"));
+                return IGraphTransformer::TStatus::Error;
+            }
+
+            if (argStruct) {
+                if (childType->GetKind() != ETypeAnnotationKind::Struct) {
+                    ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(child->Pos()),
+                        TStringBuilder() << "Expected all arguments to be of Struct type, but got: " << *childType << " for " << i << "th argument"));
+                    return IGraphTransformer::TStatus::Error;
+                }
+                auto childNames = getMemberNames(*childType->Cast<TStructExprType>());
+                if (childNames != argNames) {
+                    ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(child->Pos()),
+                        TStringBuilder() << "Struct members mismatch. Members for first argument: " << JoinSeq(",", argNames)
+                                         << ". Members for " << i << "th argument: " << JoinSeq(", ", childNames)));
+                    return IGraphTransformer::TStatus::Error;
+                }
+            } else if (argTuple) {
+                if (childType->GetKind() != ETypeAnnotationKind::Tuple) {
+                    ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(child->Pos()),
+                        TStringBuilder() << "Expected all arguments to be of Tuple type, but got: " << *childType << " for " << i << "th argument"));
+                    return IGraphTransformer::TStatus::Error;
+                }
+                auto childSize = childType->Cast<TTupleExprType>()->GetSize();
+                if (childSize != argTuple->GetSize()) {
+                    ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(child->Pos()),
+                        TStringBuilder() << "Tuple size mismatch. Got " << argTuple->GetSize()
+                                         << " for first argument and " << childSize << " for " << i << "th argument"));
+                    return IGraphTransformer::TStatus::Error;
+                }
+            } else if (childType->GetKind() == ETypeAnnotationKind::Struct) {
+                argStruct = childType->Cast<TStructExprType>();
+                argNames = getMemberNames(*argStruct);
+            } else if (childType->GetKind() == ETypeAnnotationKind::Tuple) {
+                argTuple = childType->Cast<TTupleExprType>();
+            } else {
+                ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(child->Pos()), TStringBuilder() << "Expected either struct or tuple, but got: " << *childType));
+                return IGraphTransformer::TStatus::Error;
+            }
+        }
+
+        if (argStruct) {
+            TExprNodeList asStructArgs;
+            for (auto& name : argNames) {
+                auto nameAtom = ctx.Expr.NewAtom(input->Pos(), name);
+                TExprNodeList zippedValues;
+                for (auto child : input->ChildrenList()) {
+                    zippedValues.push_back(ctx.Expr.NewCallable(child->Pos(), "Member", { child , nameAtom }));
+                }
+                auto zipped = ctx.Expr.NewList(input->Pos(), std::move(zippedValues));
+                asStructArgs.push_back(ctx.Expr.NewList(input->Pos(), { nameAtom, zipped }));
+            }
+            output = ctx.Expr.NewCallable(input->Pos(), "AsStruct", std::move(asStructArgs));
+        } else {
+            YQL_ENSURE(argTuple);
+            TExprNodeList tupleArgs;
+            for (size_t i = 0; i < argTuple->GetSize(); ++i) {
+                auto idxAtom = ctx.Expr.NewAtom(input->Pos(), ToString(i), TNodeFlags::Default);
+                TExprNodeList zippedValues;
+                for (auto child : input->ChildrenList()) {
+                    zippedValues.push_back(ctx.Expr.NewCallable(child->Pos(), "Nth", { child , idxAtom }));
+                }
+                auto zipped = ctx.Expr.NewList(input->Pos(), std::move(zippedValues));
+                tupleArgs.push_back(zipped);
+            }
+            output = ctx.Expr.NewList(input->Pos(), std::move(tupleArgs));
+        }
+
+        return IGraphTransformer::TStatus::Repeat;
+    }
+
     IGraphTransformer::TStatus TryRemoveAllOptionalsWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
         if (!EnsureArgsCount(*input, 1, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
@@ -12859,6 +12957,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         Functions[ForName] = &ForWorldWrapper;
         Functions["IfPresent"] = &IfPresentWrapper;
         Functions["StaticMap"] = &StaticMapWrapper;
+        Functions["StaticZip"] = &StaticZipWrapper;
         Functions["TryRemoveAllOptionals"] = &TryRemoveAllOptionalsWrapper;
         Functions["HasNull"] = &HasNullWrapper;
         Functions["TypeOf"] = &TypeOfWrapper;
diff --git a/ydb/library/yql/sql/v1/builtin.cpp b/ydb/library/yql/sql/v1/builtin.cpp
index a86269c290f..43ef1243f10 100644
--- a/ydb/library/yql/sql/v1/builtin.cpp
+++ b/ydb/library/yql/sql/v1/builtin.cpp
@@ -2689,6 +2689,7 @@ struct TBuiltinFuncData {
             {"combinemembers", BuildNamedBuiltinFactoryCallback<TCombineMembers>("FlattenMembers")},
             {"flattenmembers", BuildNamedBuiltinFactoryCallback<TFlattenMembers>("FlattenMembers")},
             {"staticmap", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("StaticMap", 2, 2) },
+            {"staticzip", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("StaticZip", 1, -1) },
 
             // File builtins
             {"filepath", BuildNamedBuiltinFactoryCallback<TYqlAtom>("FilePath")},