correlated sublinks in ORDER BY

3 files changed, 237 insertions, 107 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 18df3d57ce8..653606d5ca5 100644
--- a/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
@@ -64,28 +64,8 @@ std::pair<TExprNode::TPtr, TExprNode::TPtr> SplitByPredicate(TPositionHandle pos
     };
 }
 
-TSet<TString> ExtractExternalColumns(const TExprNode& select) {
-    TSet<TString> res;
-    const auto& option = select.Head();
-    auto setItems = GetSetting(option, "set_items");
-    YQL_ENSURE(setItems);
-    for (const auto& s : setItems->Tail().Children()) {
-        YQL_ENSURE(s->IsCallable("PgSetItem"));
-        auto extTypes = GetSetting(s->Head(), "final_ext_types");
-        YQL_ENSURE(extTypes);
-        for (const auto& input : extTypes->Tail().Children()) {
-            auto type = input->Tail().GetTypeAnn()->Cast<TTypeExprType>()->GetType()->Cast<TStructExprType>();
-            for (const auto& item : type->GetItems()) {
-                res.insert(NTypeAnnImpl::MakeAliasedColumn(input->Head().Content(), item->GetName()));
-            }
-        }
-    }
-
-    return res;
-}
-
 TExprNode::TPtr JoinColumns(TPositionHandle pos, const TExprNode::TPtr& list1, const TExprNode::TPtr& list2,
-    TExprNode::TPtr leftJoinColumns, ui32 subLinkId, TExprContext& ctx) {
+    TExprNode::TPtr leftJoinColumns, ui32 subLinkId, TExprContext& ctx, const TString& leftPrefix = {}) {
     auto join = ctx.Builder(pos)
         .Callable("EquiJoin")
             .List(0)
@@ -105,7 +85,7 @@ TExprNode::TPtr JoinColumns(TPositionHandle pos, const TExprNode::TPtr& list1, c
                         if (leftJoinColumns) {
                             for (ui32 i = 0; i < leftJoinColumns->ChildrenSize(); ++i) {
                                 parent.Atom(2 * i, "a");
-                                parent.Add(2* i + 1, leftJoinColumns->ChildPtr(i));
+                                parent.Atom(2* i + 1, leftPrefix + leftJoinColumns->ChildPtr(i)->Content());
                             }
                         }
 
@@ -191,7 +171,8 @@ std::pair<TExprNode::TPtr, TExprNode::TPtr> RewriteSubLinks(TPositionHandle pos,
     const TExprNode::TPtr& list, const TExprNode::TPtr& lambda,
     const TNodeMap<ui32>& subLinks, const TVector<TString>& inputAliases,
     const TExprNode::TListType& cleanedInputs,
-    const TAggregationMap* aggId, TExprContext& ctx, TOptimizeContext& optCtx) {
+    const TAggregationMap* aggId, TExprContext& ctx, TOptimizeContext& optCtx,
+    const TString& leftPrefix = {}, TVector<TString>* sublinkColumns = nullptr) {
     auto newList = list;
     auto originalRow = lambda->Head().HeadPtr();
     auto arg = ctx.NewArgument(pos, "row");
@@ -221,7 +202,7 @@ std::pair<TExprNode::TPtr, TExprNode::TPtr> RewriteSubLinks(TPositionHandle pos,
                 RewriteAggs(testLambda, *aggId, ctx, optCtx, true);
             }
 
-            auto extColumns = ExtractExternalColumns(node->Tail());
+            auto extColumns = NTypeAnnImpl::ExtractExternalColumns(node->Tail());
             if (extColumns.empty()) {
                 auto select = ExpandPgSelectSublink(node->TailPtr(), ctx, optCtx, it->second, cleanedInputs, inputAliases);
                 if (linkType == "exists") {
@@ -354,7 +335,7 @@ std::pair<TExprNode::TPtr, TExprNode::TPtr> RewriteSubLinks(TPositionHandle pos,
                         .Do([&](TExprNodeBuilder &parent) -> TExprNodeBuilder & {
                             ui32 i = 0;
                             for (const auto& c : extColumns) {
-                                parent.Atom(i++, c);
+                                parent.Atom(i++, c.first);
                             }
 
                             return parent;
@@ -467,6 +448,10 @@ std::pair<TExprNode::TPtr, TExprNode::TPtr> RewriteSubLinks(TPositionHandle pos,
                 auto columnName = "_yql_sublink_" + ToString(it->second);
                 TExprNode::TListType aggregateItems;
                 if (linkType == "exists") {
+                    if (sublinkColumns) {
+                        sublinkColumns->push_back(columnName);
+                    }
+
                     aggregateItems.push_back(ctx.Builder(node->Pos())
                         .List()
                             .Atom(0, columnName)
@@ -474,6 +459,11 @@ std::pair<TExprNode::TPtr, TExprNode::TPtr> RewriteSubLinks(TPositionHandle pos,
                         .Seal()
                         .Build());
                 } else if (linkType == "expr") {
+                    if (sublinkColumns) {
+                        sublinkColumns->push_back(columnName + "_count");
+                        sublinkColumns->push_back(columnName + "_value");
+                    }
+
                     aggregateItems.push_back(ctx.Builder(node->Pos())
                         .List()
                             .Atom(0, columnName + "_count")
@@ -487,6 +477,11 @@ std::pair<TExprNode::TPtr, TExprNode::TPtr> RewriteSubLinks(TPositionHandle pos,
                         .Seal()
                         .Build());
                 } else if (linkType == "any") {
+                    if (sublinkColumns) {
+                        sublinkColumns->push_back(columnName + "_count");
+                        sublinkColumns->push_back(columnName + "_value");
+                    }
+
                     aggregateItems.push_back(ctx.Builder(node->Pos())
                         .List()
                             .Atom(0, columnName + "_count")
@@ -501,6 +496,11 @@ std::pair<TExprNode::TPtr, TExprNode::TPtr> RewriteSubLinks(TPositionHandle pos,
                         .Build());
                 } else {
                     YQL_ENSURE(linkType == "all");
+                    if (sublinkColumns) {
+                        sublinkColumns->push_back(columnName + "_count");
+                        sublinkColumns->push_back(columnName + "_value");
+                    }
+
                     aggregateItems.push_back(ctx.Builder(node->Pos())
                         .List()
                             .Atom(0, columnName + "_count")
@@ -533,7 +533,7 @@ std::pair<TExprNode::TPtr, TExprNode::TPtr> RewriteSubLinks(TPositionHandle pos,
                     .Seal()
                     .Build();
 
-                newList = JoinColumns(pos, newList, groupedSublink, fullColList, it->second, ctx);
+                newList = JoinColumns(pos, newList, groupedSublink, fullColList, it->second, ctx, leftPrefix);
 
                 if (linkType == "exists") {
                     return ctx.Builder(node->Pos())
@@ -768,9 +768,9 @@ void AddColumnsFromType(const TTypeAnnotationNode* type, TUsedColumns& columns)
 
 void AddColumnsFromSublinks(const TNodeMap<ui32>& subLinks, TUsedColumns& columns) {
     for (const auto& s : subLinks) {
-        auto extColumns = ExtractExternalColumns(s.first->Tail());
+        auto extColumns = NTypeAnnImpl::ExtractExternalColumns(s.first->Tail());
         for (const auto& c : extColumns) {
-            columns.insert(std::make_pair(c, std::make_pair(Max<ui32>(), TString())));
+            columns.insert(std::make_pair(c.first, std::make_pair(Max<ui32>(), TString())));
         }
 
         if (!s.first->Child(2)->IsCallable("Void")) {
@@ -820,7 +820,9 @@ TUsedColumns GatherUsedColumns(const TExprNode::TPtr& result, const TExprNode::T
     if (extraSortColumns) {
         for (ui32 inputIndex = 0; inputIndex < extraSortColumns->Tail().ChildrenSize(); ++inputIndex) {
             for (const auto& x : extraSortColumns->Tail().Child(inputIndex)->Children()) {
-                usedColumns.insert(std::make_pair(TString(x->Content()), std::make_pair(inputIndex, TString())));
+                auto from = x->Content();
+                from.SkipPrefix("_yql_extra_");
+                usedColumns.insert(std::make_pair(TString(from), std::make_pair(inputIndex, TString())));
             }
         }
     }
@@ -2334,18 +2336,17 @@ TExprNode::TPtr AddExtColumns(const TExprNode::TPtr& lambda, const TExprNode::TP
         .Build();
 }
 
-void BuildExtraSortColumns(TExprNodeBuilder& parent, const TExprNode::TPtr& extraSortColumns, size_t aggIndexBegin, size_t aggIndexEnd) {
-    ui32 i = 0;
+void BuildExtraSortColumns(const TExprNode::TPtr& extraSortColumns, size_t aggIndexBegin, size_t aggIndexEnd, TVector<TString>& list) {
     if (extraSortColumns) {
         for (const auto& x : extraSortColumns->Tail().Children()) {
             for (const auto& y : x->Children()) {
-                parent.Atom(i++, y->Content());
+                list.push_back(TString(y->Content()));
             }
         }
     } 
 
     for (auto aggIndex = aggIndexBegin; aggIndex < aggIndexEnd; ++aggIndex) {
-        parent.Atom(i++, "_yql_agg_" + ToString(aggIndex));
+        list.push_back("_yql_agg_" + ToString(aggIndex));
     }
 }
 
@@ -2363,14 +2364,24 @@ TExprNode::TPtr AddExtraSortColumns(const TExprNode::TPtr& lambda, const TExprNo
                 .Seal()
                 .List(1)
                     .Atom(0, "")
-                    .Callable(1, "FilterMembers")
-                        .Arg(0, "row")
-                        .List(1)
-                            .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder & {
-                                BuildExtraSortColumns(parent, extraSortColumns, aggIndexBegin, aggIndexEnd);
-                                return parent;
-                            })
-                        .Seal()
+                    .Callable(1, "AsStruct")
+                        .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder & {
+                            TVector<TString> list;
+                            BuildExtraSortColumns(extraSortColumns, aggIndexBegin, aggIndexEnd, list);
+                            for (ui32 i = 0; i < list.size(); ++i) {
+                                TStringBuf from = list[i];
+                                from.SkipPrefix("_yql_extra_");
+                                parent.List(i)
+                                    .Atom(0, list[i])
+                                    .Callable(1, "Member")
+                                        .Arg(0, "row")
+                                        .Atom(1, from)
+                                    .Seal()
+                                    .Seal();
+                            }
+
+                            return parent;
+                        })
                     .Seal()
                 .Seal()
             .Seal()
@@ -2379,7 +2390,7 @@ TExprNode::TPtr AddExtraSortColumns(const TExprNode::TPtr& lambda, const TExprNo
 }
 
 TExprNode::TPtr RemoveExtraSortColumns(const TExprNode::TPtr& list, const TExprNode::TPtr& extraSortColumns,
-    size_t aggIndexBegin, size_t aggIndexEnd, TExprContext& ctx) {
+    size_t aggIndexBegin, size_t aggIndexEnd, const TVector<TString>& sublinkColumns, TExprContext& ctx) {
     return ctx.Builder(list->Pos())
         .Callable("OrderedMap")
             .Add(0, list)
@@ -2389,7 +2400,13 @@ TExprNode::TPtr RemoveExtraSortColumns(const TExprNode::TPtr& list, const TExprN
                     .Arg(0, "row")
                     .List(1)
                         .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder & {
-                            BuildExtraSortColumns(parent, extraSortColumns, aggIndexBegin, aggIndexEnd);
+                            TVector<TString> list;
+                            BuildExtraSortColumns(extraSortColumns, aggIndexBegin, aggIndexEnd, list);
+                            list.insert(list.end(), sublinkColumns.begin(), sublinkColumns.end());
+                            for (ui32 i = 0; i < list.size(); ++i) {
+                                parent.Atom(i, list[i]);
+                            }
+
                             return parent;
                         })
                     .Seal()
@@ -2486,7 +2503,7 @@ TExprNode::TPtr ExpandPgSelectImpl(const TExprNode::TPtr& node, TExprContext& ct
         auto distinctAll = GetSetting(setItem->Tail(), "distinct_all");
         auto distinctOn = GetSetting(setItem->Tail(), "distinct_on");
         auto sort = GetSetting(setItem->Tail(), "sort");
-        auto extraSortColumns = GetSetting(setItem->Tail(), "extra_sort_columns");
+        auto extraSortColumns = GetSetting(setItem->Tail(), "final_extra_sort_columns");
         bool oneRow = !from;
         TExprNode::TPtr list;
         if (values) {
@@ -2592,11 +2609,12 @@ TExprNode::TPtr ExpandPgSelectImpl(const TExprNode::TPtr& node, TExprContext& ct
                 list = BuildDistinctOn(node->Pos(), list, distinctOn->TailPtr(), sort, ctx);
             }
 
+            TVector<TString> sublinkColumns;
             if (sort) {
                 auto sortSubLinks = GatherSubLinks(sortLambda);
                 if (!sortSubLinks.empty()) {
                     std::tie(list, sortLambda) = RewriteSubLinks(sortLambda->Pos(), list, sortLambda,
-                        sortSubLinks, inputAliases, cleanedInputs, &aggId, ctx, optCtx);
+                        sortSubLinks, inputAliases, cleanedInputs, &aggId, ctx, optCtx, "_yql_extra_", &sublinkColumns);
                 } else {
                     RewriteAggs(sortLambda, aggId, ctx, optCtx, false);
                 }
@@ -2605,7 +2623,7 @@ TExprNode::TPtr ExpandPgSelectImpl(const TExprNode::TPtr& node, TExprContext& ct
             }
 
             if (hasExtraSortColumns) {
-                list = RemoveExtraSortColumns(list, extraSortColumns, aggsSizeBeforeSort, aggs.size(), ctx);
+                list = RemoveExtraSortColumns(list, extraSortColumns, aggsSizeBeforeSort, aggs.size(), sublinkColumns, ctx);
             }
         }
 
diff --git a/ydb/library/yql/core/type_ann/type_ann_pg.cpp b/ydb/library/yql/core/type_ann/type_ann_pg.cpp
index 2d510d01c01..d07c5254567 100644
--- a/ydb/library/yql/core/type_ann/type_ann_pg.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_pg.cpp
@@ -1560,7 +1560,7 @@ bool ScanColumns(TExprNode::TPtr root, TInputs& inputs, const THashSet<TString>&
 
 bool ScanColumnsForSublinks(bool& needRebuildSubLinks, const TNodeSet& sublinks,
     TInputs& inputs, const THashSet<TString>& possibleAliases, bool& hasColumnRef, THashSet<TString>& refs,
-    THashMap<TString, THashSet<TString>>* qualifiedRefs, TExtContext& ctx, bool scanColumnsOnly) {
+    THashMap<TString, THashSet<TString>>* qualifiedRefs, TExtContext& ctx, bool scanColumnsOnly, THashMap<ui32, TSet<TString>>* nonProjectionColumns = nullptr) {
     needRebuildSubLinks = false;
     for (const auto& s : sublinks) {
         if (s->Child(1)->IsCallable("Void")) {
@@ -1571,7 +1571,7 @@ bool ScanColumnsForSublinks(bool& needRebuildSubLinks, const TNodeSet& sublinks,
         if (!testRowLambda.IsCallable("Void")) {
             YQL_ENSURE(testRowLambda.IsLambda());
             if (!ScanColumns(testRowLambda.TailPtr(), inputs, possibleAliases, nullptr, hasColumnRef,
-                refs, qualifiedRefs, ctx, scanColumnsOnly)) {
+                refs, qualifiedRefs, ctx, scanColumnsOnly, nonProjectionColumns)) {
                 return false;
             }
         }
@@ -2006,11 +2006,97 @@ bool ValidateGroups(TInputs& inputs, const THashSet<TString>& possibleAliases,
     return true;
 }
 
-bool ValidateSort(TInputs& inputs, const THashSet<TString>& possibleAliases,
+THashMap<ui32, TSet<TString>> LoadExtraColumns(TExprNode::TPtr& value, ui32 inputsCount) {
+    THashMap<ui32, TSet<TString>> ret;
+    YQL_ENSURE(value->ChildrenSize() == inputsCount);
+    for (ui32 inputIndex = 0; inputIndex < inputsCount; ++inputIndex) {
+        auto child = value->Child(inputIndex);
+        if (child->ChildrenSize() == 0) {
+            continue;
+        }
+
+        auto& set = ret[inputIndex];
+        for (const auto& x : child->Children()) {
+            set.insert(TString(x->Content()));
+        }
+    }
+
+    return ret;
+}
+
+TExprNode::TPtr SaveExtraColumns(TPositionHandle pos, const THashMap<ui32, TSet<TString>>& columns, ui32 inputsCount, TExprContext& ctx) {
+    TExprNode::TListType groups;
+    for (ui32 i = 0; i < inputsCount; ++i) {
+        TExprNode::TListType columnsPerGroup;
+        auto it = columns.find(i);
+        if (it != columns.end()) {
+            for (const auto& x : it->second) {
+                columnsPerGroup.push_back(ctx.NewAtom(pos, x));
+            }
+        }
+
+        auto columnsPerGroupList = ctx.NewList(pos, std::move(columnsPerGroup));
+        groups.push_back(columnsPerGroupList);
+    }
+
+    return ctx.NewList(pos, std::move(groups));
+}
+
+bool AddExtraColumnsForSublinks(TExprNode::TPtr& value, const TExprNode& data, ui32 inputsCount, TExprContext& ctx) {
+    THashMap<ui32, TSet<TString>> columns = LoadExtraColumns(value, inputsCount);
+    for (auto oneSort : data.Children()) {
+        TNodeSet sublinks;
+        ScanSublinks(oneSort->Child(1)->TailPtr(), sublinks);
+        for (const auto& s : sublinks) {
+            auto c = ExtractExternalColumns(s->Tail());
+            for (const auto& [name, index] : c) {
+                YQL_ENSURE(index < inputsCount);
+                columns[index].insert("_yql_extra_" + name);
+            }
+        }
+    }
+
+    value = SaveExtraColumns(value->Pos(), columns, inputsCount, ctx);
+    for (const auto&[index, set] : columns) {
+        if (!set.empty()) {
+            return true;
+        }
+    }
+
+    return false;
+}
+
+TMap<TString, ui32> ExtractExternalColumns(const TExprNode& select) {
+    TMap<TString, ui32> res;
+    const auto& option = select.Head();
+    auto setItems = GetSetting(option, "set_items");
+    YQL_ENSURE(setItems);
+    for (const auto& s : setItems->Tail().Children()) {
+        YQL_ENSURE(s->IsCallable("PgSetItem"));
+        auto extTypes = GetSetting(s->Head(), "final_ext_types");
+        YQL_ENSURE(extTypes);
+        ui32 inputIndex = 0;
+        for (const auto& input : extTypes->Tail().Children()) {
+            auto type = input->Tail().GetTypeAnn()->Cast<TTypeExprType>()->GetType()->Cast<TStructExprType>();
+            for (const auto& item : type->GetItems()) {
+                res.insert(std::make_pair(NTypeAnnImpl::MakeAliasedColumn(input->Head().Content(), item->GetName()), inputIndex));
+            }
+
+            ++inputIndex;
+        }
+    }
+
+    return res;
+}
+
+bool ValidateSort(TInputs& inputs, TInputs& subLinkInputs, const THashSet<TString>& possibleAliases,
     const TExprNode& data, TExtContext& ctx, TExprNode::TListType& newSorts, bool scanColumnsOnly,
     THashMap<ui32, TSet<TString>>* nonProjectionColumns = nullptr) {
     newSorts.clear();
     for (auto oneSort : data.Children()) {
+        TNodeSet sublinks;
+        ScanSublinks(oneSort->Child(1)->TailPtr(), sublinks);
+
         bool hasColumnRef;
         THashSet<TString> refs;
         THashMap<TString, THashSet<TString>> qualifiedRefs;
@@ -2019,12 +2105,18 @@ bool ValidateSort(TInputs& inputs, const THashSet<TString>& possibleAliases,
             return false;
         }
 
+        bool needRebuildSubLinks;
+        if (!ScanColumnsForSublinks(needRebuildSubLinks, sublinks, subLinkInputs, possibleAliases,
+            hasColumnRef, refs, &qualifiedRefs, ctx, scanColumnsOnly, nonProjectionColumns)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
         if (scanColumnsOnly) {
             continue;
         }
 
         TVector<const TItemExprType*> items;
-        AddColumns(inputs, nullptr, refs, &qualifiedRefs, items, ctx.Expr);
+        AddColumns(needRebuildSubLinks ? subLinkInputs : inputs, nullptr, refs, &qualifiedRefs, items, ctx.Expr);
         auto effectiveType = ctx.Expr.MakeType<TStructExprType>(items);
         if (!effectiveType->Validate(oneSort->Pos(), ctx.Expr)) {
             return false;
@@ -2032,19 +2124,33 @@ bool ValidateSort(TInputs& inputs, const THashSet<TString>& possibleAliases,
 
         auto typeNode = ExpandType(oneSort->Pos(), *effectiveType, ctx.Expr);
 
-        auto argNode = ctx.Expr.NewArgument(oneSort->Pos(), "row");
-        auto arguments = ctx.Expr.NewArguments(oneSort->Pos(), { argNode });
-        TExprNode::TPtr newRoot;
-        auto status = RebuildLambdaColumns(oneSort->Child(1)->TailPtr(), argNode, newRoot, inputs, nullptr, ctx);
-        if (status == IGraphTransformer::TStatus::Error) {
-            return false;
-        }
+        auto newChildren = oneSort->ChildrenList();
+        if (needRebuildSubLinks) {
+            auto arguments = ctx.Expr.NewArguments(oneSort->Pos(), { });
+
+            TExprNode::TPtr newRoot;
+            auto status = RebuildSubLinks(oneSort->Child(1)->TailPtr(), newRoot, sublinks, subLinkInputs, typeNode, ctx);
+            if (status == IGraphTransformer::TStatus::Error) {
+                return false;
+            }
 
-        auto newLambda = ctx.Expr.NewLambda(oneSort->Pos(), std::move(arguments), std::move(newRoot));
+            auto newLambda = ctx.Expr.NewLambda(oneSort->Pos(), std::move(arguments), std::move(newRoot));
+            newChildren[1] = newLambda;
+        } else if (oneSort->Child(1)->Head().ChildrenSize() == 0) {
+            auto argNode = ctx.Expr.NewArgument(oneSort->Pos(), "row");
+            auto arguments = ctx.Expr.NewArguments(oneSort->Pos(), { argNode });
+            TExprNode::TPtr newRoot;
+            auto status = RebuildLambdaColumns(oneSort->Child(1)->TailPtr(), argNode, newRoot, inputs, nullptr, ctx);
+            if (status == IGraphTransformer::TStatus::Error) {
+                return false;
+            }
+
+            auto newLambda = ctx.Expr.NewLambda(oneSort->Pos(), std::move(arguments), std::move(newRoot));
+
+            newChildren[0] = typeNode;
+            newChildren[1] = newLambda;
+        }
 
-        auto newChildren = oneSort->ChildrenList();
-        newChildren[0] = typeNode;
-        newChildren[1] = newLambda;
         auto newSort = ctx.Expr.ChangeChildren(*oneSort, std::move(newChildren));
         newSorts.push_back(newSort);
     }
@@ -2075,6 +2181,7 @@ IGraphTransformer::TStatus PgSetItemWrapper(const TExprNode::TPtr& input, TExprN
         bool hasDistinctAll = false;
         bool hasDistinctOn = false;
         bool hasExtraSortColumns = false;
+        bool hasFinalExtraSortColumns = false;
 
         // pass 0 - from/values
         // pass 1 - join
@@ -2236,7 +2343,7 @@ IGraphTransformer::TStatus PgSetItemWrapper(const TExprNode::TPtr& input, TExprN
 
                                     auto newLambda = ctx.Expr.NewLambda(column->Pos(), std::move(arguments), std::move(newRoot));
                                     newColumnChildren[2] = newLambda;
-                                } else {
+                                } else if (column->Tail().Head().ChildrenSize() == 0) {
                                     auto argNode = ctx.Expr.NewArgument(column->Pos(), "row");
                                     auto arguments = ctx.Expr.NewArguments(column->Pos(), { argNode });
                                     auto expandedColumns = column->HeadPtr();
@@ -2475,25 +2582,25 @@ IGraphTransformer::TStatus PgSetItemWrapper(const TExprNode::TPtr& input, TExprN
                                 auto newSettings = ReplaceSetting(options, {}, TString(optionName), newWhere, ctx.Expr);
                                 output = ctx.Expr.ChangeChild(*input, 0, std::move(newSettings));
                                 return IGraphTransformer::TStatus::Repeat;
-                            }
-
-                            auto argNode = ctx.Expr.NewArgument(data.Pos(), "row");
-                            auto arguments = ctx.Expr.NewArguments(data.Pos(), { argNode });
-                            TExprNode::TPtr newRoot;
-                            auto status = RebuildLambdaColumns(data.Child(1)->TailPtr(), argNode, newRoot, joinInputs, nullptr, ctx);
-                            if (status == IGraphTransformer::TStatus::Error) {
-                                return IGraphTransformer::TStatus::Error;
-                            }
+                            } else if (data.Child(1)->Head().ChildrenSize() == 0) {
+                                auto argNode = ctx.Expr.NewArgument(data.Pos(), "row");
+                                auto arguments = ctx.Expr.NewArguments(data.Pos(), { argNode });
+                                TExprNode::TPtr newRoot;
+                                auto status = RebuildLambdaColumns(data.Child(1)->TailPtr(), argNode, newRoot, joinInputs, nullptr, ctx);
+                                if (status == IGraphTransformer::TStatus::Error) {
+                                    return IGraphTransformer::TStatus::Error;
+                                }
 
-                            auto newLambda = ctx.Expr.NewLambda(data.Pos(), std::move(arguments), std::move(newRoot));
+                                auto newLambda = ctx.Expr.NewLambda(data.Pos(), std::move(arguments), std::move(newRoot));
 
-                            auto newChildren = data.ChildrenList();
-                            newChildren[0] = typeNode;
-                            newChildren[1] = newLambda;
-                            auto newWhere = ctx.Expr.NewCallable(data.Pos(), "PgWhere", std::move(newChildren));
-                            auto newSettings = ReplaceSetting(options, {}, TString(optionName), newWhere, ctx.Expr);
-                            output = ctx.Expr.ChangeChild(*input, 0, std::move(newSettings));
-                            return IGraphTransformer::TStatus::Repeat;
+                                auto newChildren = data.ChildrenList();
+                                newChildren[0] = typeNode;
+                                newChildren[1] = newLambda;
+                                auto newWhere = ctx.Expr.NewCallable(data.Pos(), "PgWhere", std::move(newChildren));
+                                auto newSettings = ReplaceSetting(options, {}, TString(optionName), newWhere, ctx.Expr);
+                                output = ctx.Expr.ChangeChild(*input, 0, std::move(newSettings));
+                                return IGraphTransformer::TStatus::Repeat;
+                            }
                         }
                     }
                     else {
@@ -2801,7 +2908,7 @@ IGraphTransformer::TStatus PgSetItemWrapper(const TExprNode::TPtr& input, TExprN
 
                         if (needRebuildSort) {
                             TExprNode::TListType newSorts;
-                            if (!ValidateSort(joinInputs, possibleAliases, *sort, ctx, newSorts, scanColumnsOnly)) {
+                            if (!ValidateSort(joinInputs, joinInputs, possibleAliases, *sort, ctx, newSorts, scanColumnsOnly)) {
                                 return IGraphTransformer::TStatus::Error;
                             }
 
@@ -2879,10 +2986,15 @@ IGraphTransformer::TStatus PgSetItemWrapper(const TExprNode::TPtr& input, TExprN
                         return IGraphTransformer::TStatus::Error;
                     }
 
+                    bool needScan = false;
                     for (const auto& x : data.Children()) {
                         if (!x->IsCallable("PgSort")) {
                             ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(x->Pos()), "Expected PgSort"));
                         }
+
+                        if (x->Child(0)->IsCallable("Void")) {
+                            needScan = true;
+                        }
                     }
 
                     TInputs projectionInputs = joinInputs;
@@ -2891,45 +3003,44 @@ IGraphTransformer::TStatus PgSetItemWrapper(const TExprNode::TPtr& input, TExprN
                         projectionInputs.push_back(TInput{ "", outputRowType, Nothing(), TInput::Projection, {} });
                     }
 
-                    if (data.ChildrenSize() > 0 && data.Child(0)->Child(0)->IsCallable("Void")) {
+                    if (needScan) {
                         TExprNode::TListType newSortTupleItems;
                         // no effective types yet, scan lambda bodies
-                        THashMap<ui32, TSet<TString>> extraSortColumns;

-                        if (!ValidateSort(projectionInputs, possibleAliases, data, ctx, newSortTupleItems, scanColumnsOnly, &extraSortColumns)) {

-                            return IGraphTransformer::TStatus::Error;

-                        }

+                        THashMap<ui32, TSet<TString>> extraSortColumns;
+                        if (!ValidateSort(projectionInputs, joinInputs, possibleAliases, data, ctx, newSortTupleItems, scanColumnsOnly, &extraSortColumns)) {
+                            return IGraphTransformer::TStatus::Error;
+                        }
 
                         if (!scanColumnsOnly) {
-                            auto newSortTuple = ctx.Expr.NewList(data.Pos(), std::move(newSortTupleItems));

+                            auto newSortTuple = ctx.Expr.NewList(data.Pos(), std::move(newSortTupleItems));
                             auto newSettings = ReplaceSetting(options, {}, "sort", newSortTuple, ctx.Expr);
-                            if (!extraSortColumns.empty()) {
-                                TExprNode::TListType groups;
-                                for (ui32 i = 0; i < joinInputs.size(); ++i) {
-                                    TExprNode::TListType columns;
-                                    auto it = extraSortColumns.find(i);
-                                    if (it != extraSortColumns.end()) {
-                                        for (const auto& x : it->second) {
-                                            columns.push_back(ctx.Expr.NewAtom(data.Pos(), x));
-                                        }
-                                    }
+                            auto extra = SaveExtraColumns(data.Pos(), extraSortColumns, joinInputs.size(), ctx.Expr);
+                            newSettings = AddSetting(*newSettings, {}, "extra_sort_columns", extra, ctx.Expr);
 
-                                    auto columnsList = ctx.Expr.NewList(data.Pos(), std::move(columns));
-                                    groups.push_back(columnsList);
-                                }
-

-                                newSettings = AddSetting(*newSettings, {}, "extra_sort_columns", ctx.Expr.NewList(data.Pos(), std::move(groups)), ctx.Expr);

+                            output = ctx.Expr.ChangeChild(*input, 0, std::move(newSettings));
+                            return IGraphTransformer::TStatus::Repeat;
+                        }
+                    } else {
+                        auto extra = GetSetting(options, "extra_sort_columns");
+                        if (extra) {
+                            auto value = extra->TailPtr();
+                            bool hasColumns = AddExtraColumnsForSublinks(value, data, joinInputs.size(), ctx.Expr);
+                            auto newSettings = RemoveSetting(options, "extra_sort_columns", ctx.Expr);
+                            if (hasColumns) {
+                                newSettings = AddSetting(*newSettings, {}, "final_extra_sort_columns", value, ctx.Expr);
                             }
-

-                            output = ctx.Expr.ChangeChild(*input, 0, std::move(newSettings));

+
+                            output = ctx.Expr.ChangeChild(*input, 0, std::move(newSettings));
                             return IGraphTransformer::TStatus::Repeat;
                         }
                     }
-                } else if (optionName == "extra_sort_columns") {
+                } else if (optionName == "extra_sort_columns" || optionName == "final_extra_sort_columns") {
                     if (pass != 2) {
                         continue;
                     }
 
-                    hasExtraSortColumns = true;
+                    (optionName == "extra_sort_columns" ? hasExtraSortColumns : hasFinalExtraSortColumns) = true;
+
                     if (!EnsureTupleSize(*option, 2, ctx.Expr)) {
                         return IGraphTransformer::TStatus::Error;
                     }
@@ -2972,7 +3083,7 @@ IGraphTransformer::TStatus PgSetItemWrapper(const TExprNode::TPtr& input, TExprN
             return IGraphTransformer::TStatus::Error;
         }
 
-        if ((hasDistinctAll || hasDistinctOn) && hasExtraSortColumns) {
+        if ((hasDistinctAll || hasDistinctOn) && hasFinalExtraSortColumns) {
             ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()), "for SELECT DISTINCT, ORDER BY expressions must appear in select list"));
             return IGraphTransformer::TStatus::Error;
         }
@@ -3226,7 +3337,7 @@ IGraphTransformer::TStatus PgSelectWrapper(const TExprNode::TPtr& input, TExprNo
 
         if (data.ChildrenSize() > 0 && data.Child(0)->Child(0)->IsCallable("Void")) {
             // no effective types yet, scan lambda bodies
-            if (!ValidateSort(projectionInputs, {}, data, ctx, newSortTupleItems, false)) {
+            if (!ValidateSort(projectionInputs, projectionInputs, {}, data, ctx, newSortTupleItems, false)) {
                 return IGraphTransformer::TStatus::Error;
             }
 
diff --git a/ydb/library/yql/core/type_ann/type_ann_pg.h b/ydb/library/yql/core/type_ann/type_ann_pg.h
index 8d985e3ec54..5cb1a257563 100644
--- a/ydb/library/yql/core/type_ann/type_ann_pg.h
+++ b/ydb/library/yql/core/type_ann/type_ann_pg.h
@@ -13,6 +13,7 @@ const TItemExprType* AddAlias(const TString& alias, const TItemExprType* item, T
 TStringBuf RemoveAlias(TStringBuf column);
 TStringBuf RemoveAlias(TStringBuf column, TStringBuf& alias);
 const TItemExprType* RemoveAlias(const TItemExprType* item, TExprContext& ctx);
+TMap<TString, ui32> ExtractExternalColumns(const TExprNode& select);
 
 IGraphTransformer::TStatus PgStarWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
 IGraphTransformer::TStatus PgCallWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx);