YQL-15756 YQL-15555 Fix [Flat]OptionalIf constraints.

3 files changed, 94 insertions, 14 deletions

diff --git a/ydb/library/yql/ast/yql_constraint.cpp b/ydb/library/yql/ast/yql_constraint.cpp
index 81a45f9f0b9..20a3c1ac993 100644
--- a/ydb/library/yql/ast/yql_constraint.cpp
+++ b/ydb/library/yql/ast/yql_constraint.cpp
@@ -290,13 +290,33 @@ bool TSortedConstraintNode::IsPrefixOf(const TSortedConstraintNode& node) const
     return node.Includes(*this);
 }
 
-const TSortedConstraintNode::TFullSetType TSortedConstraintNode::GetAllSets() const {
+TSortedConstraintNode::TFullSetType TSortedConstraintNode::GetAllSets() const {
     TFullSetType sets;
     for (const auto& key : Content_)
         sets.insert_unique(key.first);
     return sets;
 }
 
+void TSortedConstraintNode::FilterUncompleteReferences(TSetType& references) const {
+    TSetType complete;
+    complete.reserve(references.size());
+
+    for (const auto& item : Content_) {
+        bool found = false;
+        for (const auto& path : item.first) {
+            if (references.contains(path)) {
+                found = true;
+                complete.insert_unique(path);
+            }
+        }
+
+        if (!found)
+            break;
+    }
+
+    references = std::move(complete);
+}
+
 const TSortedConstraintNode* TSortedConstraintNode::MakeCommon(const std::vector<const TConstraintSet*>& constraints, TExprContext& ctx) {
     if (constraints.empty()) {
         return nullptr;
@@ -675,6 +695,14 @@ bool TUniqueConstraintNodeBase<Distinct>::HasEqualColumns(const std::vector<std:
 }
 
 template<bool Distinct>
+void TUniqueConstraintNodeBase<Distinct>::FilterUncompleteReferences(TSetType& references) const {
+    for (const auto& set : Sets_) {
+        if (!std::all_of(set.cbegin(), set.cend(), std::bind(&TSetType::contains<TPathType>, std::cref(references), std::placeholders::_1)))
+            std::for_each(set.cbegin(), set.cend(), [&references] (const TPathType& path) { references.erase(path); });
+    }
+}
+
+template<bool Distinct>
 const TUniqueConstraintNodeBase<Distinct>*
 TUniqueConstraintNodeBase<Distinct>::FilterFields(TExprContext& ctx, const TPathFilter& predicate) const {
     auto sets = Sets_;
@@ -912,6 +940,34 @@ TPartOfConstraintNode<TOriginalConstraintNode>::RenameFields(TExprContext& ctx,
 }
 
 template<class TOriginalConstraintNode>
+const TPartOfConstraintNode<TOriginalConstraintNode>*
+TPartOfConstraintNode<TOriginalConstraintNode>::CompleteOnly(TExprContext& ctx) const {
+    TMapType mapping(Mapping_);
+
+    for (auto it = mapping.begin(); mapping.end() != it;) {
+        TSetType set;
+        set.reserve(it->second.size());
+        std::for_each(it->second.cbegin(), it->second.cend(), [&](const TPartType::value_type& pair) { set.insert_unique(pair.second); });
+
+        it->first->FilterUncompleteReferences(set);
+
+        for (auto jt = it->second.cbegin(); it->second.cend() != jt;) {
+            if (set.contains(jt->second))
+                ++jt;
+            else
+                jt = it->second.erase(jt);
+        }
+
+        if (it->second.empty())
+            it = mapping.erase(it);
+        else
+            ++it;
+    }
+
+    return mapping.empty() ? nullptr : ctx.MakeConstraint<TPartOfConstraintNode>(std::move(mapping));
+}
+
+template<class TOriginalConstraintNode>
 typename TPartOfConstraintNode<TOriginalConstraintNode>::TMapType
 TPartOfConstraintNode<TOriginalConstraintNode>::GetColumnMapping(const std::string_view& asField) const {
     auto mapping = Mapping_;
diff --git a/ydb/library/yql/ast/yql_constraint.h b/ydb/library/yql/ast/yql_constraint.h
index 6681ee18246..db9b4160f00 100644
--- a/ydb/library/yql/ast/yql_constraint.h
+++ b/ydb/library/yql/ast/yql_constraint.h
@@ -24,7 +24,6 @@ class TConstraintNode {
 protected:
     TConstraintNode(TExprContext& ctx, std::string_view name);
     TConstraintNode(TConstraintNode&& constr);
-
 public:
     using TPathType = std::deque<std::string_view>;
     using TSetType = NSorted::TSimpleSet<TPathType>;
@@ -193,7 +192,7 @@ public:
         return Content_;
     }
 
-    const TFullSetType GetAllSets() const;
+    TFullSetType GetAllSets() const;
 
     bool Equals(const TConstraintNode& node) const override;
     bool Includes(const TConstraintNode& node) const override;
@@ -204,6 +203,8 @@ public:
 
     const TSortedConstraintNode* CutPrefix(size_t newPrefixLength, TExprContext& ctx) const;
 
+    void FilterUncompleteReferences(TSetType& references) const;
+
     static const TSortedConstraintNode* MakeCommon(const std::vector<const TConstraintSet*>& constraints, TExprContext& ctx);
     const TSortedConstraintNode* MakeCommon(const TSortedConstraintNode* other, TExprContext& ctx) const;
 
@@ -241,6 +242,8 @@ public:
     bool IsOrderBy(const TSortedConstraintNode& sorted) const;
     bool HasEqualColumns(const std::vector<std::string_view>& columns) const;
 
+    void FilterUncompleteReferences(TSetType& references) const;
+
     static const TUniqueConstraintNodeBase* MakeCommon(const std::vector<const TConstraintSet*>& constraints, TExprContext& ctx);
     const TUniqueConstraintNodeBase* FilterFields(TExprContext& ctx, const TPathFilter& predicate) const;
     const TUniqueConstraintNodeBase* RenameFields(TExprContext& ctx, const TPathReduce& reduce) const;
@@ -304,6 +307,7 @@ public:
     const TPartOfConstraintNode* ExtractField(TExprContext& ctx, const std::string_view& field) const;
     const TPartOfConstraintNode* FilterFields(TExprContext& ctx, const TPathFilter& predicate) const;
     const TPartOfConstraintNode* RenameFields(TExprContext& ctx, const TPathReduce& reduce) const;
+    const TPartOfConstraintNode* CompleteOnly(TExprContext& ctx) const;
 
     static const TPartOfConstraintNode* MakeCommon(const std::vector<const TConstraintSet*>& constraints, TExprContext& ctx);
 
@@ -426,13 +430,8 @@ private:
     TMapType Mapping_;
 };
 
-class TMultiConstraintNode: public TConstraintNode {
+class TMultiConstraintNode final: public TConstraintNode {
 public:
-    struct TConstraintKey {
-        TStringBuf operator()(const TConstraintNode* node) const {
-            return node->GetName();
-        }
-    };
     using TMapType = NSorted::TSimpleMap<ui32, TConstraintSet>;
 
 public:
@@ -464,7 +463,7 @@ public:
 
     bool FilteredIncludes(const TConstraintNode& node, const THashSet<TString>& blacklist) const;
     const TConstraintNode* OnlySimpleColumns(TExprContext& ctx) const override;
-protected:
+private:
     TMapType Items_;
 };
 
diff --git a/ydb/library/yql/core/yql_expr_constraint.cpp b/ydb/library/yql/core/yql_expr_constraint.cpp
index e453826c598..f66b4a6d7c3 100644
--- a/ydb/library/yql/core/yql_expr_constraint.cpp
+++ b/ydb/library/yql/core/yql_expr_constraint.cpp
@@ -162,10 +162,10 @@ public:
         Functions["ForceRemoveMember"] = &TCallableConstraintTransformer::RemoveMemberWrap;
         Functions["ReplaceMember"] = &TCallableConstraintTransformer::ReplaceMemberWrap;
         Functions["AsList"] = &TCallableConstraintTransformer::AsListWrap;
-        Functions["OptionalIf"] = &TCallableConstraintTransformer::FromSecond<TPassthroughConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TSortedConstraintNode, TPartOfSortedConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
-        Functions["ListIf"] = &TCallableConstraintTransformer::CopyAllFrom<1>;
-        Functions["FlatListIf"] = &TCallableConstraintTransformer::CopyAllFrom<1>;
-        Functions["FlatOptionalIf"] = &TCallableConstraintTransformer::CopyAllFrom<1>;
+        Functions["OptionalIf"] = &TCallableConstraintTransformer::PassOrEmptyWrap<false, false>;
+        Functions["FlatOptionalIf"] = &TCallableConstraintTransformer::PassOrEmptyWrap<false, true>;
+        Functions["ListIf"] = &TCallableConstraintTransformer::PassOrEmptyWrap<true, false>;
+        Functions["FlatListIf"] = &TCallableConstraintTransformer::PassOrEmptyWrap<true, true>;
         Functions["EmptyIterator"] = &TCallableConstraintTransformer::FromEmpty;
         Functions["List"] = &TCallableConstraintTransformer::ListWrap;
         Functions["Dict"] = &TCallableConstraintTransformer::DictWrap;
@@ -1698,6 +1698,31 @@ private:
         return TStatus::Ok;
     }
 
+    template<bool IsList, bool IsFlat>
+    TStatus PassOrEmptyWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
+        if (const auto part = input->Tail().GetConstraint<TPartOfSortedConstraintNode>())
+            if (const auto filtered = part->CompleteOnly(ctx))
+                input->AddConstraint(filtered);
+
+        if (const auto part = input->Tail().GetConstraint<TPartOfDistinctConstraintNode>())
+            if (const auto filtered = part->CompleteOnly(ctx))
+                input->AddConstraint(filtered);
+
+        if (const auto part = input->Tail().GetConstraint<TPartOfUniqueConstraintNode>())
+            if constexpr (IsList) {
+                if (const auto filtered = part->CompleteOnly(ctx))
+                    input->AddConstraint(filtered);
+            } else
+                input->AddConstraint(part);
+
+        if constexpr (IsFlat) {
+            if (const auto empty = input->Tail().GetConstraint<TEmptyConstraintNode>())
+                input->AddConstraint(empty);
+        }
+
+        return FromSecond<TPassthroughConstraintNode, TUniqueConstraintNode, TDistinctConstraintNode, TSortedConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>(input, output, ctx);
+    }
+
     TStatus IfWrap(const TExprNode::TPtr& input, TExprNode::TPtr&, TExprContext& ctx) const {
         std::vector<const TConstraintSet*> constraints;
         constraints.reserve((input->ChildrenSize() << 1U) + 1U);