KIKIMR-21018: always allow pickling in AggregateExpander (#1742)

6 files changed, 16 insertions, 13 deletions

diff --git a/ydb/library/yql/core/yql_aggregate_expander.cpp b/ydb/library/yql/core/yql_aggregate_expander.cpp
index d88364fd8ae..56f278f0f99 100644
--- a/ydb/library/yql/core/yql_aggregate_expander.cpp
+++ b/ydb/library/yql/core/yql_aggregate_expander.cpp
@@ -273,7 +273,7 @@ bool TAggregateExpander::IsNeedPickle(const TVector<const TTypeAnnotationNode*>&
 {
     bool needPickle = false;
     for (auto type : keyItemTypes) {
-        needPickle = needPickle || AllowPickle && !IsDataOrOptionalOfData(type);
+        needPickle |= !IsDataOrOptionalOfData(type);
     }
     return needPickle;
 }
diff --git a/ydb/library/yql/core/yql_aggregate_expander.h b/ydb/library/yql/core/yql_aggregate_expander.h
index aa5adeeccf5..1bdfdef8542 100644
--- a/ydb/library/yql/core/yql_aggregate_expander.h
+++ b/ydb/library/yql/core/yql_aggregate_expander.h
@@ -8,12 +8,11 @@ namespace NYql {
 
 class TAggregateExpander {
 public:
-    TAggregateExpander(bool allowPickle, bool usePartitionsByKeys, const bool useFinalizeByKeys, const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& typesCtx,
+    TAggregateExpander(bool usePartitionsByKeys, const bool useFinalizeByKeys, const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& typesCtx,
         bool forceCompact = false, bool compactForDistinct = false, bool usePhases = false)
         : Node(node)
         , Ctx(ctx)
         , TypesCtx(typesCtx)
-        , AllowPickle(allowPickle)
         , UsePartitionsByKeys(usePartitionsByKeys)
         , UseFinalizeByKeys(useFinalizeByKeys)
         , ForceCompact(forceCompact)
@@ -93,7 +92,6 @@ private:
     const TExprNode::TPtr Node;
     TExprContext& Ctx;
     TTypeAnnotationContext& TypesCtx;
-    bool AllowPickle;
     bool UsePartitionsByKeys;
     bool UseFinalizeByKeys = false;
     bool ForceCompact;
@@ -132,7 +130,7 @@ private:
 };
 
 inline TExprNode::TPtr ExpandAggregatePeepholeImpl(const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& typesCtx, const bool useFinalizeByKey, const bool useBlocks) {
-    TAggregateExpander aggExpander(false, !useFinalizeByKey && !useBlocks, useFinalizeByKey, node, ctx, typesCtx, true);
+    TAggregateExpander aggExpander(!useFinalizeByKey && !useBlocks, useFinalizeByKey, node, ctx, typesCtx, true);
     return aggExpander.ExpandAggregate();
 }
 
diff --git a/ydb/library/yql/core/yql_expr_constraint.cpp b/ydb/library/yql/core/yql_expr_constraint.cpp
index 8a8814e8803..c2667a358b1 100644
--- a/ydb/library/yql/core/yql_expr_constraint.cpp
+++ b/ydb/library/yql/core/yql_expr_constraint.cpp
@@ -2654,9 +2654,10 @@ private:
     }
 
     template<bool Wide>
-    static TPartOfConstraintBase::TSetType GetSimpleKeys(const TExprNode& body, const TExprNode::TChildrenType& args, TExprContext& ctx) {
+    static TPartOfConstraintBase::TSetType GetSimpleKeys(const TExprNode& node, const TExprNode::TChildrenType& args, TExprContext& ctx) {
         TPartOfConstraintBase::TSetType keys;
-        if (body.IsCallable("AggrNotEquals")) {
+        if (node.IsCallable("AggrNotEquals")) {
+            const TExprNode& body = node.Head().IsCallable("StablePickle") ? node.Head() : node;
             if (body.Head().IsList() && body.Tail().IsList() && body.Head().ChildrenSize() == body.Tail().ChildrenSize()) {
                 keys.reserve(body.Tail().ChildrenSize());
                 for (auto i = 0U; i < body.Head().ChildrenSize(); ++i){
@@ -2679,10 +2680,10 @@ private:
                     keys.insert_unique(r->first);
                 }
             }
-        } else if (body.IsCallable("Or")) {
-            keys.reserve(body.ChildrenSize());
-            for (auto i = 0U; i < body.ChildrenSize(); ++i) {
-                const auto& part = GetSimpleKeys<Wide>(*body.Child(i), args, ctx);
+        } else if (node.IsCallable("Or")) {
+            keys.reserve(node.ChildrenSize());
+            for (auto i = 0U; i < node.ChildrenSize(); ++i) {
+                const auto& part = GetSimpleKeys<Wide>(*node.Child(i), args, ctx);
                 keys.insert_unique(part.cbegin(), part.cend());
             }
         }
diff --git a/ydb/library/yql/core/yql_opt_utils.cpp b/ydb/library/yql/core/yql_opt_utils.cpp
index 9342e122705..53f38546051 100644
--- a/ydb/library/yql/core/yql_opt_utils.cpp
+++ b/ydb/library/yql/core/yql_opt_utils.cpp
@@ -2021,6 +2021,8 @@ std::optional<std::pair<TPartOfConstraintBase::TPathType, ui32>> GetPathToKey(co
         } else if (body.IsCallable({"CastStruct","FilterMembers"}))  {
             return GetPathToKey(body.Head(), args);
         }
+    } else if (body.IsCallable("StablePickle")) {
+        return GetPathToKey(body.Head(), args);
     }
 
     return std::nullopt;
@@ -2054,6 +2056,8 @@ std::optional<TPartOfConstraintBase::TPathType> GetPathToKey(const TExprNode& bo
         return GetPathToKey(body.Head().Tail().Head(), arg);
     if (IsTransparentIfPresent(body) && &body.Head() == &arg)
         return GetPathToKey(body.Child(1)->Tail().Head(), body.Child(1)->Head().Head());
+    if (body.IsCallable("StablePickle"))
+        return GetPathToKey(body.Head(), arg);
 
     return std::nullopt;
 }
diff --git a/ydb/library/yql/dq/opt/dq_opt_log.cpp b/ydb/library/yql/dq/opt/dq_opt_log.cpp
index b86fc9da74d..cf86e2a7f21 100644
--- a/ydb/library/yql/dq/opt/dq_opt_log.cpp
+++ b/ydb/library/yql/dq/opt/dq_opt_log.cpp
@@ -22,7 +22,7 @@ TExprBase DqRewriteAggregate(TExprBase node, TExprContext& ctx, TTypeAnnotationC
     if (!node.Maybe<TCoAggregateBase>()) {
         return node;
     }
-    TAggregateExpander aggExpander(true, !typesCtx.IsBlockEngineEnabled() && !useFinalizeByKey, useFinalizeByKey, node.Ptr(), ctx, typesCtx, false, compactForDistinct, usePhases);
+    TAggregateExpander aggExpander(!typesCtx.IsBlockEngineEnabled() && !useFinalizeByKey, useFinalizeByKey, node.Ptr(), ctx, typesCtx, false, compactForDistinct, usePhases);
     auto result = aggExpander.ExpandAggregate();
     YQL_ENSURE(result);
 
diff --git a/ydb/library/yql/providers/yt/provider/yql_yt_logical_optimize.cpp b/ydb/library/yql/providers/yt/provider/yql_yt_logical_optimize.cpp
index 35ca38e8c0d..b9d080afe1a 100644
--- a/ydb/library/yql/providers/yt/provider/yql_yt_logical_optimize.cpp
+++ b/ydb/library/yql/providers/yt/provider/yql_yt_logical_optimize.cpp
@@ -354,7 +354,7 @@ protected:
 
         auto usePhases = State_->Configuration->UseAggPhases.Get().GetOrElse(false);
         auto usePartitionsByKeys = State_->Configuration->UsePartitionsByKeysForFinalAgg.Get().GetOrElse(true);
-        TAggregateExpander aggExpander(true, usePartitionsByKeys, false, node.Ptr(), ctx, *State_->Types, false, false, usePhases);
+        TAggregateExpander aggExpander(usePartitionsByKeys, false, node.Ptr(), ctx, *State_->Types, false, false, usePhases);
         return aggExpander.ExpandAggregate();
     }