YQL-8971 YQL-15555 Constraints for FlattenMembers.

3 files changed, 96 insertions, 19 deletions

diff --git a/ydb/library/yql/ast/yql_constraint.cpp b/ydb/library/yql/ast/yql_constraint.cpp
index fc556132c5d..4ffe62e4b6b 100644
--- a/ydb/library/yql/ast/yql_constraint.cpp
+++ b/ydb/library/yql/ast/yql_constraint.cpp
@@ -994,6 +994,24 @@ TPartOfConstraintNode<TOriginalConstraintNode>::GetColumnMapping(const std::stri
 }
 
 template<class TOriginalConstraintNode>
+typename TPartOfConstraintNode<TOriginalConstraintNode>::TMapType
+TPartOfConstraintNode<TOriginalConstraintNode>::GetColumnMapping(TExprContext& ctx, const std::string_view& prefix) const {
+    auto mapping = Mapping_;
+    if (!prefix.empty()) {
+        const TString str(prefix);
+        for (auto& item : mapping) {
+            for (auto& part : item.second) {
+                if (part.first.empty())
+                    part.first.emplace_front(prefix);
+                else
+                    part.first.front() = ctx.AppendString(str + part.first.front());
+            }
+        }
+    }
+    return mapping;
+}
+
+template<class TOriginalConstraintNode>
 const TPartOfConstraintNode<TOriginalConstraintNode>*
 TPartOfConstraintNode<TOriginalConstraintNode>::MakeCommon(const std::vector<const TConstraintSet*>& constraints, TExprContext& ctx) {
     if (constraints.empty()) {
@@ -1274,6 +1292,17 @@ void TPassthroughConstraintNode::ToJson(NJson::TJsonWriter& out) const {
     out.CloseMap();
 }
 
+void TPassthroughConstraintNode::UniqueMerge(TMapType& output, TMapType&& input) {
+    output.merge(std::move(input));
+    while (!input.empty()) {
+        const auto exists = input.extract(input.cbegin());
+        auto& target = output[exists.key()];
+        target.reserve(target.size() + exists.mapped().size());
+        for (auto& item : exists.mapped())
+            target.insert_unique(std::move(item));
+    }
+}
+
 const TPassthroughConstraintNode* TPassthroughConstraintNode::ExtractField(TExprContext& ctx, const std::string_view& field) const {
     TMapType passtrought;
     for (const auto& part : Mapping_) {
@@ -1388,7 +1417,7 @@ const TPassthroughConstraintNode::TMapType& TPassthroughConstraintNode::GetColum
     return Mapping_;
 }
 
-TPassthroughConstraintNode::TMapType TPassthroughConstraintNode::GetMappingForField(const std::string_view& field) const {
+TPassthroughConstraintNode::TMapType TPassthroughConstraintNode::GetColumnMapping(const std::string_view& field) const {
     TMapType mapping(Mapping_.size());
     for (const auto& map : Mapping_) {
         TPartType part;
@@ -1402,6 +1431,24 @@ TPassthroughConstraintNode::TMapType TPassthroughConstraintNode::GetMappingForFi
     return mapping;
 }
 
+TPassthroughConstraintNode::TMapType TPassthroughConstraintNode::GetColumnMapping(TExprContext& ctx, const std::string_view& prefix) const {
+    TMapType mapping(Mapping_.size());
+    for (const auto& map : Mapping_) {
+        TPartType part;
+        part.reserve(map.second.size());
+        const TString str(prefix);
+        std::transform(map.second.cbegin(), map.second.cend(), std::back_inserter(part), [&](TPartType::value_type item) {
+            if (item.first.empty())
+                item.first.emplace_front(prefix);
+            else
+                item.first.front() = ctx.AppendString(str + item.first.front());
+            return item;
+        });
+        mapping.emplace(map.first ? map.first : this, std::move(part));
+    }
+    return mapping;
+}
+
 TPassthroughConstraintNode::TReverseMapType TPassthroughConstraintNode::GetReverseMapping() const {
     if (1U == Mapping_.size() && 1U == Mapping_.cbegin()->second.size() && Mapping_.cbegin()->second.cbegin()->first.empty())
         return {{Mapping_.cbegin()->second.cbegin()->second, Mapping_.cbegin()->second.cbegin()->second}};
diff --git a/ydb/library/yql/ast/yql_constraint.h b/ydb/library/yql/ast/yql_constraint.h
index 947b5f6281f..743583b08d8 100644
--- a/ydb/library/yql/ast/yql_constraint.h
+++ b/ydb/library/yql/ast/yql_constraint.h
@@ -299,6 +299,7 @@ public:
 
     const TMapType& GetColumnMapping() const;
     TMapType GetColumnMapping(const std::string_view& asField) const;
+    TMapType GetColumnMapping(TExprContext& ctx, const std::string_view& prefix) const;
 
     bool Equals(const TConstraintNode& node) const override;
     bool Includes(const TConstraintNode& node) const override;
@@ -360,15 +361,17 @@ public:
     }
 
     const TMapType& GetColumnMapping() const;
-    TReverseMapType GetReverseMapping() const;
+    TMapType GetColumnMapping(const std::string_view& asField) const;
+    TMapType GetColumnMapping(TExprContext& ctx, const std::string_view& prefix) const;
 
-    TMapType GetMappingForField(const std::string_view& field) const;
+    TReverseMapType GetReverseMapping() const;
 
     bool Equals(const TConstraintNode& node) const override;
     bool Includes(const TConstraintNode& node) const override;
     void Out(IOutputStream& out) const override;
     void ToJson(NJson::TJsonWriter& out) const override;
 
+    static void UniqueMerge(TMapType& output, TMapType&& input);
     const TPassthroughConstraintNode* ExtractField(TExprContext& ctx, const std::string_view& field) const;
 
     static const TPassthroughConstraintNode* MakeCommon(const std::vector<const TConstraintSet*>& constraints, TExprContext& ctx);
diff --git a/ydb/library/yql/core/yql_expr_constraint.cpp b/ydb/library/yql/core/yql_expr_constraint.cpp
index c29ef767565..d0db56f85de 100644
--- a/ydb/library/yql/core/yql_expr_constraint.cpp
+++ b/ydb/library/yql/core/yql_expr_constraint.cpp
@@ -112,6 +112,7 @@ public:
         Functions["ExtractMembers"] = &TCallableConstraintTransformer::ExtractMembersWrap;
         Functions["RemoveSystemMembers"] = &TCallableConstraintTransformer::RemovePrefixMembersWrap;
         Functions["RemovePrefixMembers"] = &TCallableConstraintTransformer::RemovePrefixMembersWrap;
+        Functions["FlattenMembers"] = &TCallableConstraintTransformer::FlattenMembersWrap;
         Functions["SelectMembers"] = &TCallableConstraintTransformer::SelectMembersWrap;
         Functions["FilterMembers"] = &TCallableConstraintTransformer::SelectMembersWrap;
         Functions["CastStruct"] = &TCallableConstraintTransformer::SelectMembersWrap;
@@ -1366,13 +1367,8 @@ private:
         for (auto i = 0U; i < input->ChildrenSize(); ++i) {
             const auto child = input->Child(i);
             const auto& name = ctx.GetIndexAsString(i);
-            if (const auto pass = child->GetConstraint<TPassthroughConstraintNode>()) {
-                for (auto part : pass->GetColumnMapping()) {
-                    std::transform(part.second.cbegin(), part.second.cend(), std::back_inserter(passthrough[part.first ? part.first : pass]), [&name](TPassthroughConstraintNode::TPartType::value_type item) {
-                        item.first.emplace_front(name);
-                        return item;
-                    });
-                }
+            if (const auto part = child->GetConstraint<TPassthroughConstraintNode>()) {
+                TPassthroughConstraintNode::UniqueMerge(passthrough, part->GetColumnMapping(name));
             }
 
             if (const auto part = child->GetConstraint<TPartOfSortedConstraintNode>()) {
@@ -1421,13 +1417,8 @@ private:
         std::vector<const TConstraintSet*> structConstraints;
         for (const auto& child : input->Children()) {
             const auto& name = child->Head().Content();
-            if (const auto pass = child->Tail().GetConstraint<TPassthroughConstraintNode>()) {
-                for (auto part : pass->GetColumnMapping()) {
-                    std::transform(part.second.cbegin(), part.second.cend(), std::back_inserter(passthrough[part.first ? part.first : pass]), [&name](TPassthroughConstraintNode::TPartType::value_type item) {
-                        item.first.emplace_front(name);
-                        return item;
-                    });
-                }
+            if (const auto part = child->Tail().GetConstraint<TPassthroughConstraintNode>()) {
+                TPassthroughConstraintNode::UniqueMerge(passthrough, part->GetColumnMapping(name));
             }
 
             if (const auto part = child->Tail().GetConstraint<TPartOfSortedConstraintNode>()) {
@@ -1465,6 +1456,42 @@ private:
         return TStatus::Ok;
     }
 
+    TStatus FlattenMembersWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
+        TPassthroughConstraintNode::TMapType passthrough;
+        TPartOfSortedConstraintNode::TMapType sorted;
+        TPartOfUniqueConstraintNode::TMapType uniques;
+        TPartOfDistinctConstraintNode::TMapType distincts;
+
+        for (const auto& child : input->Children()) {
+            const auto& prefix = child->Head().Content();
+            if (const auto part = child->Tail().GetConstraint<TPassthroughConstraintNode>()) {
+                TPassthroughConstraintNode::UniqueMerge(passthrough, part->GetColumnMapping(ctx, prefix));
+            }
+            if (const auto part = child->Tail().GetConstraint<TPartOfSortedConstraintNode>()) {
+                TPartOfSortedConstraintNode::UniqueMerge(sorted, part->GetColumnMapping(ctx, prefix));
+            }
+            if (const auto part = child->Tail().GetConstraint<TPartOfUniqueConstraintNode>()) {
+                TPartOfUniqueConstraintNode::UniqueMerge(uniques, part->GetColumnMapping(ctx, prefix));
+            }
+            if (const auto part = child->Tail().GetConstraint<TPartOfDistinctConstraintNode>()) {
+                TPartOfDistinctConstraintNode::UniqueMerge(distincts, part->GetColumnMapping(ctx, prefix));
+            }
+        }
+        if (!passthrough.empty()) {
+            input->AddConstraint(ctx.MakeConstraint<TPassthroughConstraintNode>(std::move(passthrough)));
+        }
+        if (!sorted.empty()) {
+            input->AddConstraint(ctx.MakeConstraint<TPartOfSortedConstraintNode>(std::move(sorted)));
+        }
+        if (!uniques.empty()) {
+            input->AddConstraint(ctx.MakeConstraint<TPartOfUniqueConstraintNode>(std::move(uniques)));
+        }
+        if (!distincts.empty()) {
+            input->AddConstraint(ctx.MakeConstraint<TPartOfDistinctConstraintNode>(std::move(distincts)));
+        }
+        return TStatus::Ok;
+    }
+
     template<class TPartOfConstraint>
     static void AddPartOf(const TExprNode::TPtr& input, TExprContext& ctx) {
         typename TPartOfConstraint::TMapType map;
@@ -2636,9 +2663,9 @@ private:
 
         TPassthroughConstraintNode::TMapType passthrough;
         if (const auto keysPassthrough = GetConstraintFromLambda<TPassthroughConstraintNode, false>(*input->Child(1), ctx))
-            passthrough.merge(keysPassthrough->GetMappingForField(ctx.GetIndexAsString(0U)));
+            passthrough.merge(keysPassthrough->GetColumnMapping(ctx.GetIndexAsString(0U)));
         if (const auto paysPassthrough = GetConstraintFromLambda<TPassthroughConstraintNode, false>(*input->Child(2), ctx))
-            passthrough.merge(paysPassthrough->GetMappingForField(ctx.GetIndexAsString(1U)));
+            passthrough.merge(paysPassthrough->GetColumnMapping(ctx.GetIndexAsString(1U)));
 
         if (!passthrough.empty())
             input->AddConstraint(ctx.MakeConstraint<TPassthroughConstraintNode>(std::move(passthrough)));