Fix WindowTraits rebuild with different input type

1 files changed, 17 insertions, 8 deletions

diff --git a/ydb/library/yql/core/yql_opt_window.cpp b/ydb/library/yql/core/yql_opt_window.cpp
index 6b2c8819934..6f0d88e3bba 100644
--- a/ydb/library/yql/core/yql_opt_window.cpp
+++ b/ydb/library/yql/core/yql_opt_window.cpp
@@ -268,12 +268,16 @@ TCalcOverWindowTraits ExtractCalcOverWindowTraits(const TExprNode::TPtr& frames,
                 rawTraits.CalculateLambdaLead = lead;
                 rawTraits.OutputType = traits->Child(1)->GetTypeAnn();
                 YQL_ENSURE(rawTraits.OutputType);
+            } else if (traits->IsCallable({"Rank", "DenseRank"})) {
+                rawTraits.OutputType = traits->Child(1)->GetTypeAnn();
+                auto lambdaInputType =
+                    traits->Child(0)->GetTypeAnn()->Cast<TTypeExprType>()->GetType()->Cast<TListExprType>()->GetItemType();
+                auto lambda = ReplaceFirstLambdaArgWithCastStruct(*traits->Child(1), *lambdaInputType, ctx);
+                rawTraits.CalculateLambda = ctx.ChangeChild(*traits, 1, std::move(lambda));
             } else {
-                YQL_ENSURE(traits->IsCallable({"RowNumber", "Rank", "DenseRank"}));
-
+                YQL_ENSURE(traits->IsCallable("RowNumber"));
                 rawTraits.CalculateLambda = traits;
-                rawTraits.OutputType = traits->IsCallable("RowNumber") ?
-                    ctx.MakeType<TDataExprType>(EDataSlot::Uint64) : traits->Child(1)->GetTypeAnn();
+                rawTraits.OutputType = traits->GetTypeAnn();
             }
         }
     }
@@ -2804,12 +2808,14 @@ TExprNode::TPtr RebuildCalcOverWindowGroup(TPositionHandle pos, const TExprNode:
                 auto traits = kvTuple->ChildPtr(1);
                 YQL_ENSURE(traits->IsCallable({"Lag", "Lead", "RowNumber", "Rank", "DenseRank", "WindowTraits"}));
                 if (traits->IsCallable("WindowTraits")) {
+                    YQL_ENSURE(traits->Head().GetTypeAnn());
+                    const TTypeAnnotationNode& oldItemType = *traits->Head().GetTypeAnn()->Cast<TTypeExprType>()->GetType();
                     traits = ctx.Builder(traits->Pos())
                         .Callable(traits->Content())
                             .Add(0, inputItemType)
-                            .Add(1, ctx.DeepCopyLambda(*traits->Child(1)))
-                            .Add(2, ctx.DeepCopyLambda(*traits->Child(2)))
-                            .Add(3, ctx.DeepCopyLambda(*traits->Child(3)))
+                            .Add(1, ctx.DeepCopyLambda(*ReplaceFirstLambdaArgWithCastStruct(*traits->Child(1), oldItemType, ctx)))
+                            .Add(2, ctx.DeepCopyLambda(*ReplaceFirstLambdaArgWithCastStruct(*traits->Child(2), oldItemType, ctx)))
+                            .Add(3, ctx.DeepCopyLambda(*ReplaceFirstLambdaArgWithCastStruct(*traits->Child(3), oldItemType, ctx)))
                             .Add(4, ctx.DeepCopyLambda(*traits->Child(4)))
                             .Add(5, traits->Child(5)->IsLambda() ? ctx.DeepCopyLambda(*traits->Child(5)) : traits->ChildPtr(5))
                         .Seal()
@@ -2818,7 +2824,10 @@ TExprNode::TPtr RebuildCalcOverWindowGroup(TPositionHandle pos, const TExprNode:
                     TExprNodeList args;
                     args.push_back(inputType);
                     if (traits->ChildrenSize() > 1) {
-                        args.push_back(ctx.DeepCopyLambda(*traits->Child(1)));
+                        YQL_ENSURE(traits->Head().GetTypeAnn());
+                        const TTypeAnnotationNode& oldItemType = *traits->Head().GetTypeAnn()->Cast<TTypeExprType>()->GetType()
+                            ->Cast<TListExprType>()->GetItemType();
+                        args.push_back(ctx.DeepCopyLambda(*ReplaceFirstLambdaArgWithCastStruct(*traits->Child(1), oldItemType, ctx)));
                     }
                     if (traits->ChildrenSize() > 2) {
                         args.push_back(traits->ChildPtr(2));