YQL-8971 YQL-15555 Unique/Distinct/Chopped some fixes.

2 files changed, 106 insertions, 48 deletions

diff --git a/ydb/library/yql/ast/yql_constraint.cpp b/ydb/library/yql/ast/yql_constraint.cpp
index 338e4268c7..6aca0bb779 100644
--- a/ydb/library/yql/ast/yql_constraint.cpp
+++ b/ydb/library/yql/ast/yql_constraint.cpp
@@ -536,39 +536,44 @@ const TChoppedConstraintNode* TChoppedConstraintNode::MakeCommon(const std::vect
 
 const TChoppedConstraintNode*
 TChoppedConstraintNode::FilterFields(TExprContext& ctx, const TPathFilter& predicate) const {
-    auto sets = Sets_;
-    for (auto it = sets.cbegin(); sets.cend() != it;) {
-        if (std::all_of(it->cbegin(), it->cend(), predicate))
-            ++it;
-        else
-            it = sets.erase(it);
+    TFullSetType chopped;
+    chopped.reserve(Sets_.size());
+    for (const auto& set : Sets_) {
+        auto newSet = set;
+        for (auto it = newSet.cbegin(); newSet.cend() != it;) {
+            if (predicate(*it))
+                ++it;
+            else
+                it = newSet.erase(it);
+        }
+
+        if (newSet.empty())
+            return nullptr;;
+
+        chopped.insert_unique(std::move(newSet));
     }
-    return sets.empty() ? nullptr : ctx.MakeConstraint<TChoppedConstraintNode>(std::move(sets));
+    return ctx.MakeConstraint<TChoppedConstraintNode>(std::move(chopped));
 }
 
 const TChoppedConstraintNode*
 TChoppedConstraintNode::RenameFields(TExprContext& ctx, const TPathReduce& reduce) const {
-    TFullSetType sets;
-    sets.reserve(Sets_.size());
+    TFullSetType chopped;
+    chopped.reserve(Sets_.size());
     for (const auto& set : Sets_) {
-        std::vector<TSetType> newSets(1U);
-        newSets.front().reserve(set.size());
+        TSetType newSet;
+        newSet.reserve(set.size());
         for (const auto& path : set) {
-            auto newPaths = reduce(path);
-            if (newPaths.empty())
-                break;
-            auto tmpSets(std::move(newSets));
-            for (const auto& newPath : newPaths) {
-                for (const auto& oldSet : tmpSets) {
-                    newSets.emplace_back(oldSet);
-                    newSets.back().insert_unique(newPath);
-                }
-            }
+            if (const auto& newPaths = reduce(path); !newPaths.empty())
+                newSet.insert_unique(newPaths.cbegin(), newPaths.cend());
         }
-        if (set.size() == newSets.front().size())
-            sets.insert_unique(newSets.cbegin(), newSets.cend());
+
+        if (newSet.empty())
+            return nullptr;
+
+        chopped.insert_unique(std::move(newSet));
     }
-    return sets.empty() ? nullptr : ctx.MakeConstraint<TChoppedConstraintNode>(std::move(sets));
+
+    return ctx.MakeConstraint<TChoppedConstraintNode>(std::move(chopped));
 }
 
 const TChoppedConstraintNode*
@@ -1990,6 +1995,11 @@ void Out<NYql::TPartOfSortedConstraintNode>(IOutputStream& out, const NYql::TPar
 }
 
 template<>
+void Out<NYql::TPartOfChoppedConstraintNode>(IOutputStream& out, const NYql::TPartOfChoppedConstraintNode& c) {
+    c.Out(out);
+}
+
+template<>
 void Out<NYql::TPartOfUniqueConstraintNode>(IOutputStream& out, const NYql::TPartOfUniqueConstraintNode& c) {
     c.Out(out);
 }
diff --git a/ydb/library/yql/core/yql_expr_constraint.cpp b/ydb/library/yql/core/yql_expr_constraint.cpp
index 075240f49e..b54ea754bf 100644
--- a/ydb/library/yql/core/yql_expr_constraint.cpp
+++ b/ydb/library/yql/core/yql_expr_constraint.cpp
@@ -150,10 +150,10 @@ public:
         Functions["Limit"] = &TCallableConstraintTransformer::TakeWrap;
         Functions["Member"] = &TCallableConstraintTransformer::MemberWrap;
         Functions["AsStruct"] = &TCallableConstraintTransformer::AsStructWrap;
-        Functions["Just"] = &TCallableConstraintTransformer::FromFirst<TPassthroughConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TPartOfSortedConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
-        Functions["Unwrap"] = &TCallableConstraintTransformer::FromFirst<TPassthroughConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TPartOfSortedConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
+        Functions["Just"] = &TCallableConstraintTransformer::FromFirst<TPassthroughConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TPartOfSortedConstraintNode, TPartOfChoppedConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
+        Functions["Unwrap"] = &TCallableConstraintTransformer::FromFirst<TPassthroughConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TPartOfSortedConstraintNode, TPartOfChoppedConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
         Functions["Ensure"] = &TCallableConstraintTransformer::CopyAllFrom<0>;
-        Functions["ToList"] = &TCallableConstraintTransformer::FromFirst<TPassthroughConstraintNode, TEmptyConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TPartOfSortedConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
+        Functions["ToList"] = &TCallableConstraintTransformer::FromFirst<TPassthroughConstraintNode, TEmptyConstraintNode, TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TPartOfSortedConstraintNode, TPartOfChoppedConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
         Functions["ToOptional"] = &TCallableConstraintTransformer::FromFirst<TPassthroughConstraintNode, TEmptyConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
         Functions["Head"] = &TCallableConstraintTransformer::FromFirst<TPassthroughConstraintNode, TEmptyConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
         Functions["Last"] = &TCallableConstraintTransformer::FromFirst<TPassthroughConstraintNode, TEmptyConstraintNode, TVarIndexConstraintNode, TMultiConstraintNode>;
@@ -203,15 +203,11 @@ public:
         Functions["DictItems"] = &TCallableConstraintTransformer::FromFirst<TPassthroughConstraintNode, TUniqueConstraintNode, TDistinctConstraintNode, TEmptyConstraintNode>;
         Functions["DictKeys"] = &TCallableConstraintTransformer::DictHalfWrap<true>;
         Functions["DictPayloads"] = &TCallableConstraintTransformer::DictHalfWrap<false>;
-        Functions["FoldMap"] = &TCallableConstraintTransformer::InheriteEmptyFromInput; // TODO: passthrough
-        Functions["Fold1Map"] = &TCallableConstraintTransformer::InheriteEmptyFromInput; // TODO: passthrough
-        Functions["Chain1Map"] = &TCallableConstraintTransformer::InheriteEmptyFromInput; // TODO: passthrough, sorted, unique
-        Functions["WideChain1Map"] = &TCallableConstraintTransformer::InheriteEmptyFromInput; // TODO: passthrough, sorted, unique
+        Functions["Chain1Map"] = &TCallableConstraintTransformer::Chain1MapWrap<false>;
+        Functions["WideChain1Map"] = &TCallableConstraintTransformer::Chain1MapWrap<true>;
         Functions["IsKeySwitch"] = &TCallableConstraintTransformer::IsKeySwitchWrap;
         Functions["Condense"] = &TCallableConstraintTransformer::CondenseWrap;
         Functions["Condense1"] = &TCallableConstraintTransformer::Condense1Wrap<false>;
-        Functions["Squeeze"] = &TCallableConstraintTransformer::InheriteEmptyFromInput;
-        Functions["Squeeze1"] = &TCallableConstraintTransformer::InheriteEmptyFromInput;
         Functions["GroupingCore"] = &TCallableConstraintTransformer::InheriteEmptyFromInput;
         Functions["Chopper"] = &TCallableConstraintTransformer::InheriteEmptyFromInput;
         Functions["WideChopper"] = &TCallableConstraintTransformer::InheriteEmptyFromInput;
@@ -527,6 +523,12 @@ private:
                 }
             }
 
+            if (const auto part = input->Head().GetConstraint<TPartOfChoppedConstraintNode>()) {
+                if (const auto filtered = part->FilterFields(ctx, filter)) {
+                    input->AddConstraint(filtered);
+                }
+            }
+
             if (const auto part = input->Head().GetConstraint<TPartOfUniqueConstraintNode>()) {
                 if (const auto filtered = part->FilterFields(ctx, filter)) {
                     input->AddConstraint(filtered);
@@ -767,7 +769,7 @@ private:
     template<class TConstraintType>
     static const TConstraintType* GetLambdaConstraint(const TExprNode& lambda, TExprContext& ctx) {
         if (2U == lambda.ChildrenSize())
-            return lambda.GetConstraint<TConstraintType>();
+            return lambda.Tail().GetConstraint<TConstraintType>();
 
         TVector<const TConstraintSet*> constraints;
         constraints.reserve(lambda.ChildrenSize() - 1U);
@@ -807,6 +809,12 @@ private:
             }
         }
 
+        if (body.IsCallable({"CastStruct","FilterMembers"}))
+            return GetPathToKey(body.Head(), arg);
+        if (body.IsCallable("Member") && body.Head().IsCallable("AsStruct"))
+            return GetPathToKey(GetLiteralStructMember(body.Head(), body.Tail()), arg);
+        if (body.IsCallable("Nth") && body.Head().IsList())
+            return GetPathToKey(*body.Head().Child(FromString<ui32>(body.Tail().Content())), arg);
         if (IsTransparentIfPresent(body) && &body.Head() == &arg)
             return GetPathToKey(body.Child(1)->Tail().Head(), body.Child(1)->Head().Head());
 
@@ -2610,6 +2618,8 @@ private:
                 return GetPathToKey(GetLiteralStructMember(body.Head(), body.Tail()), args);
             } else if (body.IsCallable("Nth") && body.Head().IsList()) {
                 return GetPathToKey(*body.Head().Child(FromString<ui32>(body.Tail().Content())), args);
+            } else if (body.IsCallable({"CastStruct","FilterMembers"}))  {
+                return GetPathToKey(body.Head(), args);
             }
         }
 
@@ -2662,10 +2672,11 @@ private:
                     if (auto path = GetPathToKey(*body.Child(i), arg))
                         keys.insert_unique(std::move(*path));
             }
-        } else if (body.IsCallable({"StablePickle"})) {
+        } else if (body.IsCallable("StablePickle")) {
             return GetSimpleKeys(body.Head(), arg);
-        } else if (auto path = GetPathToKey(body, arg))
+        } else if (auto path = GetPathToKey(body, arg)) {
             keys.insert_unique(std::move(*path));
+        }
 
         return keys;
     }
@@ -2723,10 +2734,12 @@ private:
     template<class TConstraint, bool Wide>
     static void GetCommonFromBothLambdas(const TExprNode::TPtr& input, const typename TConstraint::TMainConstraint* original, TExprContext& ctx) {
         if (original)
-            if (const auto init = TConstraint::MakeComplete(ctx, GetConstraintFromLambda<TConstraint, Wide>(*input->Child(1), ctx)->GetColumnMapping(), original))
-                if (const auto update = TConstraint::MakeComplete(ctx, GetConstraintFromLambda<TConstraint, Wide>(input->Tail(), ctx)->GetColumnMapping(), original))
-                    if (const auto common = init->MakeCommon(update, ctx))
-                        input->AddConstraint(common);
+            if (const auto initPart = GetConstraintFromLambda<TConstraint, Wide>(*input->Child(1), ctx))
+                if (const auto init = TConstraint::MakeComplete(ctx, initPart->GetColumnMapping(), original))
+                    if (const auto updatePart = GetConstraintFromLambda<TConstraint, Wide>(input->Tail(), ctx))
+                        if (const auto update = TConstraint::MakeComplete(ctx, updatePart->GetColumnMapping(), original))
+                            if (const auto common = init->MakeCommon(update, ctx))
+                                input->AddConstraint(common);
     }
 
     template<bool Wide>
@@ -2808,6 +2821,41 @@ private:
         return FromFirst<TEmptyConstraintNode>(input, output, ctx);
     }
 
+    template<class TConstraint, bool Wide>
+    static void GetCommonFromBothLambdas(const TExprNode::TPtr& input, TExprContext& ctx) {
+        if (const auto original = input->Head().GetConstraint<typename TConstraint::TMainConstraint>())
+            if (const auto initPart = GetConstraintFromLambda<TConstraint, Wide>(*input->Child(1), ctx))
+                if (const auto init = TConstraint::MakeComplete(ctx, initPart->GetColumnMapping(), original))
+                    if (const auto updatePart = GetConstraintFromLambda<TConstraint, Wide>(input->Tail(), ctx))
+                        if (const auto update = TConstraint::MakeComplete(ctx, updatePart->GetColumnMapping(), original))
+                            if (const auto common = init->MakeCommon(update, ctx))
+                                input->AddConstraint(common);
+    }
+
+    template<bool Wide>
+    TStatus Chain1MapWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
+        auto argsConstraints = GetConstraintsForInputArgument<true, Wide>(*input, ctx);
+        if (const auto status = UpdateLambdaConstraints(input->ChildRef(1), ctx, argsConstraints); status != TStatus::Ok) {
+            return status;
+        }
+
+        const auto initLambda = input->Child(1);
+        argsConstraints.reserve(argsConstraints.size() + initLambda->ChildrenSize() - 1U);
+        for (ui32 i = 1U; i < initLambda->ChildrenSize(); ++i) {
+            argsConstraints.emplace_back(initLambda->Child(i)->GetAllConstraints());
+        }
+
+        if (const auto status = UpdateLambdaConstraints(input->ChildRef(2), ctx, argsConstraints); status != TStatus::Ok) {
+            return status;
+        }
+
+        GetCommonFromBothLambdas<TPartOfSortedConstraintNode, Wide>(input, ctx);
+        GetCommonFromBothLambdas<TPartOfChoppedConstraintNode, Wide>(input, ctx);
+        GetCommonFromBothLambdas<TPartOfUniqueConstraintNode, Wide>(input, ctx);
+        GetCommonFromBothLambdas<TPartOfDistinctConstraintNode, Wide>(input, ctx);
+        return FromFirst<TEmptyConstraintNode>(input, output, ctx);
+    }
+
     template<bool Distinct>
     static void GetUniquesForPayloads(const TExprNode::TPtr& input, TExprContext& ctx) {
         typename TUniqueConstraintNodeBase<Distinct>::TFullSetType sets{TConstraintNode::TSetType{TConstraintNode::TPathType{ctx.GetIndexAsString(0U)}}};
@@ -2870,11 +2918,17 @@ private:
             return status;
         }
 
+        TConstraintNode::TSetType keys;
         if constexpr (Partitions) {
+            keys = GetSimpleKeys(input->Child(TCoBase::idx_KeySelectorLambda)->Tail(), input->Child(TCoBase::idx_KeySelectorLambda)->Head().Head());
             if (const auto sortKeySelector = input->Child(TCoBase::idx_SortKeySelectorLambda); sortKeySelector->IsLambda()) {
                 if (const auto status = UpdateLambdaConstraints(*sortKeySelector); status != TStatus::Ok) {
                     return status;
                 }
+
+                auto sortKeys = GetSimpleKeys(sortKeySelector->Tail(), sortKeySelector->Head().Head());
+                std::move(sortKeys.begin(), sortKeys.end(), std::back_inserter(keys));
+                std::sort(keys.begin(), keys.end());
             }
         }
 
@@ -2891,14 +2945,8 @@ private:
                 argConstraints.emplace_back(filtered);
 
         if constexpr (Partitions) {
-            TVector<TStringBuf> partitionKeys; // TODO: Replace on TSetType
-            ExtractKeys(*input->Child(TCoBase::idx_KeySelectorLambda), partitionKeys);
-            if (!partitionKeys.empty()) {
-                TChoppedConstraintNode::TFullSetType sets;
-                sets.reserve(partitionKeys.size());
-                std::transform(partitionKeys.cbegin(), partitionKeys.cend(), std::back_inserter(sets), [](const TStringBuf& column) { return TConstraintNode::TSetType{TConstraintNode::TPathType(1U, column)}; });
-                argConstraints.emplace_back(ctx.MakeConstraint<TChoppedConstraintNode>(std::move(sets)));
-            }
+            if (!keys.empty())
+                argConstraints.emplace_back(ctx.MakeConstraint<TChoppedConstraintNode>(keys));
         }
 
         if (const auto status = UpdateLambdaConstraints(input->ChildRef(TCoBase::idx_ListHandlerLambda), ctx, {argConstraints}); status != TStatus::Ok) {