Improve sorted constraint.

7 files changed, 188 insertions, 167 deletions

diff --git a/ydb/library/yql/ast/yql_constraint.cpp b/ydb/library/yql/ast/yql_constraint.cpp
index 80b4c0ada66..e6d28354da7 100644
--- a/ydb/library/yql/ast/yql_constraint.cpp
+++ b/ydb/library/yql/ast/yql_constraint.cpp
@@ -29,9 +29,9 @@ void TConstraintNode::Out(IOutputStream& out) const {
     out.Write(Name_);
 }
 
-bool TConstraintNode::PathExistsInType(const TPathType& path, const TTypeAnnotationNode& type) {
+const TTypeAnnotationNode* TConstraintNode::GetSubTypeByPath(const TPathType& path, const TTypeAnnotationNode& type) {
     if (path.empty())
-        return true;
+        return &type;
 
     const auto tail = [](const TPathType& path) {
         auto p(path);
@@ -40,25 +40,25 @@ bool TConstraintNode::PathExistsInType(const TPathType& path, const TTypeAnnotat
     };
     switch (type.GetKind()) {
         case ETypeAnnotationKind::Optional:
-            return PathExistsInType(path, *type.Cast<TOptionalExprType>()->GetItemType());
+            return GetSubTypeByPath(path, *type.Cast<TOptionalExprType>()->GetItemType());
         case ETypeAnnotationKind::Struct:
             if (const auto itemType = type.Cast<TStructExprType>()->FindItemType(path.front()))
-                return PathExistsInType(tail(path), *itemType);
+                return GetSubTypeByPath(tail(path), *itemType);
             break;
         case ETypeAnnotationKind::Tuple:
             if (const auto index = TryFromString<ui64>(TStringBuf(path.front())))
                 if (const auto typleType = type.Cast<TTupleExprType>(); typleType->GetSize() > *index)
-                    return PathExistsInType(tail(path), *typleType->GetItems()[*index]);
+                    return GetSubTypeByPath(tail(path), *typleType->GetItems()[*index]);
             break;
         case ETypeAnnotationKind::Multi:
             if (const auto index = TryFromString<ui64>(TStringBuf(path.front())))
                 if (const auto multiType = type.Cast<TMultiExprType>(); multiType->GetSize() > *index)
-                    return PathExistsInType(tail(path), *multiType->GetItems()[*index]);
+                    return GetSubTypeByPath(tail(path), *multiType->GetItems()[*index]);
             break;
         default:
             break;
     }
-    return false;
+    return nullptr;
 }
 
 //////////////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -209,7 +209,8 @@ TSortedConstraintNode::TSortedConstraintNode(TExprContext& ctx, TContainerType&&
     YQL_ENSURE(!Content_.empty());
     for (const auto& c : Content_) {
         YQL_ENSURE(!c.first.empty());
-        Hash_ = std::accumulate(c.first.cbegin(), c.first.cend(), c.second ? Hash_ : ~Hash_, [](ui64 hash, const std::string_view& field) { return MurmurHash<ui64>(field.data(), field.size(), hash); });
+        for (const auto& path : c.first)
+            Hash_ = std::accumulate(path.cbegin(), path.cend(), c.second ? Hash_ : ~Hash_, [](ui64 hash, const std::string_view& field) { return MurmurHash<ui64>(field.data(), field.size(), hash); });
     }
 }
 
@@ -249,13 +250,16 @@ bool TSortedConstraintNode::Includes(const TConstraintNode& node) const {
 void TSortedConstraintNode::Out(IOutputStream& out) const {
     TConstraintNode::Out(out);
     out.Write('(');
-    for (size_t i = 0; i < Content_.size(); ++i) {
-        if (i > 0) {
+    bool first = true;
+    for (const auto& c : Content_) {
+        if (first)
+            first = false;
+        else
             out.Write(';');
-        }
-        out.Write(JoinSeq(',', Content_[i].first));
+
+        out.Write(JoinSeq(',', c.first));
         out.Write('[');
-        out.Write(Content_[i].second ? "asc" : "desc");
+        out.Write(c.second ? "asc" : "desc");
         out.Write(']');
     }
     out.Write(')');
@@ -263,10 +267,10 @@ void TSortedConstraintNode::Out(IOutputStream& out) const {
 
 void TSortedConstraintNode::ToJson(NJson::TJsonWriter& out) const {
     out.OpenArray();
-    for (size_t i = 0; i < Content_.size(); ++i) {
+    for (const auto& c : Content_) {
         out.OpenArray();
-        out.Write(JoinSeq(';', Content_[i].first));
-        out.Write(Content_[i].second);
+        out.Write(JoinSeq(';', c.first));
+        out.Write(c.second);
         out.CloseArray();
     }
     out.CloseArray();
@@ -279,16 +283,15 @@ bool TSortedConstraintNode::IsPrefixOf(const TSortedConstraintNode& node) const
 bool TSortedConstraintNode::IsOrderBy(const TUniqueConstraintNode& unique) const {
     NSorted::TSimpleSet<TColumnsSet> columns;
     for (const auto& key : Content_)
-        if (!key.first.empty())
-            columns.insert_unique(key.first);
+        columns.insert_unique(key.first);
 
     const auto ordered = columns;
     for (const auto& set : unique.GetAllSets()) {
         if (std::all_of(set.cbegin(), set.cend(), [&ordered](const TPathType& path) {
-            return !path.empty() && std::any_of(ordered.cbegin(), ordered.cend(), [&path](const TColumnsSet& s) { return s.contains(path.front()); });
+            return !path.empty() && std::any_of(ordered.cbegin(), ordered.cend(), [&path](const TColumnsSet& s) { return s.contains(path); });
         })) {
             std::for_each(set.cbegin(), set.cend(), [&columns](const TPathType& path) {
-                if (const auto it = std::find_if(columns.cbegin(), columns.cend(), [&path](const TColumnsSet& s) { return s.contains(path.front()); }); columns.cend() != it)
+                if (const auto it = std::find_if(columns.cbegin(), columns.cend(), [&path](const TColumnsSet& s) { return s.contains(path); }); columns.cend() != it)
                     columns.erase(it);
             });
             if (columns.empty())
@@ -372,7 +375,7 @@ const TSortedConstraintNode* TSortedConstraintNode::FilterByType(const TSortedCo
         auto content = sorted->GetContent();
         for (size_t i = 0; i < content.size(); ++i) {
             for (auto it = content[i].first.cbegin(); content[i].first.cend() != it;) {
-                if (outItemType->FindItem(*it))
+                if (GetSubTypeByPath(*it, *outItemType))
                     ++it;
                 else
                     it = content[i].first.erase(it);
@@ -403,6 +406,13 @@ const TSortedConstraintNode* TSortedConstraintNode::CutPrefix(size_t newPrefixLe
     return ctx.MakeConstraint<TSortedConstraintNode>(std::move(content));
 }
 
+bool TSortedConstraintNode::IsApplicableToType(const TTypeAnnotationNode& type) const {
+    const auto& itemType = GetSeqItemType(type);
+    return std::all_of(Content_.cbegin(), Content_.cend(), [&itemType](const std::pair<TColumnsSet, bool>& pair) {
+        return std::all_of(pair.first.cbegin(), pair.first.cend(), std::bind(&GetSubTypeByPath, std::placeholders::_1, std::cref(itemType)));
+    });
+}
+
 //////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
 TGroupByConstraintNode::TGroupByConstraintNode(TExprContext& ctx, const std::vector<TStringBuf>& columns)
@@ -546,19 +556,11 @@ void TUniqueConstraintNode::Out(IOutputStream& out) const {
         out.Write('(');
         bool first = true;
         for (const auto& path : set) {
-            if (!first) {
+            if (first)
+                first = false;
+            else
                 out.Write(',');
-            }
-            if (!path.empty()) {
-                auto it = path.cbegin();
-                out.Write(*it);
-                while (path.cend() > ++it) {
-                    out.Write('#');
-                    out.Write(*it);
-                }
-            }
-
-            first = false;
+            out << path;
         }
         out.Write(')');
     }
@@ -667,7 +669,7 @@ const TUniqueConstraintNode* TUniqueConstraintNode::RenameFields(TExprContext& c
 bool TUniqueConstraintNode::IsApplicableToType(const TTypeAnnotationNode& type) const {
     const auto& itemType = GetSeqItemType(type);
     return std::all_of(Sets_.cbegin(), Sets_.cend(), [&itemType](const TSetType& set) {
-        return std::all_of(set.cbegin(), set.cend(), std::bind(&PathExistsInType, std::placeholders::_1, std::cref(itemType)));
+        return std::all_of(set.cbegin(), set.cend(), std::bind(&GetSubTypeByPath, std::placeholders::_1, std::cref(itemType)));
     });
 }
 
@@ -726,32 +728,17 @@ bool TPartOfUniqueConstraintNode::Includes(const TConstraintNode& node) const {
 void TPartOfUniqueConstraintNode::Out(IOutputStream& out) const {
     TConstraintNode::Out(out);
     out.Write('(');
-
     bool first = true;
     for (const auto& part : Mapping_) {
         for (const auto& item : part.second) {
-            if (!first) {
+            if (first)
+                first = false;
+            else
                 out.Write(',');
-            }
-            if (!item.first.empty()) {
-                auto it = item.first.cbegin();
-                out.Write(*it);
-                while (item.first.cend() > ++it) {
-                    out.Write('#');
-                    out.Write(*it);
-                }
-            }
-            out.Write(':');
-            if (!item.second.empty()) {
-                auto it = item.second.cbegin();
-                out.Write(*it);
-                while (item.second.cend() > ++it) {
-                    out.Write('#');
-                    out.Write(*it);
-                }
-            }
 
-            first = false;
+            out << item.first;
+            out.Write(':');
+            out << item.second;
         }
     }
     out.Write(')');
@@ -1596,6 +1583,22 @@ const TMultiConstraintNode* TMultiConstraintNode::MakeCommon(const std::vector<c
 //////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
 template<>
+void Out<NYql::TConstraintNode::TPathType>(IOutputStream& out, const NYql::TConstraintNode::TPathType& path) {
+    if (path.empty())
+        out.Write('/');
+    else {
+        bool first = true;
+        for (const auto& c : path) {
+            if (first)
+                first = false;
+            else
+                out.Write('/');
+            out.Write(c);
+        }
+    }
+}
+
+template<>
 void Out<NYql::TConstraintNode>(IOutputStream& out, const NYql::TConstraintNode& c) {
     c.Out(out);
 }
diff --git a/ydb/library/yql/ast/yql_constraint.h b/ydb/library/yql/ast/yql_constraint.h
index e274cc6d08a..9c14a276b68 100644
--- a/ydb/library/yql/ast/yql_constraint.h
+++ b/ydb/library/yql/ast/yql_constraint.h
@@ -86,7 +86,8 @@ public:
         return Name_;
     }
 
-    static bool PathExistsInType(const TPathType& path, const TTypeAnnotationNode& type);
+    static const TTypeAnnotationNode* GetSubTypeByPath(const TPathType& path, const TTypeAnnotationNode& type);
+    static TString PathToString(const TPathType& path);
 protected:
     ui64 Hash_;
     std::string_view Name_;
@@ -205,7 +206,7 @@ private:
 
 class TSortedConstraintNode final: public TConstraintNode {
 public:
-    using TColumnsSet = NSorted::TSimpleSet<std::string_view>;
+    using TColumnsSet = NSorted::TSimpleSet<TPathType>;
     using TContainerType = TSmallVec<std::pair<TColumnsSet, bool>>;
 private:
     friend struct TExprContext;
@@ -235,6 +236,7 @@ public:
     const TSortedConstraintNode* MakeCommon(const TSortedConstraintNode* other, TExprContext& ctx) const;
     static const TSortedConstraintNode* FilterByType(const TSortedConstraintNode* sorted, const TStructExprType* outItemType, TExprContext& ctx);
 
+    bool IsApplicableToType(const TTypeAnnotationNode& type) const override;
 protected:
     TContainerType Content_;
 };
diff --git a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
index 7f98ac8f2eb..88d5c4153f4 100644
--- a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
@@ -119,10 +119,15 @@ TExprNode::TPtr KeepSortedConstraint(TExprNode::TPtr node, const TSortedConstrai
                     .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& {
                         size_t index = 0;
                         for (auto c : constent) {
-                            parent.Callable(index++, "Member")
-                                .Arg(0, "item")
-                                .Atom(1, c.first.front())
-                            .Seal();
+                            if (c.first.front().empty())
+                                parent.Arg(index++, "item");
+                            else {
+                                YQL_ENSURE(c.first.front().size() == 1U, "Just column expected.");
+                                parent.Callable(index++, "Member")
+                                    .Arg(0, "item")
+                                    .Atom(1, c.first.front().front())
+                                .Seal();
+                            }
                         }
                         return parent;
                     })
@@ -2726,10 +2731,15 @@ void FixSortness(const TExprNode& origNode, TExprNode::TPtr& node, TExprContext&
                         .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& {
                             size_t index = 0;
                             for (auto c : content) {
-                                parent.Callable(index++, "Member")
-                                        .Arg(0, "item")
-                                        .Atom(1, c.first.front())
-                                        .Seal();
+                                if (c.first.empty())
+                                    parent.Arg(0, "item");
+                                else {
+                                    YQL_ENSURE(c.first.front().size() == 1U, "Just column expected.");
+                                    parent.Callable(index++, "Member")
+                                            .Arg(0, "item")
+                                            .Atom(1, c.first.front().front())
+                                            .Seal();
+                                }
                             }
                             return parent;
                         })
diff --git a/ydb/library/yql/core/yql_expr_constraint.cpp b/ydb/library/yql/core/yql_expr_constraint.cpp
index 3fe4032a4d8..67fd91c4fc5 100644
--- a/ydb/library/yql/core/yql_expr_constraint.cpp
+++ b/ydb/library/yql/core/yql_expr_constraint.cpp
@@ -83,7 +83,7 @@ public:
         Functions["Unordered"] = &TCallableConstraintTransformer::FromFirst<TPassthroughConstraintNode, TEmptyConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
         Functions["UnorderedSubquery"] = &TCallableConstraintTransformer::FromFirst<TPassthroughConstraintNode, TEmptyConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
         Functions["Sort"] = &TCallableConstraintTransformer::SortWrap;
-        Functions["AssumeSorted"] = &TCallableConstraintTransformer::AssumeSortedWrap;
+        Functions["AssumeSorted"] = &TCallableConstraintTransformer::SortWrap;
         Functions["AssumeUnique"] = &TCallableConstraintTransformer::AssumeUniqueWrap;
         Functions["AssumeColumnOrder"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
         Functions["AssumeAllMembersNullableAtOnce"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
@@ -339,30 +339,17 @@ private:
     }
 
     TStatus SortWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
-        auto status = UpdateLambdaConstraints(*input->Child(2));
-        if (status != TStatus::Ok) {
+        if (const auto status = UpdateLambdaConstraints(input->Tail()); status != TStatus::Ok) {
             return status;
         }
 
-        if (auto sorted = DeduceSortConstraint(*input->Child(0), *input->Child(1), *input->Child(2), ctx)) {
+        if (const auto sorted = DeduceSortConstraint(*input->Child(1), *input->Child(2), ctx)) {
             input->AddConstraint(sorted);
         }
 
         return FromFirst<TPassthroughConstraintNode, TEmptyConstraintNode, TUniqueConstraintNode, TVarIndexConstraintNode>(input, output, ctx);
     }
 
-    TStatus AssumeSortedWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
-        auto status = UpdateLambdaConstraints(*input->Child(2));
-        if (status != TStatus::Ok) {
-            return status;
-        }
-
-        if (auto assumeConstr = DeduceSortConstraint(*input->Child(0), *input->Child(1), *input->Child(2), ctx)) {
-            input->AddConstraint(assumeConstr);
-        }
-        return FromFirst<TPassthroughConstraintNode, TEmptyConstraintNode, TUniqueConstraintNode, TVarIndexConstraintNode>(input, output, ctx);
-    }
-
     TStatus AssumeUniqueWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
         TUniqueConstraintNode::TFullSetType sets;
         for (auto i = 1U; i < input->ChildrenSize(); ++i) {
@@ -382,13 +369,12 @@ private:
 
     template <bool UseSort>
     TStatus TopWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
-        auto status = UpdateLambdaConstraints(*input->Child(3));
-        if (status != TStatus::Ok) {
+        if (const auto status = UpdateLambdaConstraints(input->Tail()); status != TStatus::Ok) {
             return status;
         }
 
         if constexpr (UseSort) {
-            if (auto sorted = DeduceSortConstraint(*input->Child(0), *input->Child(2), *input->Child(3), ctx)) {
+            if (const auto sorted = DeduceSortConstraint(*input->Child(2), *input->Child(3), ctx)) {
                 input->AddConstraint(sorted);
             }
         }
@@ -727,6 +713,56 @@ private:
             return GetLambdaConstraint<TConstraintType>(lambda, ctx);
     }
 
+    static std::vector<std::pair<TConstraintNode::TPathType, bool>>
+    ExtractSimpleSortTraits(const TExprNode& sortDirections, const TExprNode& keySelectorLambda) {
+        const auto& keySelectorBody = keySelectorLambda.Tail();
+        const auto& keySelectorArg = keySelectorLambda.Head().Head();
+        std::vector<std::pair<TConstraintNode::TPathType, bool>> columns;
+        if (sortDirections.IsCallable("Bool"))
+            columns.emplace_back(TConstraintNode::TPathType(), IsTrue(sortDirections.Tail().Content()));
+        else if (sortDirections.IsList())
+            if (const auto size = keySelectorBody.ChildrenSize()) {
+                columns.reserve(size);
+                for (auto i = 0U; i < size; ++i)
+                    if (const auto child = sortDirections.Child(i); child->IsCallable("Bool"))
+                        columns.emplace_back(TConstraintNode::TPathType(), IsTrue(child->Tail().Content()));
+                    else
+                        return {};
+            } else
+                return {};
+        else
+            return  {};
+
+        if (&keySelectorBody == &keySelectorArg)
+            columns.resize(1U);
+        else if (keySelectorBody.IsCallable({"Member", "Nth"}) && &keySelectorBody.Head() == &keySelectorArg)
+            if (columns.size() == 1U)
+                columns.front().first.emplace_back(keySelectorBody.Tail().Content());
+            else
+                return {};
+        else if (keySelectorBody.IsList())
+            if (const auto size = keySelectorBody.ChildrenSize()) {
+                std::unordered_set<std::string_view> set(size);
+                columns.resize(size, std::make_pair(TConstraintNode::TPathType(), columns.back().second));
+                auto it = columns.begin();
+                for (auto i = 0U; i < columns.size(); ++i) {
+                    if (const auto child = keySelectorBody.Child(i); child->IsCallable({"Member", "Nth"}) && &child->Head() == &keySelectorArg) {
+                        if (set.emplace(child->Tail().Content()).second)
+                            it++->first.emplace_back(child->Tail().Content());
+                        else if (columns.cend() != it)
+                            it = columns.erase(it);
+                    } else {
+                        columns.resize(i);
+                    }
+                }
+            } else
+                return {};
+        else
+            return {};
+
+        return columns;
+    }
+
     static TConstraintNode::TListType GetConstraintsForInputArgument(const TExprNode& node, std::unordered_set<const TPassthroughConstraintNode*>& explicitPasstrought, TExprContext& ctx) {
         TConstraintNode::TListType constraints;
         if (const auto inItemType = GetSeqItemType(node.Head().GetTypeAnn())) {
@@ -1081,30 +1117,43 @@ private:
         if (auto sort = MakeCommonConstraint<TSortedConstraintNode>(input, 0, ctx)) {
             if (Union && input->ChildrenSize() > 1) {
                 // Check and exclude modified keys from final constraint
-                auto resultStruct = input->GetTypeAnn()->Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>();
-                std::vector<const TStructExprType*> inputStructs;
-                for (auto& child: input->Children()) {
-                    inputStructs.push_back(child->GetTypeAnn()->Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>());
-                }
-                auto commonPrefixLength = sort->GetContent().size();
-                for (size_t i = 0; i < commonPrefixLength; ++i) {
-                    auto column = sort->GetContent()[i].first.front();
-                    auto pos = resultStruct->FindItem(column);
-                    YQL_ENSURE(pos, "Missing column " << TString{column}.Quote() << " in result type");
-                    auto resultItemType = resultStruct->GetItems()[*pos];
-                    for (size_t childNdx = 0; childNdx < input->ChildrenSize(); ++childNdx) {
-                        const auto inputStruct = inputStructs[childNdx];
-                        if (auto pos = inputStruct->FindItem(column)) {
-                            if (resultItemType != inputStruct->GetItems()[*pos]) {
-                                commonPrefixLength = i;
-                                break;
+                const auto resultItemType = input->GetTypeAnn()->Cast<TListExprType>()->GetItemType();
+                std::vector<const TTypeAnnotationNode*> inputs;
+                for (const auto& child: input->Children()) {
+                    inputs.emplace_back(child->GetTypeAnn()->Cast<TListExprType>()->GetItemType());
+                }
+
+                auto content = sort->GetContent();
+                for (auto i = 0U; i < content.size(); ++i) {
+                    for (auto it = content[i].first.cbegin(); content[i].first.cend() != it;) {
+                        const auto resultItemSubType = TConstraintNode::GetSubTypeByPath(*it, *resultItemType);
+                        YQL_ENSURE(resultItemSubType, "Missing " << *it << " in result type");
+                        auto childNdx = 0U;
+                        while (childNdx < input->ChildrenSize()) {
+                            if (const auto inputItemSubType = TConstraintNode::GetSubTypeByPath(*it, *inputs[childNdx])) {
+                                if (IsSameAnnotation(*inputItemSubType, *resultItemSubType)) {
+                                    ++childNdx;
+                                    continue;
+                                }
+                            } else {
+                                YQL_ENSURE(input->Child(childNdx)->GetConstraint<TEmptyConstraintNode>(), "Missing column " << *it << " in non empty input type");
                             }
-                        } else {
-                            YQL_ENSURE(input->Child(childNdx)->GetConstraint<TEmptyConstraintNode>(), "Missing column " << TString{column}.Quote() << " in non empty input type");
+                            break;
                         }
+
+                        if (childNdx < input->ChildrenSize())
+                            ++it;
+                        else
+                            it = content[i].first.erase(it);
+                    }
+
+                    if (content[i].first.empty()) {
+                        content.resize(i);
+                        break;
                     }
                 }
-                sort = sort->CutPrefix(commonPrefixLength, ctx);
+
+                sort = content.empty() ? nullptr : ctx.MakeConstraint<TSortedConstraintNode>(std::move(content));
             }
             if (sort) {
                 input->AddConstraint(sort);
@@ -2505,15 +2554,13 @@ private:
         return structType->GetSize() ? structType : nullptr;
     }
 
-    static const TSortedConstraintNode* DeduceSortConstraint(const TExprNode& list, const TExprNode& directions, const TExprNode& keyExtractor, TExprContext& ctx) {
-        if (GetSeqItemType(*list.GetTypeAnn()).GetKind() == ETypeAnnotationKind::Struct) {
-            if (const auto& columns = ExtractSimpleSortTraits(directions, keyExtractor); !columns.empty()) {
-                TSortedConstraintNode::TContainerType content(columns.size());
-                std::transform(columns.cbegin(), columns.cend(), content.begin(), [](const std::pair<std::string_view, bool>& item) {
-                   return TSortedConstraintNode::TContainerType::value_type({item.first}, item.second);
-                });
-                return ctx.MakeConstraint<TSortedConstraintNode>(std::move(content));
-            }
+    static const TSortedConstraintNode* DeduceSortConstraint(const TExprNode& directions, const TExprNode& keyExtractor, TExprContext& ctx) {
+        if (const auto& columns = ExtractSimpleSortTraits(directions, keyExtractor); !columns.empty()) {
+            TSortedConstraintNode::TContainerType content(columns.size());
+            std::transform(columns.cbegin(), columns.cend(), content.begin(), [](const std::pair<TConstraintNode::TPathType, bool>& item) {
+                return std::make_pair(TSortedConstraintNode::TColumnsSet{item.first}, item.second);
+            });
+            return ctx.MakeConstraint<TSortedConstraintNode>(std::move(content));
         }
         return nullptr;
     }
@@ -2542,11 +2589,13 @@ private:
         TSortedConstraintNode::TContainerType filtered;
         for (auto i = 0U; i < content.size(); ++i) {
             TSortedConstraintNode::TContainerType::value_type nextItem;
-            for (const auto& column : content[i].first) {
-                auto range = reverseMapping.equal_range(column);
-                if (range.first != range.second) {
-                    for (auto it = range.first; it != range.second; ++it) {
-                        nextItem.first.emplace_back(it->second);
+            for (const auto& path : content[i].first) {
+                if (path.size() == 1U) {
+                    auto range = reverseMapping.equal_range(path.front());
+                    if (range.first != range.second) {
+                        for (auto it = range.first; it != range.second; ++it) {
+                            nextItem.first.insert_unique(TConstraintNode::TPathType{it->second});
+                        }
                     }
                 }
             }
diff --git a/ydb/library/yql/core/yql_opt_utils.cpp b/ydb/library/yql/core/yql_opt_utils.cpp
index 12f2193c8b0..c24cd01ed82 100644
--- a/ydb/library/yql/core/yql_opt_utils.cpp
+++ b/ydb/library/yql/core/yql_opt_utils.cpp
@@ -1003,53 +1003,6 @@ void ExtractSimpleKeys(const TExprNode* keySelectorBody, const TExprNode* keySel
     }
 }
 
-std::vector<std::pair<std::string_view, bool>> ExtractSimpleSortTraits(const TExprNode& sortDirections, const TExprNode& keySelectorLambda) {
-    const auto& keySelectorBody = keySelectorLambda.Tail();
-    const auto& keySelectorArg = keySelectorLambda.Head().Head();
-    std::vector<std::pair<std::string_view, bool>> columns;
-    if (sortDirections.IsCallable("Bool"))
-        columns.emplace_back(std::string_view(), IsTrue(sortDirections.Tail().Content()));
-    else if (sortDirections.IsList())
-        if (const auto size = keySelectorBody.ChildrenSize()) {
-            columns.reserve(size);
-            for (auto i = 0U; i < size; ++i)
-                if (const auto child = sortDirections.Child(i); child->IsCallable("Bool"))
-                    columns.emplace_back(std::string_view(), IsTrue(child->Tail().Content()));
-                else
-                    return {};
-        } else
-            return {};
-    else
-        return  {};
-
-    if (keySelectorBody.IsCallable("Member") && &keySelectorBody.Head() == &keySelectorArg)
-        if (columns.size() == 1U)
-            columns.front().first = keySelectorBody.Tail().Content();
-        else
-            return {};
-    else if (keySelectorBody.IsList())
-        if (const auto size = keySelectorBody.ChildrenSize()) {
-            std::unordered_set<std::string_view> set(size);
-            columns.resize(size, std::make_pair(std::string_view(), columns.back().second));
-            auto it = columns.begin();
-            for (auto i = 0U; i < columns.size(); ++i) {
-                if (const auto child = keySelectorBody.Child(i); child->IsCallable("Member") && &child->Head() == &keySelectorArg) {
-                    if (set.emplace(child->Tail().Content()).second)
-                        it++->first = child->Tail().Content();
-                    else if (columns.cend() != it)
-                        it = columns.erase(it);
-                } else {
-                    columns.resize(i);
-                }
-            }
-        } else
-            return {};
-    else
-        return {};
-
-    return columns;
-}
-
 const TExprNode& SkipCallables(const TExprNode& node, const std::initializer_list<std::string_view>& skipCallables) {
     const TExprNode* p = &node;
     while (p->IsCallable(skipCallables)) {
diff --git a/ydb/library/yql/core/yql_opt_utils.h b/ydb/library/yql/core/yql_opt_utils.h
index 049f458fb74..e199ba26788 100644
--- a/ydb/library/yql/core/yql_opt_utils.h
+++ b/ydb/library/yql/core/yql_opt_utils.h
@@ -73,7 +73,6 @@ void ExtractSimpleKeys(const TExprNode* keySelectorBody, const TExprNode* keySel
 inline void ExtractSimpleKeys(const TExprNode& keySelectorLambda, TVector<TStringBuf>& columns) {
     ExtractSimpleKeys(keySelectorLambda.Child(1), keySelectorLambda.Child(0)->Child(0), columns);
 }
-std::vector<std::pair<std::string_view, bool>> ExtractSimpleSortTraits(const TExprNode& sortDirections, const TExprNode& keySelectorLambda);
 
 TExprNode::TPtr MakeNull(TPositionHandle position, TExprContext& ctx);
 TExprNode::TPtr MakeConstMap(TPositionHandle position, const TExprNode::TPtr& input, const TExprNode::TPtr& value, TExprContext& ctx);
diff --git a/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp b/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp
index c0f12545863..fb6ecfa7459 100644
--- a/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp
+++ b/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp
@@ -2608,7 +2608,12 @@ TMkqlCommonCallableCompiler::TShared::TShared() {
                 std::vector<TRuntimeNode> keys;
                 keys.reserve(content.size());
                 for (const auto& c : content) {
-                    keys.push_back(ctx.ProgramBuilder.Member(item, c.first.front()));
+                    if (c.first.front().empty())
+                         keys.push_back(item);
+                    else {
+                        MKQL_ENSURE(c.first.front().size() == 1U, "Just column expected.");
+                        keys.push_back(ctx.ProgramBuilder.Member(item, c.first.front().front()));
+                    }
                 }
                 return ctx.ProgramBuilder.NewTuple(keys);
             };