Correct support many sets of unique columns.

3 files changed, 82 insertions, 34 deletions

diff --git a/ydb/library/yql/ast/yql_constraint.cpp b/ydb/library/yql/ast/yql_constraint.cpp
index ce474125331..74715b02baf 100644
--- a/ydb/library/yql/ast/yql_constraint.cpp
+++ b/ydb/library/yql/ast/yql_constraint.cpp
@@ -404,12 +404,14 @@ TUniqueConstraintNode::TSetType ColumnsListToSet(const std::vector<std::string_v
 }
 
 TUniqueConstraintNode::TFullSetType DedupSets(TUniqueConstraintNode::TFullSetType&& sets) {
-    if (sets.size() > 1U) {
-        for (const auto& set : sets) {
+    for (bool found = true; found && sets.size() > 1U;) {
+        found = false;
+        for (auto ot = sets.cbegin(); !found && sets.cend() != ot; ++ot) {
             for (auto it = sets.cbegin(); sets.cend() != it;) {
-                if (set.size() < it->size() && std::all_of(set.cbegin(), set.cend(), [it](const TConstraintNode::TPathType& path) { return it->cend() == it->find(path); }))
+                if (ot->size() < it->size() && std::all_of(ot->cbegin(), ot->cend(), [it](const TConstraintNode::TPathType& path) { return it->cend() != it->find(path); })) {
                     it = sets.erase(it);
-                else
+                    found = true;
+                } else
                     ++it;
             }
         }
@@ -532,13 +534,23 @@ const TUniqueConstraintNode* TUniqueConstraintNode::MakeCommon(const std::vector
 }
 
 bool TUniqueConstraintNode::HasEqualColumns(const std::vector<std::string_view>& columns) const {
-    // TODO: from combination of all set instead of any set.
-    return !columns.empty() && std::any_of(Sets_.cbegin(), Sets_.cend(), [&columns](const TSetType& set) {
-        return columns.size() == set.size() && std::all_of(columns.cbegin(), columns.cend(), [&set](const std::string_view& col) {
-            const auto it = set.lower_bound(TPathType(1U, col));
-            return set.cend() != it && !it->empty() && it->front() == col;
-        });
-    });
+    if (columns.empty())
+        return false;
+
+    const std::unordered_set<std::string_view> ordered(columns.cbegin(), columns.cend());
+    std::unordered_set<std::string_view> uniques(columns.size());
+    for (const auto& set : Sets_) {
+        if (std::all_of(set.cbegin(), set.cend(), [&ordered](const TPathType& path) { return !path.empty() && ordered.contains(path.front()); })) {
+            for (const auto& path : set) {
+                if (!path.empty()) {
+                    uniques.emplace(path.front());
+                }
+            }
+            if (uniques.size() == ordered.size())
+                return true;
+        }
+    }
+    return false;
 }
 
 const TUniqueConstraintNode* TUniqueConstraintNode::FilterFields(TExprContext& ctx, const std::function<bool(const TPathType& front)>& predicate) const {
@@ -834,20 +846,28 @@ TPartOfUniqueConstraintNode::ExtractField(const TMapType& mapping, const std::st
 const TUniqueConstraintNode* TPartOfUniqueConstraintNode::MakeComplete(TExprContext& ctx, const TMapType& mapping, const TUniqueConstraintNode* original) {
     if (const auto it = mapping.find(original); mapping.cend() != it) {
         TUniqueConstraintNode::TFullSetType newSets;
-
-        auto sets = it->first->GetAllSets();
-        for (auto s = sets.begin(); sets.end() != s;) {
-            TUniqueConstraintNode::TSetType newSet;
-            for (const auto& map : it->second) {
-                newSet.insert_unique(map.first);
-                s->erase(map.second);
+        for (const auto& set : it->first->GetAllSets()) {
+            TReversePartType reverseMap;
+            reverseMap.reserve(it->second.size());
+            for (const auto& map : it->second)
+                reverseMap[map.second].insert_unique(map.first);
+
+            if (std::all_of(set.cbegin(), set.cend(), [reverseMap] (const TPathType& path) { return reverseMap.cend() != reverseMap.find(path); })) {
+                TUniqueConstraintNode::TFullSetType addSets = {TUniqueConstraintNode::TSetType()};
+                for (const auto& old : set) {
+                    const auto& renamed = reverseMap[old];
+                    TUniqueConstraintNode::TFullSetType incSets;
+                    for (auto s = addSets.begin(); addSets.end() != s; s = addSets.erase(s)) {
+                        for (const auto& item : renamed) {
+                            auto newSet = std::move(*s);
+                            newSet.insert_unique(item);
+                            incSets.insert_unique(std::move(newSet));
+                        }
+                    }
+                    incSets.swap(addSets);
+                }
+                newSets.insert(addSets.cbegin(), addSets.cend());
             }
-
-            if (s->empty()) {
-                newSets.insert_unique(std::move(newSet));
-                s = sets.erase(s);
-            } else
-                ++s;
         }
 
         if (!newSets.empty())
diff --git a/ydb/library/yql/ast/yql_constraint.h b/ydb/library/yql/ast/yql_constraint.h
index 873c0fa9e14..ec5928d6f1e 100644
--- a/ydb/library/yql/ast/yql_constraint.h
+++ b/ydb/library/yql/ast/yql_constraint.h
@@ -258,6 +258,7 @@ private:
 class TPartOfUniqueConstraintNode final: public TConstraintNode {
 public:
     using TPartType = NSorted::TSimpleMap<TPathType, TPathType>;
+    using TReversePartType = NSorted::TSimpleMap<TPathType, NSorted::TSimpleSet<TPathType>>;
     using TMapType = std::unordered_map<const TUniqueConstraintNode*, TPartType>;
 private:
     friend struct TExprContext;
diff --git a/ydb/library/yql/core/yql_expr_constraint.cpp b/ydb/library/yql/core/yql_expr_constraint.cpp
index 06525cd91b7..425c53a017a 100644
--- a/ydb/library/yql/core/yql_expr_constraint.cpp
+++ b/ydb/library/yql/core/yql_expr_constraint.cpp
@@ -1968,12 +1968,12 @@ private:
         }
 
         if (const auto switchLambda = input->Child(2); switchLambda->Tail().IsCallable(TCoBool::CallableName()) && IsFalse(switchLambda->Tail().Head().Content())) {
-            if (const auto outItemType = GetNonEmptyStructItemType(*input->GetTypeAnn())) {
-                std::vector<std::string_view> columns;
-                columns.reserve(outItemType->GetSize());
-                for (const auto& item: outItemType->GetItems())
-                    columns.emplace_back(item->GetName());
-                input->AddConstraint(ctx.MakeConstraint<TUniqueConstraintNode>(columns));
+            if (const auto& fields = GetAllItemTypeFields(*input->GetTypeAnn(), ctx); !fields.empty()) {
+                TUniqueConstraintNode::TFullSetType sets;
+                sets.reserve(fields.size());
+                for (const auto& field: fields)
+                    sets.insert_unique(TUniqueConstraintNode::TSetType{TConstraintNode::TPathType(1U, field)});
+                input->AddConstraint(ctx.MakeConstraint<TUniqueConstraintNode>(std::move(sets)));
             }
         }
         else {
@@ -2056,10 +2056,13 @@ private:
         }
 
         if (const auto switchLambda = input->Child(2); switchLambda->Tail().IsCallable(TCoBool::CallableName()) && IsFalse(switchLambda->Tail().Head().Content())) {
-            std::vector<std::string_view> fields(initLambda->Head().ChildrenSize());
-            ui32 i = 0U;
-            std::generate_n(fields.begin(), initLambda->Head().ChildrenSize(), [&i, &ctx](){ return ctx.GetIndexAsString(i++); });
-            input->AddConstraint(ctx.MakeConstraint<TUniqueConstraintNode>(fields));
+            if (const auto width = initLambda->Head().ChildrenSize()) {
+                TUniqueConstraintNode::TFullSetType sets;
+                sets.reserve(width);
+                for (ui32 i = 0U; i < width; ++i)
+                    sets.insert_unique(TUniqueConstraintNode::TSetType{TConstraintNode::TPathType(1U, ctx.GetIndexAsString(i))});
+                input->AddConstraint(ctx.MakeConstraint<TUniqueConstraintNode>(std::move(sets)));
+            }
         } else {
             TVector<TStringBuf> groupByKeys;
             if (const auto groupBy = switchLambda->GetConstraint<TGroupByConstraintNode>()) {
@@ -2435,6 +2438,30 @@ private:
         ExtractSimpleKeys(body, arg, columns);
     }
 
+    static std::vector<std::string_view> GetAllItemTypeFields(const TTypeAnnotationNode& type, TExprContext& ctx) {
+        std::vector<std::string_view> fields;
+        if (const auto itemType = GetItemType(type)) {
+            switch (itemType->GetKind()) {
+                case ETypeAnnotationKind::Struct:
+                    if (const auto structType = itemType->Cast<TStructExprType>()) {
+                        fields.reserve(structType->GetSize());
+                        std::transform(structType->GetItems().cbegin(), structType->GetItems().cend(), std::back_inserter(fields), std::bind(&TItemExprType::GetName, std::placeholders::_1));
+                    }
+                    break;
+                case ETypeAnnotationKind::Tuple:
+                    if (const auto size = itemType->Cast<TTupleExprType>()->GetSize()) {
+                        fields.resize(size);
+                        ui32 i = 0U;
+                        std::generate(fields.begin(), fields.end(), [&]() { return ctx.GetIndexAsString(i++); });
+                    }
+                    break;
+                default:
+                    break;
+            }
+        }
+        return fields;
+    }
+
     static const TStructExprType* GetNonEmptyStructItemType(const TTypeAnnotationNode& type) {
         const TTypeAnnotationNode* itemType = GetItemType(type);
         if (!itemType || itemType->GetKind() != ETypeAnnotationKind::Struct) {