KIKIMR-19512 Setting for memory limit in first stage combiner.

4 files changed, 45 insertions, 2 deletions

diff --git a/ydb/core/kqp/opt/peephole/kqp_opt_peephole.cpp b/ydb/core/kqp/opt/peephole/kqp_opt_peephole.cpp
index 16b322b860..b8ac524ff9 100644
--- a/ydb/core/kqp/opt/peephole/kqp_opt_peephole.cpp
+++ b/ydb/core/kqp/opt/peephole/kqp_opt_peephole.cpp
@@ -174,13 +174,53 @@ private:
     TKikimrConfiguration::TPtr Config;
 };
 
+class TKqpPeepholeFinalTransformer : public TOptimizeTransformerBase {
+public:
+    TKqpPeepholeFinalTransformer(TTypeAnnotationContext& ctx, TKikimrConfiguration::TPtr config)
+        : TOptimizeTransformerBase(&ctx, NYql::NLog::EComponent::ProviderKqp, {}), Config(config)
+    {
+#define HNDL(name) "KqpPeepholeFinal-"#name, Hndl(&TKqpPeepholeFinalTransformer::name)
+        AddHandler(0, &TCoWideCombiner::Match, HNDL(SetCombinerMemoryLimit));
+#undef HNDL
+    }
+private:
+    TMaybeNode<TExprBase> SetCombinerMemoryLimit(TExprBase node, TExprContext& ctx) {
+        if (const auto limit = node.Ref().Child(TCoWideCombiner::idx_MemLimit); limit->IsAtom("0")) {
+            if (const auto limitSetting = Config->_KqpYqlCombinerMemoryLimit.Get(); limitSetting && *limitSetting) {
+                return ctx.ChangeChild(node.Ref(), TCoWideCombiner::idx_MemLimit, ctx.RenameNode(*limit, ToString(-i64(*limitSetting))));
+            }
+        }
+        return node;
+    }
+
+    const TKikimrConfiguration::TPtr Config;
+};
+
+struct TKqpPeepholePipelineFinalConfigurator : IPipelineConfigurator {
+    TKqpPeepholePipelineFinalConfigurator(TKikimrConfiguration::TPtr config)
+        : Config(config)
+    {}
+
+    void AfterCreate(TTransformationPipeline*) const override {}
+
+    void AfterTypeAnnotation(TTransformationPipeline* pipeline) const override {
+        pipeline->Add(new TKqpPeepholeFinalTransformer(*pipeline->GetTypeAnnotationContext(), Config), "KqpPeepholeFinal");
+    }
+
+    void AfterOptimize(TTransformationPipeline*) const override {}
+private:
+    const TKikimrConfiguration::TPtr Config;
+};
+
 TStatus PeepHoleOptimize(const TExprBase& program, TExprNode::TPtr& newProgram, TExprContext& ctx,
     IGraphTransformer& typeAnnTransformer, TTypeAnnotationContext& typesCtx, TKikimrConfiguration::TPtr config,
     bool allowNonDeterministicFunctions, bool withFinalStageRules)
 {
     TKqpPeepholePipelineConfigurator kqpPeephole(config);
+    TKqpPeepholePipelineFinalConfigurator kqpPeepholeFinal(config);
     TPeepholeSettings peepholeSettings;
     peepholeSettings.CommonConfig = &kqpPeephole;
+    peepholeSettings.FinalConfig = &kqpPeepholeFinal;
     peepholeSettings.WithFinalStageRules = withFinalStageRules;
     peepholeSettings.WithNonDeterministicRules = false;
 
diff --git a/ydb/core/kqp/provider/yql_kikimr_settings.cpp b/ydb/core/kqp/provider/yql_kikimr_settings.cpp
index f1f86737b6..8f4a9229c0 100644
--- a/ydb/core/kqp/provider/yql_kikimr_settings.cpp
+++ b/ydb/core/kqp/provider/yql_kikimr_settings.cpp
@@ -1,6 +1,7 @@
 #include "yql_kikimr_settings.h"
 
 #include <ydb/core/protos/config.pb.h>
+#include <util/generic/size_literals.h>
 
 namespace NYql {
 
@@ -43,6 +44,7 @@ TKikimrConfiguration::TKikimrConfiguration() {
     REGISTER_SETTING(*this, _KqpMaxComputeActors);
     REGISTER_SETTING(*this, _KqpEnableSpilling);
     REGISTER_SETTING(*this, _KqpDisableLlvmForUdfStages);
+    REGISTER_SETTING(*this, _KqpYqlCombinerMemoryLimit).Lower(0ULL).Upper(1_GB);
 
     REGISTER_SETTING(*this, KqpPushOlapProcess);
 
@@ -65,7 +67,7 @@ TKikimrConfiguration::TKikimrConfiguration() {
     REGISTER_SETTING(*this, OptEnablePredicateExtract);
     REGISTER_SETTING(*this, OptEnableOlapPushdown);
     REGISTER_SETTING(*this, OptEnableOlapProvideComputeSharding);
-    
+
     REGISTER_SETTING(*this, OptUseFinalizeByKey);
     REGISTER_SETTING(*this, OptEnableCostBasedOptimization);
     REGISTER_SETTING(*this, MaxDPccpDPTableSize);
diff --git a/ydb/core/kqp/provider/yql_kikimr_settings.h b/ydb/core/kqp/provider/yql_kikimr_settings.h
index 13dbdc4140..c773ad299b 100644
--- a/ydb/core/kqp/provider/yql_kikimr_settings.h
+++ b/ydb/core/kqp/provider/yql_kikimr_settings.h
@@ -36,6 +36,7 @@ struct TKikimrSettings {
     NCommon::TConfSetting<ui32, false> _KqpMaxComputeActors;
     NCommon::TConfSetting<bool, false> _KqpEnableSpilling;
     NCommon::TConfSetting<bool, false> _KqpDisableLlvmForUdfStages;
+    NCommon::TConfSetting<ui64, false> _KqpYqlCombinerMemoryLimit;
 
     /* No op just to avoid errors in Cloud Logging until they remove this from their queries */
     NCommon::TConfSetting<bool, false> KqpPushOlapProcess;
diff --git a/ydb/library/yql/core/type_ann/type_ann_wide.cpp b/ydb/library/yql/core/type_ann/type_ann_wide.cpp
index f6f1c922e3..6e6e71b0d2 100644
--- a/ydb/library/yql/core/type_ann/type_ann_wide.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_wide.cpp
@@ -314,7 +314,7 @@ IGraphTransformer::TStatus WideCombinerWrapper(const TExprNode::TPtr& input, TEx
     }
 
     if (const auto& limit = input->Child(1U)->Content(); !limit.empty()) {
-        ui64 memLimit = 0ULL;
+        i64 memLimit = 0LL;
         if (!TryFromString(limit, memLimit)) {
             ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Child(1U)->Pos()), TStringBuilder() <<
                 "Bad memLimit value: " << limit));