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

3 files changed, 9 insertions, 6 deletions

diff --git a/ydb/library/yql/core/yql_aggregate_expander.h b/ydb/library/yql/core/yql_aggregate_expander.h
index 8161ad81751..aa261458dbb 100644
--- 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 @@ public:
         , HaveSessionSetting(false)
         , OriginalRowType(nullptr)
         , RowType(nullptr)
-        , UseBlocks(typesCtx.IsBlockEngineEnabled() && !allowSpilling)
+        , UseBlocks(useBlocks)
     {
         PreMap = Ctx.Builder(node->Pos())
             .Lambda()
@@ -135,7 +135,7 @@ private:
 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
index 0688e4df460..cf544dfa221 100644
--- 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
index 7653a1e6a23..66bf92da61f 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,8 @@ protected:
 
         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 @@ protected:
             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 {