Implement column order for window functions

initial

3 files changed, 25 insertions, 0 deletions

diff --git a/ydb/library/yql/core/type_ann/type_ann_columnorder.cpp b/ydb/library/yql/core/type_ann/type_ann_columnorder.cpp
index 3a6a71ced74..9befa13e82e 100644
--- a/ydb/library/yql/core/type_ann/type_ann_columnorder.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_columnorder.cpp
@@ -266,6 +266,29 @@ IGraphTransformer::TStatus OrderForEquiJoin(const TExprNode::TPtr& node, TExprNo
     return ctx.Types.SetColumnOrder(*node, resultColumnOrder, ctx.Expr);
 };
 
+IGraphTransformer::TStatus OrderForCalcOverWindow(const TExprNode::TPtr& node, TExprNode::TPtr& output, TExtContext& ctx) {
+    Y_UNUSED(output);
+
+    auto inputOrder = ctx.Types.LookupColumnOrder(node->Head());
+    if (!inputOrder) {
+        return IGraphTransformer::TStatus::Ok;
+    }
+
+    const TStructExprType* inputType = node->Head().GetTypeAnn()->Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>();
+    const TStructExprType* outputType = node->GetTypeAnn()->Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>();
+
+    // we simply add new CalcOverWindow columns after original input columns
+    TVector<TString> resultOrder = *inputOrder;
+    for (auto& item : outputType->GetItems()) {
+        auto col = item->GetName();
+        if (!inputType->FindItem(col)) {
+            resultOrder.emplace_back(col);
+        }
+    }
+
+    return ctx.Types.SetColumnOrder(*node, resultOrder, ctx.Expr);
+}
+
 IGraphTransformer::TStatus OrderFromFirst(const TExprNode::TPtr& node, TExprNode::TPtr& output, TExtContext& ctx) {
     Y_UNUSED(output);
     if (auto columnOrder = ctx.Types.LookupColumnOrder(node->Head())) {
diff --git a/ydb/library/yql/core/type_ann/type_ann_columnorder.h b/ydb/library/yql/core/type_ann/type_ann_columnorder.h
index dc9e9c0a564..90c17236a30 100644
--- a/ydb/library/yql/core/type_ann/type_ann_columnorder.h
+++ b/ydb/library/yql/core/type_ann/type_ann_columnorder.h
@@ -14,6 +14,7 @@ IGraphTransformer::TStatus OrderForSqlProject(const TExprNode::TPtr& node, TExpr
 IGraphTransformer::TStatus OrderForMergeExtend(const TExprNode::TPtr& node, TExprNode::TPtr& output, TExtContext& ctx);
 IGraphTransformer::TStatus OrderForUnionAll(const TExprNode::TPtr& node, TExprNode::TPtr& output, TExtContext& ctx);
 IGraphTransformer::TStatus OrderForEquiJoin(const TExprNode::TPtr& node, TExprNode::TPtr& output, TExtContext& ctx);
+IGraphTransformer::TStatus OrderForCalcOverWindow(const TExprNode::TPtr& node, TExprNode::TPtr& output, TExtContext& ctx);
 
 IGraphTransformer::TStatus OrderFromFirst(const TExprNode::TPtr& node, TExprNode::TPtr& output, TExtContext& ctx);
 IGraphTransformer::TStatus OrderFromFirstAndOutputType(const TExprNode::TPtr& node, TExprNode::TPtr& output, TExtContext& ctx);
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 f5335777945..81d50565c68 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -12000,6 +12000,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         ColumnOrderFunctions[RightName] = &OrderFromFirst;
         ColumnOrderFunctions["UnionAll"] = &OrderForUnionAll;
         ColumnOrderFunctions["EquiJoin"] = &OrderForEquiJoin;
+        ColumnOrderFunctions["CalcOverWindow"] = &OrderForCalcOverWindow;
 
         ColumnOrderFunctions["RemovePrefixMembers"] = &OrderFromFirstAndOutputType;
         ColumnOrderFunctions["Sort"] = ColumnOrderFunctions["Take"] = ColumnOrderFunctions["Skip"] =