Avoid block aggregation in auto mode for YT (#8894)

ydb-platform · Sep 7, 2024 · e9c2717 · e9c2717
1 parent c6d4d1b
commit e9c2717
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 6 deletions.
diff --git a/ydb/library/yql/core/yql_aggregate_expander.h b/ydb/library/yql/core/yql_aggregate_expander.h
@@ -9,7 +9,7 @@ namespace NYql {
 class TAggregateExpander {
 public:
     TAggregateExpander(bool usePartitionsByKeys, const bool useFinalizeByKeys, const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& typesCtx,
-        bool forceCompact = false, bool compactForDistinct = false, bool usePhases = false, bool allowSpilling = false)
+        bool forceCompact = false, bool compactForDistinct = false, bool usePhases = false, bool useBlocks = false)
         : Node(node)
         , Ctx(ctx)
         , TypesCtx(typesCtx)
@@ -25,7 +25,7 @@ class TAggregateExpander {
         , HaveSessionSetting(false)
         , OriginalRowType(nullptr)
         , RowType(nullptr)
-        , UseBlocks(typesCtx.IsBlockEngineEnabled() && !allowSpilling)
+        , UseBlocks(useBlocks)
     {
         PreMap = Ctx.Builder(node->Pos())
             .Lambda()
@@ -135,7 +135,7 @@ class TAggregateExpander {
 inline TExprNode::TPtr ExpandAggregatePeepholeImpl(const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& typesCtx,
     const bool useFinalizeByKey, const bool useBlocks, const bool allowSpilling) {
     TAggregateExpander aggExpander(!useFinalizeByKey && !useBlocks, useFinalizeByKey, node, ctx, typesCtx,
-        true, false, false, allowSpilling);
+        true, false, false, typesCtx.IsBlockEngineEnabled() && !allowSpilling);
     return aggExpander.ExpandAggregate();
 }
 

diff --git a/ydb/library/yql/dq/opt/dq_opt_log.cpp b/ydb/library/yql/dq/opt/dq_opt_log.cpp
@@ -23,7 +23,8 @@ TExprBase DqRewriteAggregate(TExprBase node, TExprContext& ctx, TTypeAnnotationC
         return node;
     }
     TAggregateExpander aggExpander(!typesCtx.IsBlockEngineEnabled() && !useFinalizeByKey,
-        useFinalizeByKey, node.Ptr(), ctx, typesCtx, false, compactForDistinct, usePhases, allowSpilling);
+        useFinalizeByKey, node.Ptr(), ctx, typesCtx, false, compactForDistinct, usePhases,
+        typesCtx.IsBlockEngineEnabled() && !allowSpilling);
     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
@@ -354,7 +354,8 @@ class TYtLogicalOptProposalTransformer : public TOptimizeTransformerBase {
 
         auto usePhases = State_->Configuration->UseAggPhases.Get().GetOrElse(false);
         auto usePartitionsByKeys = State_->Configuration->UsePartitionsByKeysForFinalAgg.Get().GetOrElse(true);
-        TAggregateExpander aggExpander(usePartitionsByKeys, false, node.Ptr(), ctx, *State_->Types, false, false, usePhases);
+        TAggregateExpander aggExpander(usePartitionsByKeys, false, node.Ptr(), ctx, *State_->Types, false, false,
+            usePhases, State_->Types->UseBlocks || State_->Types->BlockEngineMode == EBlockEngineMode::Force);
         return aggExpander.ExpandAggregate();
     }
 
@@ -2518,7 +2519,8 @@ class TYtLogicalOptProposalTransformer : public TOptimizeTransformerBase {
             return node;
         }
 
-        return TAggregateExpander::CountAggregateRewrite(aggregate, ctx, State_->Types->IsBlockEngineEnabled());
+        return TAggregateExpander::CountAggregateRewrite(aggregate, ctx,
+            State_->Types->UseBlocks || State_->Types->BlockEngineMode == EBlockEngineMode::Force);
     }
 
     TMaybeNode<TExprBase> ZeroSampleToZeroLimit(TExprBase node, TExprContext& ctx) const {