expressions in distinct on

2 files changed, 95 insertions, 38 deletions

diff --git a/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp b/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
index bc071ee04ef..13ab88bf879 100644
--- a/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
@@ -2082,12 +2082,12 @@ TExprNode::TPtr BuildSort(TPositionHandle pos, const TExprNode::TPtr& sort, cons
         .Build();
 }
 
-TExprNode::TPtr BuildDistinctOn(TPositionHandle pos, const TExprNode::TPtr& list, const TExprNode::TPtr& distinctOn,
+TExprNode::TPtr BuildDistinctOn(TPositionHandle pos, TExprNode::TPtr list, const TExprNode::TPtr& distinctOn,
     const TExprNode::TPtr& sort, TExprContext& ctx) {
     // filter by RowNumber() == 1
 
     TExprNode::TListType args;
-    auto begin = ctx.NewCallable(pos, "Void", {});

+    auto begin = ctx.NewCallable(pos, "Void", {});
     auto end = ctx.NewCallable(pos, "Int32", { ctx.NewAtom(pos, "0") });
     args.push_back(ctx.Builder(pos)
         .List()
@@ -2124,31 +2124,86 @@ TExprNode::TPtr BuildDistinctOn(TPositionHandle pos, const TExprNode::TPtr& list
         .Seal()
         .Build();
 
-    TExprNode::TListType keys;
+    bool rebuildKey = false;
     for (auto p : distinctOn->Children()) {
         YQL_ENSURE(p->IsCallable("PgGroup"));
-        const auto& member = p->Tail().Tail();
-        YQL_ENSURE(member.IsCallable("Member"));
-        keys.push_back(member.TailPtr());
+        NNodes::TCoLambda lambda(&p->Tail());
+        if (!lambda.Body().Ref().IsCallable("Member") ||
+            &lambda.Body().Ref().Head() != &lambda.Args().Arg(0).Ref()) {
+            rebuildKey = true;
+            break;
+        }
+    }
+
+    TExprNode::TPtr keysNode;
+    auto originalList = list;
+    if (!rebuildKey) {
+        TExprNode::TListType keys;
+        for (auto p : distinctOn->Children()) {
+            NNodes::TCoLambda lambda(&p->Tail());
+            const auto& member = lambda.Body().Ref();
+            YQL_ENSURE(member.IsCallable("Member"));
+            keys.push_back(member.TailPtr());
+        }
+
+        keysNode = ctx.NewList(pos, std::move(keys));
+    } else {
+        TExprNode::TListType keys;
+        for (ui32 i = 0; i < distinctOn->ChildrenSize(); ++i) {
+            keys.push_back(ctx.NewAtom(pos, "_yql_distinct_on_" + ToString(i)));
+        }
+
+        keysNode = ctx.NewList(pos, std::move(keys));
+
+        list = ctx.Builder(pos)
+            .Callable("OrderedMap")
+                .Add(0, list)
+                .Lambda(1)
+                    .Param("row")
+                    .Callable("FlattenMembers")
+                        .List(0)
+                            .Atom(0, "")
+                            .Arg(1, "row")
+                        .Seal()
+                        .List(1)
+                            .Atom(0, "")
+                            .Callable(1, "AsStruct")
+                                .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder & {
+                                    for (ui32 i = 0; i < distinctOn->ChildrenSize(); ++i) {
+                                        parent.List(i)
+                                            .Atom(0, "_yql_distinct_on_" + ToString(i))
+                                            .Apply(1, distinctOn->Child(i)->Tail())
+                                                .With(0, "row")
+                                             .Seal()
+                                        .Seal();
+                                    }
+
+                                    return parent;
+                                })
+                            .Seal()
+                        .Seal()
+                    .Seal()
+                .Seal()
+            .Seal()
+            .Build();
     }
 
-    auto keysNode = ctx.NewList(pos, std::move(keys));
     auto sortNode = ctx.NewCallable(pos, "Void", {});
     if (sort && sort->Tail().ChildrenSize() > 0) {
         sortNode = BuildSortTraits(pos, sort->Tail(), list, ctx);
     }
 
-    auto ret = ctx.Builder(pos)

-        .Callable("CalcOverWindow")

-            .Add(0, list)

-            .Add(1, keysNode)

-            .Add(2, sortNode)

-            .Add(3, frames)

-        .Seal()

+    auto ret = ctx.Builder(pos)
+        .Callable("CalcOverWindow")
+            .Add(0, list)
+            .Add(1, keysNode)
+            .Add(2, sortNode)
+            .Add(3, frames)
+        .Seal()
         .Build();
 
     ret = ctx.Builder(pos)
-        .Callable("Filter")
+        .Callable("OrderedFilter")
             .Add(0, ret)
             .Lambda(1)
                 .Param("row")
@@ -2166,19 +2221,23 @@ TExprNode::TPtr BuildDistinctOn(TPositionHandle pos, const TExprNode::TPtr& list
         .Build();
 
     ret = ctx.Builder(pos)
-        .Callable("Map")
+        .Callable("OrderedMap")
             .Add(0, ret)
             .Lambda(1)
                 .Param("row")
-                .Callable("RemoveMember")
+                .Callable("CastStruct")
                     .Arg(0, "row")
-                    .Atom(1, "_yql_row_number")
+                    .Callable(1, "ListItemType")
+                        .Callable(0, "TypeOf")
+                            .Add(0, originalList)
+                        .Seal()
+                    .Seal()
                 .Seal()
             .Seal()
         .Seal()
         .Build();
 
-    return ret;

+    return ret;
 }
 
 TExprNode::TPtr BuildOffset(TPositionHandle pos, const TExprNode::TPtr& offset, const TExprNode::TPtr& list, TExprContext& ctx) {
diff --git a/ydb/library/yql/sql/pg/pg_sql.cpp b/ydb/library/yql/sql/pg/pg_sql.cpp
index d66f14229f7..3a3a3704f6c 100644
--- a/ydb/library/yql/sql/pg/pg_sql.cpp
+++ b/ydb/library/yql/sql/pg/pg_sql.cpp
@@ -313,17 +313,15 @@ public:
                 } else {
                     for (int i = 0; i < ListLength(x->distinctClause); ++i) {
                         auto node = ListNodeNth(x->distinctClause, i);
-                        if (NodeTag(node) != T_ColumnRef) {
-                            NodeNotImplemented(x->distinctClause, node);
+                        TExprSettings settings;
+                        settings.AllowColumns = true;
+                        settings.Scope = "DISTINCT ON";
+                        auto expr = ParseExpr(node, settings);
+                        if (!expr) {
                             return nullptr;
                         }
 
-                        auto ref = ParseColumnRef(CAST_NODE(ColumnRef, node));
-                        if (!ref) {
-                            return nullptr;
-                        }
-
-                        auto lambda = L(A("lambda"), QL(), ref);
+                        auto lambda = L(A("lambda"), QL(), expr);
                         distinctOnItems.push_back(L(A("PgGroup"), L(A("Void")), lambda));
                     }
                 }
@@ -788,14 +786,14 @@ public:
         TView view;
         view.Name = value->ctename;
 
-        for (int i = 0; i < ListLength(value->aliascolnames); ++i) {

-            auto node = ListNodeNth(value->aliascolnames, i);

-            if (NodeTag(node) != T_String) {

-                NodeNotImplemented(value, node);

-                return false;

-            }

-

-            view.ColNames.push_back(StrVal(node));

+        for (int i = 0; i < ListLength(value->aliascolnames); ++i) {
+            auto node = ListNodeNth(value->aliascolnames, i);
+            if (NodeTag(node) != T_String) {
+                NodeNotImplemented(value, node);
+                return false;
+            }
+
+            view.ColNames.push_back(StrVal(node));
         }
 
         if (NodeTag(value->ctequery) != T_SelectStmt) {
@@ -803,9 +801,9 @@ public:
             return false;
         }
 
-        view.Source = ParseSelectStmt(CAST_NODE(SelectStmt, value->ctequery), true);

-        if (!view.Source) {

-            return false;

+        view.Source = ParseSelectStmt(CAST_NODE(SelectStmt, value->ctequery), true);
+        if (!view.Source) {
+            return false;
         }
 
         auto& currentCTEs = CTE.back();