From e5d9b3d5d07d40539a0114c52d02d1b28a7e1298 Mon Sep 17 00:00:00 2001
From: Vladislav Lukachik <vladluk@ydb.tech>
Date: Mon, 28 Oct 2024 12:00:21 +0000
Subject: [PATCH] Stop state from growing and spill growed is spilling is
 enabled

---
 .../minikql/comp_nodes/mkql_wide_combine.cpp  | 25 ++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/ydb/library/yql/minikql/comp_nodes/mkql_wide_combine.cpp b/ydb/library/yql/minikql/comp_nodes/mkql_wide_combine.cpp
index d823ad5bd5e0..637c95a76831 100644
--- a/ydb/library/yql/minikql/comp_nodes/mkql_wide_combine.cpp
+++ b/ydb/library/yql/minikql/comp_nodes/mkql_wide_combine.cpp
@@ -274,12 +274,21 @@ class TState : public TComputationValue<TState> {
             Tongue = CurrentPage->data() + CurrentPosition;
         }
         Throat = States.GetKey(itInsert) + KeyWidth;
-        if (isNew) {
-            States.CheckGrow();
+        if (isNew && !IsOutOfMemory) {
+            try {
+                States.CheckGrow();
+            } catch (const TMemoryLimitExceededException& e) {
+                Cerr << "State " << (void *)this << " no longer growing\n";
+                IsOutOfMemory = true;
+            }
         }
         return isNew;
     }
 
+    bool CheckIsOutOfMemory() const {
+        return IsOutOfMemory;
+    }
+
     template<bool SkipYields>
     bool ReadMore() {
         if constexpr (SkipYields) {
@@ -331,6 +340,7 @@ class TState : public TComputationValue<TState> {
 private:
     std::optional<TStorageIterator> ExtractIt;
     const ui32 KeyWidth, StateWidth;
+    bool IsOutOfMemory = false;
     ui64 CurrentPosition = 0;
     TRow* CurrentPage = nullptr;
     TStorage Storage;
@@ -451,6 +461,9 @@ class TSpillingSupportState : public TComputationValue<TSpillingSupportState> {
     ETasteResult TasteIt() {
         if (GetMode() == EOperatingMode::InMemory) {
             bool isNew = InMemoryProcessingState.TasteIt();
+            if (InMemoryProcessingState.CheckIsOutOfMemory()) {
+                StateWantsToSpill = true;
+            }
             Throat = InMemoryProcessingState.Throat;
             return isNew ? ETasteResult::Init : ETasteResult::Update;
         }
@@ -649,7 +662,11 @@ class TSpillingSupportState : public TComputationValue<TSpillingSupportState> {
     }
 
     bool CheckMemoryAndSwitchToSpilling() {
-        if (AllowSpilling && Ctx.SpillerFactory && IsSwitchToSpillingModeCondition()) {
+        if (!(AllowSpilling && Ctx.SpillerFactory)) {
+            return false;
+        }
+        if (StateWantsToSpill || IsSwitchToSpillingModeCondition()) {
+            StateWantsToSpill = false;
             LogMemoryUsage();
 
             SwitchMode(EOperatingMode::SplittingState);
@@ -831,6 +848,7 @@ class TSpillingSupportState : public TComputationValue<TSpillingSupportState> {
     }
 
     bool IsSwitchToSpillingModeCondition() const {
+        return false;
         return !HasMemoryForProcessing() || TlsAllocState->GetMaximumLimitValueReached();
     }
 
@@ -840,6 +858,7 @@ class TSpillingSupportState : public TComputationValue<TSpillingSupportState> {
     NUdf::TUnboxedValuePod* Throat = nullptr;
 
 private:
+    bool StateWantsToSpill = false;
     bool IsEverythingExtracted = false;
 
     TState InMemoryProcessingState;