AsumeUnique support many sets of columns.

5 files changed, 60 insertions, 47 deletions

diff --git a/ydb/library/yql/core/common_opt/yql_co_extr_members.cpp b/ydb/library/yql/core/common_opt/yql_co_extr_members.cpp
index afbbf1a586b..e3142abfe32 100644
--- a/ydb/library/yql/core/common_opt/yql_co_extr_members.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_extr_members.cpp
@@ -237,31 +237,24 @@ TExprNode::TPtr ApplyExtractMembersToSort(const TExprNode::TPtr& node, const TEx
 }
 
 TExprNode::TPtr ApplyExtractMembersToAssumeUnique(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix) {
-    TCoAssumeUnique assumeUnique(node);
-    TSet<TStringBuf> extractFields;
+    std::unordered_set<std::string_view> extractFields(members->ChildrenSize());
     for (const auto& x : members->ChildrenList()) {
         extractFields.emplace(x->Content());
     }
-    const bool allExist = AllOf(assumeUnique.UniqueBy(), [&extractFields] (const TCoAtom& u) { return extractFields.contains(u.Value()); });
-    if (allExist) {
+
+    auto children = node->ChildrenList();
+    children.front() = ctx.NewCallable(node->Pos(), TCoExtractMembers::CallableName(), {std::move(children.front()), members});
+    auto it = children.begin();
+    it = std::remove_if(++it, children.end(), [&extractFields](const TExprNode::TPtr& list) { return AnyOf(list->Children(), [&extractFields] (const TExprNode::TPtr& u) { return !extractFields.contains(u->Content()); }); });
+    children.erase(it, children.cend());
+
+    if (children.size() > 1U) {
         YQL_CLOG(DEBUG, Core) << "Move ExtractMembers over " << node->Content() << logSuffix;
-        return ctx.Builder(assumeUnique.Pos())
-            .Callable(node->Content())
-                .Callable(0, TCoExtractMembers::CallableName())
-                    .Add(0, assumeUnique.Input().Ptr())
-                    .Add(1, members)
-                .Seal()
-                .Add(1, assumeUnique.UniqueBy().Ptr())
-            .Seal()
-            .Build();
+        return ctx.ChangeChildren(*node, std::move(children));
+    } else {
+        YQL_CLOG(DEBUG, Core) << "Drop " << node->Content() << " after ExtractMembers" << logSuffix;
+        return std::move(children.front());
     }
-    YQL_CLOG(DEBUG, Core) << "Drop " << node->Content() << " after ExtractMembers" << logSuffix;
-    return ctx.Builder(assumeUnique.Pos())
-        .Callable(TCoExtractMembers::CallableName())
-            .Add(0, assumeUnique.Input().Ptr())
-            .Add(1, members)
-        .Seal()
-        .Build();
 }
 
 TExprNode::TPtr ApplyExtractMembersToTop(const TExprNode::TPtr& node, const TExprNode::TPtr& members, const TParentsMap& parentsMap, TExprContext& ctx, TStringBuf logSuffix) {
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 d2bbe59db5e..b5de86a3a8e 100644
--- a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
@@ -140,8 +140,9 @@ TExprNode::TPtr KeepConstraints(TExprNode::TPtr node, const TExprNode& src, TExp
             for (const auto& path : set)
                 if (!path.empty())
                     columns.emplace_back(ctx.NewAtom(node->Pos(), path.front()));
-        if (!columns.empty()) { // TODO: AssumeUnique must support any set of columns.
-            res = ctx.Builder(node->Pos())
+        res = columns.empty() ?
+            ctx.NewCallable(node->Pos(), "AssumeUnique", {std::move(res)}):
+            ctx.Builder(node->Pos())
                 .Callable("AssumeUnique")
                     .Add(0, std::move(res))
                     .List(1)
@@ -149,7 +150,6 @@ TExprNode::TPtr KeepConstraints(TExprNode::TPtr node, const TExprNode& src, TExp
                     .Seal()
                 .Seal()
                 .Build();
-        }
     }
     return res;
 }
diff --git a/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json b/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json
index 0e83c9c9c98..696777c89e5 100644
--- a/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json
+++ b/ydb/library/yql/core/expr_nodes/yql_expr_nodes.json
@@ -1632,10 +1632,10 @@
         },
         {
             "Name": "TCoAssumeUnique",
-            "Base": "TCoInputBase",
+            "Base": "TFreeArgCallable",
             "Match": {"Type": "Callable", "Name": "AssumeUnique"},
             "Children": [
-                {"Index": 1, "Name": "UniqueBy", "Type": "TCoAtomList"}
+                {"Index": 0, "Name": "Input", "Type": "TExprBase"}
             ]
         },
         {
diff --git a/ydb/library/yql/core/type_ann/type_ann_list.cpp b/ydb/library/yql/core/type_ann/type_ann_list.cpp
index 84fd259f675..040dbb9e048 100644
--- a/ydb/library/yql/core/type_ann/type_ann_list.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_list.cpp
@@ -4241,31 +4241,33 @@ namespace {
     }
 
     IGraphTransformer::TStatus AssumeUniqueWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
-        if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
+        if (!EnsureMinArgsCount(*input, 1, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
         }
 
-        const auto status = NormalizeTupleOfAtoms(input, 1, output, ctx.Expr);
-        if (status != IGraphTransformer::TStatus::Ok) {
-            return status;
-        }
-
         const TTypeAnnotationNode* inputItemType = nullptr;
         if (!EnsureNewSeqType<false>(input->Head(), ctx.Expr, &inputItemType)) {
             return IGraphTransformer::TStatus::Error;
         }
 
-        if (!EnsureStructType(input->Head().Pos(), *inputItemType, ctx.Expr)) {
-            return IGraphTransformer::TStatus::Error;
-        }
-
-        const auto structType = inputItemType->Cast<TStructExprType>();
-        for (auto& x : input->Tail().Children()) {
-            YQL_ENSURE(x->IsAtom());
-            auto pos = FindOrReportMissingMember(x->Content(), input->Head().Pos(), *structType, ctx.Expr);
-            if (!pos) {
+        if (input->ChildrenSize() > 1U) {
+            if (!EnsureStructType(input->Head().Pos(), *inputItemType, ctx.Expr)) {
                 return IGraphTransformer::TStatus::Error;
             }
+
+            for (auto i = 1U; i < input->ChildrenSize(); ++i) {
+                if (const auto status = NormalizeTupleOfAtoms(input, i, output, ctx.Expr); status != IGraphTransformer::TStatus::Ok) {
+                    return status;
+                }
+
+                const auto structType = inputItemType->Cast<TStructExprType>();
+                for (const auto& x : input->Child(i)->Children()) {
+                    YQL_ENSURE(x->IsAtom());
+                    if (!FindOrReportMissingMember(x->Content(), x->Pos(), *structType, ctx.Expr)) {
+                        return IGraphTransformer::TStatus::Error;
+                    }
+                }
+            }
         }
 
         input->SetTypeAnn(input->Head().GetTypeAnn());
diff --git a/ydb/library/yql/core/yql_expr_constraint.cpp b/ydb/library/yql/core/yql_expr_constraint.cpp
index c7c0ced7fc2..31b6fc84ef4 100644
--- a/ydb/library/yql/core/yql_expr_constraint.cpp
+++ b/ydb/library/yql/core/yql_expr_constraint.cpp
@@ -363,12 +363,19 @@ private:
     }
 
     TStatus AssumeUniqueWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
-        std::vector<std::string_view> columns;
-        columns.reserve(input->Child(1)->ChildrenSize());
-        for (const auto& column: input->Child(1)->Children()) {
-            columns.emplace_back(column->Content());
+        TUniqueConstraintNode::TFullSetType sets;
+        for (auto i = 1U; i < input->ChildrenSize(); ++i) {
+            TUniqueConstraintNode::TSetType columns;
+            columns.reserve(input->Child(i)->ChildrenSize());
+            for (const auto& column: input->Child(i)->Children())
+                columns.insert_unique(TConstraintNode::TPathType(1U, column->Content()));
+            sets.insert_unique(std::move(columns));
         }
-        input->AddConstraint(ctx.MakeConstraint<TUniqueConstraintNode>(columns));
+
+        if (sets.empty())
+            sets.insert_unique(TUniqueConstraintNode::TSetType{TConstraintNode::TPathType()});
+
+        input->AddConstraint(ctx.MakeConstraint<TUniqueConstraintNode>(std::move(sets)));
         return FromFirst<TPassthroughConstraintNode, TSortedConstraintNode, TEmptyConstraintNode, TVarIndexConstraintNode>(input, output, ctx);
     }
 
@@ -777,7 +784,7 @@ private:
                         break;
                 }
 
-                if (auto mapping = TPartOfUniqueConstraintNode::GetCommonMapping(node.Head().GetConstraint<TUniqueConstraintNode>(), node.Head().GetConstraint<TPartOfUniqueConstraintNode>()); !mapping.empty()) {
+                if (auto mapping = TPartOfUniqueConstraintNode::GetCommonMapping(GetDetailedUinique(node.Head().GetConstraint<TUniqueConstraintNode>(), *node.Head().GetTypeAnn(), ctx), node.Head().GetConstraint<TPartOfUniqueConstraintNode>()); !mapping.empty()) {
                     constraints.emplace_back(ctx.MakeConstraint<TPartOfUniqueConstraintNode>(std::move(mapping)));
                 }
                 if (const auto groupBy = node.Head().GetConstraint<TGroupByConstraintNode>()) {
@@ -846,7 +853,7 @@ private:
         }
 
         if (const auto lambdaUnique = GetConstraintFromLambda<TPartOfUniqueConstraintNode, WideOutput>(input->Tail(), ctx)) {
-            if (const auto unique = input->Head().GetConstraint<TUniqueConstraintNode>()) {
+            if (const auto unique = GetDetailedUinique(input->Head().GetConstraint<TUniqueConstraintNode>(), *input->Head().GetTypeAnn(), ctx)) {
                 if (const auto complete = TPartOfUniqueConstraintNode::MakeComplete(ctx, lambdaUnique->GetColumnMapping(), unique)) {
                     input->AddConstraint(complete);
                 }
@@ -2462,6 +2469,17 @@ private:
         return fields;
     }
 
+    static const TUniqueConstraintNode* GetDetailedUinique(const TUniqueConstraintNode* unique,  const TTypeAnnotationNode& type, TExprContext& ctx) {
+        if (!unique)
+            return nullptr;
+
+        if (const auto& sets = unique->GetAllSets(); sets.size() != 1U || sets.cbegin()->size() != 1U || !sets.cbegin()->cbegin()->empty())
+            return unique;
+
+        const auto& columns = GetAllItemTypeFields(type, ctx);
+        return columns.empty() ? nullptr : ctx.MakeConstraint<TUniqueConstraintNode>(columns);
+    }
+
     static const TStructExprType* GetNonEmptyStructItemType(const TTypeAnnotationNode& type) {
         const TTypeAnnotationNode* itemType = GetItemType(type);
         if (!itemType || itemType->GetKind() != ETypeAnnotationKind::Struct) {