Save distinct constraint for PruneKeys

commit_hash:c571530f0a74edc44a65da0388e8e4c543121044

1 files changed, 41 insertions, 10 deletions

diff --git a/yql/essentials/core/yql_expr_constraint.cpp b/yql/essentials/core/yql_expr_constraint.cpp
index 7fcf560a2cf..dede320a721 100644
--- a/yql/essentials/core/yql_expr_constraint.cpp
+++ b/yql/essentials/core/yql_expr_constraint.cpp
@@ -822,23 +822,54 @@ private:
         }
 
         if constexpr (Adjacent) {
-            if (const auto status = CopyAllFrom<0>(input, output, ctx); status != TStatus::Ok) {
-                return status;
-            }
+            TSet<TStringBuf> except = {
+                TUniqueConstraintNode::Name(),
+                TDistinctConstraintNode::Name(),
+            };
+            CopyExcept(*input, *input->Child(0), except);
+
+            auto uniqueConstraint = input->Child(0)->GetConstraint<TUniqueConstraintNode>();
+            auto distinctConstraint = input->Child(0)->GetConstraint<TDistinctConstraintNode>();
 
             TPartOfConstraintBase::TSetType keys = GetPathsToKeys<true>(input->Child(1)->Tail(), input->Child(1)->Head().Head());
-            TPartOfConstraintBase::TSetOfSetsType uniqueKeys;
-            for (const auto& elem : keys) {
-                uniqueKeys.insert(TPartOfConstraintBase::TSetType{elem});
-            }
             if (!keys.empty()) {
-                input->AddConstraint(ctx.MakeConstraint<TUniqueConstraintNode>(TUniqueConstraintNode::TContentType{uniqueKeys}));
-                input->AddConstraint(ctx.MakeConstraint<TDistinctConstraintNode>(TDistinctConstraintNode::TContentType{uniqueKeys}));
+                TPartOfConstraintBase::TSetOfSetsType uniqueKeys;
+                for (const auto& elem : keys) {
+                    uniqueKeys.insert(TPartOfConstraintBase::TSetType{elem});
+                }
+
+                auto newUniqueConstraint = ctx.MakeConstraint<TUniqueConstraintNode>(TUniqueConstraintNode::TContentType{uniqueKeys});
+                if (uniqueConstraint) {
+                    uniqueConstraint = TUniqueConstraintNode::Merge(
+                        newUniqueConstraint,
+                        dynamic_cast<const TUniqueConstraintNode*>(uniqueConstraint),
+                        ctx);
+                } else {
+                    uniqueConstraint = newUniqueConstraint;
+                }
+
+                auto newDistinctConstraint = ctx.MakeConstraint<TDistinctConstraintNode>(TDistinctConstraintNode::TContentType{uniqueKeys});
+                if (distinctConstraint) {
+                    distinctConstraint = TDistinctConstraintNode::Merge(
+                        newDistinctConstraint,
+                        dynamic_cast<const TDistinctConstraintNode*>(distinctConstraint),
+                        ctx);
+                } else {
+                    distinctConstraint = newDistinctConstraint;
+                }
+            }
+
+            if (uniqueConstraint) {
+                input->AddConstraint(uniqueConstraint);
+            }
+            if (distinctConstraint) {
+                input->AddConstraint(distinctConstraint);
             }
+
             return TStatus::Ok;
         }
 
-        return FromFirst<TEmptyConstraintNode>(input, output, ctx);
+        return FromFirst<TEmptyConstraintNode, TUniqueConstraintNode, TDistinctConstraintNode>(input, output, ctx);
     }
 
     template<class TConstraint>