dont allocate iterators on mkql allocator

13 files changed, 198 insertions, 126 deletions

diff --git a/ydb/core/engine/minikql/minikql_engine_host.cpp b/ydb/core/engine/minikql/minikql_engine_host.cpp
index 4e5a4688c34..2442fe7bfd7 100644
--- a/ydb/core/engine/minikql/minikql_engine_host.cpp
+++ b/ydb/core/engine/minikql/minikql_engine_host.cpp
@@ -344,7 +344,9 @@ NUdf::TUnboxedValue TEngineHost::SelectRow(const TTableId& tableId, const TArray
 }
 
 template<class TTableIt>
-class TSelectRangeLazyRow : public TComputationValue<TSelectRangeLazyRow<TTableIt>>{
+class TSelectRangeLazyRow : public TComputationValue<TSelectRangeLazyRow<TTableIt>> {
+private:
+    using TBase = TComputationValue<TSelectRangeLazyRow<TTableIt>>;
     NUdf::TUnboxedValue GetElement(ui32 index) const override {
         BuildValue(index);
         return GetPtr()[index];
@@ -362,7 +364,7 @@ public:
     void operator delete(void *mem, std::size_t sz) {
         auto ptr = (TSelectRangeLazyRow*)mem;
         auto extraSize = ptr->Size() * sizeof(NUdf::TUnboxedValue) + ptr->GetMaskSize() * sizeof(ui64);
-        MKQLFreeWithSize(mem, sz + extraSize);
+        TBase::FreeWithSize(mem, sz + extraSize);
     }
 
     void operator delete[](void *mem, std::size_t sz) = delete;
@@ -372,7 +374,7 @@ public:
         ui32 size = dbData.ColumnCount + systemColumnTags.size();
 
         ui32 maskSize = size > 0 ? (size - 1) / 64 + 1 : 0;
-        void* buffer = MKQLAllocWithSize(sizeof(TSelectRangeLazyRow)
+        void* buffer = TBase::AllocWithSize(sizeof(TSelectRangeLazyRow)
             + size * sizeof(NUdf::TUnboxedValue)
             + maskSize * sizeof(ui64));
 
diff --git a/ydb/core/tx/datashard/datashard_ut_minikql.cpp b/ydb/core/tx/datashard/datashard_ut_minikql.cpp
index b788107f873..e7280b78f9a 100644
--- a/ydb/core/tx/datashard/datashard_ut_minikql.cpp
+++ b/ydb/core/tx/datashard/datashard_ut_minikql.cpp
@@ -1579,8 +1579,9 @@ Y_UNIT_TEST(MemoryUsageImmediateSmallTx) {
     }
 
     auto counters3 = ReadTxMemoryCounters(t, TTestTxConfig::TxTablet0);
-    // Expect one allocation on prepare.
-    CheckCounters(counters3, counters2, 1, 0, 0);
+    // Expect one allocation on prepare for default pool + one allocation for temporary pool after exception on limit extension unavailable.
+    // on execute_data_tx_unit.cpp:127 we set limit and on temporary object allocation (on reading) we add new page with exhausted limit.
+    CheckCounters(counters3, counters2, 2, 0, 0);
 }
 
 Y_UNIT_TEST(MemoryUsageImmediateMediumTx) {
diff --git a/ydb/library/yql/dq/runtime/dq_output_consumer.h b/ydb/library/yql/dq/runtime/dq_output_consumer.h
index 2b717bb4171..6aee7c6df0a 100644
--- a/ydb/library/yql/dq/runtime/dq_output_consumer.h
+++ b/ydb/library/yql/dq/runtime/dq_output_consumer.h
@@ -10,7 +10,8 @@ class TTypeEnvironment;
 
 namespace NYql::NDq {
 
-class IDqOutputConsumer : public TSimpleRefCount<IDqOutputConsumer>, public NKikimr::NMiniKQL::TWithMiniKQLAlloc {
+class IDqOutputConsumer : public TSimpleRefCount<IDqOutputConsumer>,
+    public NKikimr::NMiniKQL::TWithDefaultMiniKQLAlloc {
 public:
     using TPtr = TIntrusivePtr<IDqOutputConsumer>;
 
diff --git a/ydb/library/yql/minikql/computation/mkql_computation_list_adapter.h b/ydb/library/yql/minikql/computation/mkql_computation_list_adapter.h
index 95798ca383e..a237c143a06 100644
--- a/ydb/library/yql/minikql/computation/mkql_computation_list_adapter.h
+++ b/ydb/library/yql/minikql/computation/mkql_computation_list_adapter.h
@@ -14,10 +14,10 @@ public:
     typedef std::function<NUdf::TUnboxedValue(const TItem&)> TItemFactory;
     typedef TComputationValue<TVectorListAdapter<TVectorType>> TBase;
 
-    class TIterator: public TComputationValue<TIterator> {
+    class TIterator: public TTemporaryComputationValue<TIterator> {
     public:
         TIterator(TMemoryUsageInfo* memInfo, const TVectorType& list, TItemFactory itemFactory, ui64 start, ui64 finish, bool reversed)
-            : TComputationValue<TIterator>(memInfo)
+            : TTemporaryComputationValue<TIterator>(memInfo)
             , List(list)
             , ItemFactory(itemFactory)
             , Start(start)
@@ -57,10 +57,10 @@ public:
         ui64 Index;
     };
 
-    class TDictIterator : public TComputationValue<TDictIterator> {
+    class TDictIterator : public TTemporaryComputationValue<TDictIterator> {
     public:
         TDictIterator(TMemoryUsageInfo* memInfo, THolder<TIterator>&& iter)
-            : TComputationValue<TDictIterator>(memInfo)
+            : TTemporaryComputationValue<TDictIterator>(memInfo)
             , Iter(std::move(iter))
             , Index(Max<ui64>())
         {}
diff --git a/ydb/library/yql/minikql/computation/mkql_computation_node_holders.cpp b/ydb/library/yql/minikql/computation/mkql_computation_node_holders.cpp
index ff8be740309..dd51799d1c5 100644
--- a/ydb/library/yql/minikql/computation/mkql_computation_node_holders.cpp
+++ b/ydb/library/yql/minikql/computation/mkql_computation_node_holders.cpp
@@ -255,7 +255,7 @@ public:
     void* operator new[](size_t sz) = delete;
     void operator delete(void *mem, std::size_t sz) {
         const auto pSize = static_cast<void*>(static_cast<ui8*>(mem) + offsetof(TDirectArrayHolderInplace, Size));
-        MKQLFreeWithSize(mem, sz + *static_cast<ui32*>(pSize) * sizeof(NUdf::TUnboxedValue));
+        FreeWithSize(mem, sz + *static_cast<ui32*>(pSize) * sizeof(NUdf::TUnboxedValue));
     }
 
     void operator delete[](void *mem, std::size_t sz) = delete;
@@ -285,10 +285,10 @@ public:
     }
 
 private:
-    class TIterator : public TComputationValue<TIterator> {
+    class TIterator : public TTemporaryComputationValue<TIterator> {
     public:
         TIterator(const TDirectArrayHolderInplace* parent)
-            : TComputationValue(parent->GetMemInfo()), Parent(const_cast<TDirectArrayHolderInplace*>(parent))
+            : TTemporaryComputationValue(parent->GetMemInfo()), Parent(const_cast<TDirectArrayHolderInplace*>(parent))
         {}
 
     private:
@@ -314,10 +314,10 @@ private:
         ui32 Current = ~0U;
     };
 
-    class TKeysIterator : public TComputationValue<TKeysIterator> {
+    class TKeysIterator : public TTemporaryComputationValue<TKeysIterator> {
     public:
         TKeysIterator(const TDirectArrayHolderInplace& parent)
-            : TComputationValue(parent.GetMemInfo()), Size(parent.GetSize())
+            : TTemporaryComputationValue(parent.GetMemInfo()), Size(parent.GetSize())
         {}
     private:
         bool Skip() final {
@@ -1832,10 +1832,10 @@ private:
 class THashedDictHolder: public TComputationValue<THashedDictHolder> {
 public:
     template <bool NoSwap>
-    class TIterator: public TComputationValue<TIterator<NoSwap>> {
+    class TIterator: public TTemporaryComputationValue<TIterator<NoSwap>> {
     public:
         TIterator(const THashedDictHolder* parent)
-            : TComputationValue<TIterator<NoSwap>>(parent->GetMemInfo())
+            : TTemporaryComputationValue<TIterator<NoSwap>>(parent->GetMemInfo())
             , Parent(const_cast<THashedDictHolder*>(parent))
             , Iterator(Parent->Map.begin())
             , End(Parent->Map.end())
@@ -3140,7 +3140,8 @@ NUdf::TUnboxedValuePod THolderFactory::CreateDirectArrayHolder(ui32 size, NUdf::
         return GetEmptyContainer();
     }
 
-    const auto buffer = MKQLAllocFastWithSize(sizeof(TDirectArrayHolderInplace) + size * sizeof(NUdf::TUnboxedValue), CurrentAllocState);
+    const auto buffer = MKQLAllocFastWithSize(
+        sizeof(TDirectArrayHolderInplace) + size * sizeof(NUdf::TUnboxedValue), CurrentAllocState, EMemorySubPool::Default);
     const auto h = ::new(buffer) TDirectArrayHolderInplace(&MemInfo, size);
 
     auto res = NUdf::TUnboxedValuePod(h);
diff --git a/ydb/library/yql/minikql/computation/mkql_computation_node_impl.h b/ydb/library/yql/minikql/computation/mkql_computation_node_impl.h
index 0668bc5d809..9eddfed4252 100644
--- a/ydb/library/yql/minikql/computation/mkql_computation_node_impl.h
+++ b/ydb/library/yql/minikql/computation/mkql_computation_node_impl.h
@@ -789,20 +789,19 @@ protected:
     const EValueRepresentation Kind;
 };
 
-template <typename TDerived, typename TBase = NUdf::IBoxedValue>
-class TComputationValue: public TBase, public TWithMiniKQLAlloc
+template <typename TDerived, typename TBaseExt = NYql::NUdf::IBoxedValue>
+class TComputationValueBase: public TBaseExt
 {
+private:
+    using TBase = TBaseExt;
 public:
     template <typename... Args>
-    TComputationValue(TMemoryUsageInfo* memInfo, Args&&... args)
+    TComputationValueBase(Args&&... args)
         : TBase(std::forward<Args>(args)...)
     {
-        M_.MemInfo = memInfo;
-        MKQL_MEM_TAKE(memInfo, this, sizeof(TDerived), __MKQL_LOCATION__);
     }
 
-    ~TComputationValue() {
-        MKQL_MEM_RETURN(GetMemInfo(), this, sizeof(TDerived));
+    ~TComputationValueBase() {
     }
 
 private:
@@ -1020,20 +1019,55 @@ public:
     }
 
 protected:
+    [[noreturn]] void ThrowNotSupported(const char* func) const {
+        THROW yexception() << "Unsupported access to '" << func << "' method of: " << TypeName(*this);
+    }
+};
+
+template <typename TDerived, typename TBaseExt, EMemorySubPool MemoryPool>
+class TComputationValueImpl: public TComputationValueBase<TDerived, TBaseExt>,
+    public TWithMiniKQLAlloc<MemoryPool> {
+private:
+    using TBase = TComputationValueBase<TDerived, TBaseExt>;
+protected:
     inline TMemoryUsageInfo* GetMemInfo() const {
         return static_cast<TMemoryUsageInfo*>(M_.MemInfo);
     }
-
-    [[noreturn]] void ThrowNotSupported(const char* func) const {
-        THROW yexception() << "Unsupported access to '" << func << "' method of: " << TypeName(*this);
+    using TWithMiniKQLAlloc<MemoryPool>::AllocWithSize;
+    using TWithMiniKQLAlloc<MemoryPool>::FreeWithSize;
+public:
+    template <typename... Args>
+    TComputationValueImpl(TMemoryUsageInfo* memInfo, Args&&... args)
+        : TBase(std::forward<Args>(args)...) {
+        M_.MemInfo = memInfo;
+        MKQL_MEM_TAKE(memInfo, this, sizeof(TDerived), __MKQL_LOCATION__);
     }
 
+    ~TComputationValueImpl() {
+        MKQL_MEM_RETURN(GetMemInfo(), this, sizeof(TDerived));
+    }
 private:
     struct {
         void* MemInfo; // used for tracking memory usage during execution
     } M_;
 };
 
+template <typename TDerived, typename TBaseExt = NUdf::IBoxedValue>
+class TTemporaryComputationValue: public TComputationValueImpl<TDerived, TBaseExt, EMemorySubPool::Temporary> {
+private:
+    using TBase = TComputationValueImpl<TDerived, TBaseExt, EMemorySubPool::Temporary>;
+public:
+    using TBase::TBase;
+};
+
+template <typename TDerived, typename TBaseExt = NUdf::IBoxedValue>
+class TComputationValue: public TComputationValueImpl<TDerived, TBaseExt, EMemorySubPool::Default> {
+private:
+    using TBase = TComputationValueImpl<TDerived, TBaseExt, EMemorySubPool::Default>;
+public:
+    using TBase::TBase;
+};
+
 template<bool IsStream>
 struct TThresher;
 
diff --git a/ydb/library/yql/minikql/computation/mkql_computation_node_list.h b/ydb/library/yql/minikql/computation/mkql_computation_node_list.h
index d87fe871bf0..880098f59ba 100644
--- a/ydb/library/yql/minikql/computation/mkql_computation_node_list.h
+++ b/ydb/library/yql/minikql/computation/mkql_computation_node_list.h
@@ -27,7 +27,7 @@ namespace NKikimr {
             }
 
             static TSelf* AllocateChunk(ui64 length) {
-                const auto block = MKQLAllocWithSize(length * sizeof(T) + sizeof(TListChunk));
+                const auto block = TWithDefaultMiniKQLAlloc::AllocWithSize(length * sizeof(T) + sizeof(TListChunk));
                 const auto ptr = new (block) TListChunk(length);
                 PlacementNew(reinterpret_cast<T*>(ptr + 1), length);
                 return ptr;
@@ -45,7 +45,8 @@ namespace NKikimr {
                 for (auto it = DataBegin(); it != DataEnd(); it++) {
                     it->~T();
                 }
-                MKQLFreeWithSize(static_cast<void*>(this), sizeof(TListChunk) + sizeof(T) * (DataEnd() - DataBegin()));
+                TWithDefaultMiniKQLAlloc::FreeWithSize(
+                    static_cast<void*>(this), sizeof(TListChunk) + sizeof(T) * (DataEnd() - DataBegin()));
             }
 
             inline T* DataBegin() {
diff --git a/ydb/library/yql/minikql/mkql_alloc.cpp b/ydb/library/yql/minikql/mkql_alloc.cpp
index d2549a31e04..9c7d2646059 100644
--- a/ydb/library/yql/minikql/mkql_alloc.cpp
+++ b/ydb/library/yql/minikql/mkql_alloc.cpp
@@ -10,6 +10,7 @@ namespace NMiniKQL {
 Y_POD_THREAD(TAllocState*) TlsAllocState;
 
 TAllocPageHeader TAllocState::EmptyPageHeader = { 0, 0, 0, 0, nullptr, nullptr };
+TAllocState::TCurrentPages TAllocState::EmptyCurrentPages = { &TAllocState::EmptyPageHeader, &TAllocState::EmptyPageHeader };
 
 void TAllocState::TListEntry::Link(TAllocState::TListEntry* root) noexcept {
     Left = root;
@@ -115,7 +116,7 @@ void TScopedAlloc::Release() {
     }
 }
 
-void* MKQLAllocSlow(size_t sz, TAllocState* state) {
+void* MKQLAllocSlow(size_t sz, TAllocState* state, const EMemorySubPool mPool) {
     auto roundedSize = AlignUp(sz + sizeof(TAllocPageHeader), MKQL_ALIGNMENT);
     auto capacity = Max(ui64(TAlignedPagePool::POOL_PAGE_SIZE), roundedSize);
     auto currPage = (TAllocPageHeader*)state->GetBlock(capacity);
@@ -123,11 +124,12 @@ void* MKQLAllocSlow(size_t sz, TAllocState* state) {
     currPage->Capacity = capacity;
     currPage->Offset = roundedSize;
 
+    auto& mPage = state->CurrentPages[(TMemorySubPoolIdx)mPool];
     auto newPageAvailable = capacity - roundedSize;
-    auto curPageAvailable = state->CurrentPage->Capacity - state->CurrentPage->Offset;
+    auto curPageAvailable = mPage->Capacity - mPage->Offset;
 
     if (newPageAvailable > curPageAvailable) {
-        state->CurrentPage = currPage;
+        mPage = currPage;
     }
 
     void* ret = (char*)currPage + sizeof(TAllocPageHeader);
@@ -137,39 +139,42 @@ void* MKQLAllocSlow(size_t sz, TAllocState* state) {
     return ret;
 }
 
-void MKQLFreeSlow(TAllocPageHeader* header, TAllocState *state) noexcept {
+void MKQLFreeSlow(TAllocPageHeader* header, TAllocState *state, const EMemorySubPool mPool) noexcept {
     Y_VERIFY_DEBUG(state);
     Y_VERIFY_DEBUG(header->MyAlloc == state, "%s", (TStringBuilder() << "wrong allocator was used; "
         "allocated with: " << header->MyAlloc->GetInfo() << " freed with: " << TlsAllocState->GetInfo()).data());
     state->ReturnBlock(header, header->Capacity);
-    if (header == state->CurrentPage) {
-        state->CurrentPage = &TAllocState::EmptyPageHeader;
+    if (header == state->CurrentPages[(TMemorySubPoolIdx)mPool]) {
+        state->CurrentPages[(TMemorySubPoolIdx)mPool] = &TAllocState::EmptyPageHeader;
     }
 }
 
-void* TPagedArena::AllocSlow(size_t sz) {
-    auto prevLink = CurrentPage_;
+void* TPagedArena::AllocSlow(const size_t sz, const EMemorySubPool mPool) {
+    auto& currentPage = CurrentPages_[(TMemorySubPoolIdx)mPool];
+    auto prevLink = currentPage;
     auto roundedSize = AlignUp(sz + sizeof(TAllocPageHeader), MKQL_ALIGNMENT);
     auto capacity = Max(ui64(TAlignedPagePool::POOL_PAGE_SIZE), roundedSize);
-    CurrentPage_ = (TAllocPageHeader*)PagePool_->GetBlock(capacity);
-    CurrentPage_->Capacity = capacity;
-    void* ret = (char*)CurrentPage_ + sizeof(TAllocPageHeader);
-    CurrentPage_->Offset = roundedSize;
-    CurrentPage_->UseCount = 0;
-    CurrentPage_->MyAlloc = PagePool_;
-    CurrentPage_->Link = prevLink;
+    currentPage = (TAllocPageHeader*)PagePool_->GetBlock(capacity);
+    currentPage->Capacity = capacity;
+    void* ret = (char*)currentPage + sizeof(TAllocPageHeader);
+    currentPage->Offset = roundedSize;
+    currentPage->UseCount = 0;
+    currentPage->MyAlloc = PagePool_;
+    currentPage->Link = prevLink;
     return ret;
 }
 
 void TPagedArena::Clear() noexcept {
-    auto current = CurrentPage_;
-    while (current != &TAllocState::EmptyPageHeader) {
-        auto next = current->Link;
-        PagePool_->ReturnBlock(current, current->Capacity);
-        current = next;
-    }
+    for (auto&& i : CurrentPages_) {
+        auto current = i;
+        while (current != &TAllocState::EmptyPageHeader) {
+            auto next = current->Link;
+            PagePool_->ReturnBlock(current, current->Capacity);
+            current = next;
+        }
 
-    CurrentPage_ = &TAllocState::EmptyPageHeader;
+        i = &TAllocState::EmptyPageHeader;
+    }
 }
 
 } // NMiniKQL
diff --git a/ydb/library/yql/minikql/mkql_alloc.h b/ydb/library/yql/minikql/mkql_alloc.h
index a1cdc735f4b..eaa9ee0337a 100644
--- a/ydb/library/yql/minikql/mkql_alloc.h
+++ b/ydb/library/yql/minikql/mkql_alloc.h
@@ -28,6 +28,14 @@ struct TAllocPageHeader {
     TAllocPageHeader* Link;
 };
 
+using TMemorySubPoolIdx = ui32;
+enum class EMemorySubPool: TMemorySubPoolIdx {
+    Default = 0,
+    Temporary = 1,
+
+    Count
+};
+
 constexpr ui32 MaxPageUserData = TAlignedPagePool::POOL_PAGE_SIZE - sizeof(TAllocPageHeader);
 
 static_assert(sizeof(TAllocPageHeader) % MKQL_ALIGNMENT == 0, "Incorrect size of header");
@@ -56,8 +64,12 @@ struct TAllocState : public TAlignedPagePool
         Free(ptr, size);
     }
 
+    using TCurrentPages = std::array<TAllocPageHeader*, (TMemorySubPoolIdx)EMemorySubPool::Count>;
+
     static TAllocPageHeader EmptyPageHeader;
-    TAllocPageHeader* CurrentPage = &EmptyPageHeader;
+    static TCurrentPages EmptyCurrentPages;
+
+    std::array<TAllocPageHeader*, (TMemorySubPoolIdx)EMemorySubPool::Count> CurrentPages = EmptyCurrentPages;
     TListEntry OffloadedBlocksRoot;
     TListEntry GlobalPAllocList;
     TListEntry* CurrentPAllocList;
@@ -164,51 +176,52 @@ class TPagedArena {
 public:
     TPagedArena(TAlignedPagePool* pagePool) noexcept
         : PagePool_(pagePool)
-        , CurrentPage_(&TAllocState::EmptyPageHeader)
+        , CurrentPages_(TAllocState::EmptyCurrentPages)
     {}
 
     TPagedArena(const TPagedArena&) = delete;
     TPagedArena(TPagedArena&& other) noexcept
         : PagePool_(other.PagePool_)
-        , CurrentPage_(other.CurrentPage_)
+        , CurrentPages_(other.CurrentPages_)
     {
-        other.CurrentPage_ = &TAllocState::EmptyPageHeader;
+        other.CurrentPages_ = TAllocState::EmptyCurrentPages;
     }
 
     void operator=(const TPagedArena&) = delete;
     void operator=(TPagedArena&& other) noexcept {
         Clear();
         PagePool_ = other.PagePool_;
-        CurrentPage_ = other.CurrentPage_;
-        other.CurrentPage_ = &TAllocState::EmptyPageHeader;
+        CurrentPages_ = other.CurrentPages_;
+        other.CurrentPages_ = TAllocState::EmptyCurrentPages;
     }
 
     ~TPagedArena() noexcept {
         Clear();
     }
 
-    void* Alloc(size_t sz) {
-        if (Y_LIKELY(CurrentPage_->Offset + sz <= CurrentPage_->Capacity)) {
-            void* ret = (char*)CurrentPage_ + CurrentPage_->Offset;
-            CurrentPage_->Offset = AlignUp(CurrentPage_->Offset + sz, MKQL_ALIGNMENT);
+    void* Alloc(size_t sz, const EMemorySubPool pagePool = EMemorySubPool::Default) {
+        auto& currentPage = CurrentPages_[(TMemorySubPoolIdx)pagePool];
+        if (Y_LIKELY(currentPage->Offset + sz <= currentPage->Capacity)) {
+            void* ret = (char*)currentPage + currentPage->Offset;
+            currentPage->Offset = AlignUp(currentPage->Offset + sz, MKQL_ALIGNMENT);
             return ret;
         }
 
-        return AllocSlow(sz);
+        return AllocSlow(sz, pagePool);
     }
 
     void Clear() noexcept;
 
 private:
-    void* AllocSlow(size_t sz);
+    void* AllocSlow(const size_t sz, const EMemorySubPool pagePool);
 
 private:
     TAlignedPagePool* PagePool_;
-    TAllocPageHeader* CurrentPage_ = &TAllocState::EmptyPageHeader;
+    TAllocState::TCurrentPages CurrentPages_ = TAllocState::EmptyCurrentPages;
 };
 
-void* MKQLAllocSlow(size_t sz, TAllocState* state);
-inline void* MKQLAllocFastDeprecated(size_t sz, TAllocState* state) {
+void* MKQLAllocSlow(size_t sz, TAllocState* state, const EMemorySubPool mPool);
+inline void* MKQLAllocFastDeprecated(size_t sz, TAllocState* state, const EMemorySubPool mPool) {
     Y_VERIFY_DEBUG(state);
 
 #ifdef PROFILE_MEMORY_ALLOCATIONS
@@ -221,7 +234,7 @@ inline void* MKQLAllocFastDeprecated(size_t sz, TAllocState* state) {
     return ret + 1;
 #endif
 
-    auto currPage = state->CurrentPage;
+    auto currPage = state->CurrentPages[(TMemorySubPoolIdx)mPool];
     if (Y_LIKELY(currPage->Offset + sz <= currPage->Capacity)) {
         void* ret = (char*)currPage + currPage->Offset;
         currPage->Offset = AlignUp(currPage->Offset + sz, MKQL_ALIGNMENT);
@@ -229,10 +242,10 @@ inline void* MKQLAllocFastDeprecated(size_t sz, TAllocState* state) {
         return ret;
     }
 
-    return MKQLAllocSlow(sz, state);
+    return MKQLAllocSlow(sz, state, mPool);
 }
 
-inline void* MKQLAllocFastWithSize(size_t sz, TAllocState* state) {
+inline void* MKQLAllocFastWithSize(size_t sz, TAllocState* state, const EMemorySubPool mPool) {
     Y_VERIFY_DEBUG(state);
 
     bool useMemalloc = state->SupportsSizedAllocators && sz > MaxPageUserData;
@@ -252,7 +265,7 @@ inline void* MKQLAllocFastWithSize(size_t sz, TAllocState* state) {
         return ret + 1;
     }
 
-    auto currPage = state->CurrentPage;
+    auto currPage = state->CurrentPages[(TMemorySubPoolIdx)mPool];
     if (Y_LIKELY(currPage->Offset + sz <= currPage->Capacity)) {
         void* ret = (char*)currPage + currPage->Offset;
         currPage->Offset = AlignUp(currPage->Offset + sz, MKQL_ALIGNMENT);
@@ -260,12 +273,12 @@ inline void* MKQLAllocFastWithSize(size_t sz, TAllocState* state) {
         return ret;
     }
 
-    return MKQLAllocSlow(sz, state);
+    return MKQLAllocSlow(sz, state, mPool);
 }
 
-void MKQLFreeSlow(TAllocPageHeader* header, TAllocState *state) noexcept;
+void MKQLFreeSlow(TAllocPageHeader* header, TAllocState *state, const EMemorySubPool mPool) noexcept;
 
-inline void MKQLFreeDeprecated(const void* mem) noexcept {
+inline void MKQLFreeDeprecated(const void* mem, const EMemorySubPool mPool = EMemorySubPool::Default) noexcept {
     if (!mem) {
         return;
     }
@@ -287,10 +300,10 @@ inline void MKQLFreeDeprecated(const void* mem) noexcept {
         return;
     }
 
-    MKQLFreeSlow(header, TlsAllocState);
+    MKQLFreeSlow(header, TlsAllocState, mPool);
 }
 
-inline void MKQLFreeFastWithSize(const void* mem, size_t sz, TAllocState* state) noexcept {
+inline void MKQLFreeFastWithSize(const void* mem, size_t sz, TAllocState* state, const EMemorySubPool mPool) noexcept {
     if (!mem) {
         return;
     }
@@ -319,19 +332,19 @@ inline void MKQLFreeFastWithSize(const void* mem, size_t sz, TAllocState* state)
         return;
     }
 
-    MKQLFreeSlow(header, state);
+    MKQLFreeSlow(header, state, mPool);
 }
 
-inline void* MKQLAllocDeprecated(size_t sz) {
-    return MKQLAllocFastDeprecated(sz, TlsAllocState);
+inline void* MKQLAllocDeprecated(size_t sz, const EMemorySubPool mPool = EMemorySubPool::Default) {
+    return MKQLAllocFastDeprecated(sz, TlsAllocState, mPool);
 }
 
-inline void* MKQLAllocWithSize(size_t sz) {
-    return MKQLAllocFastWithSize(sz, TlsAllocState);
+inline void* MKQLAllocWithSize(size_t sz, const EMemorySubPool mPool) {
+    return MKQLAllocFastWithSize(sz, TlsAllocState, mPool);
 }
 
-inline void MKQLFreeWithSize(const void* mem, size_t sz) noexcept {
-    return MKQLFreeFastWithSize(mem, sz, TlsAllocState);
+inline void MKQLFreeWithSize(const void* mem, size_t sz, const EMemorySubPool mPool) noexcept {
+    return MKQLFreeFastWithSize(mem, sz, TlsAllocState, mPool);
 }
 
 inline void MKQLRegisterObject(NUdf::TBoxedValue* value) noexcept {
@@ -342,33 +355,44 @@ inline void MKQLUnregisterObject(NUdf::TBoxedValue* value) noexcept {
     return value->Unlink();
 }
 
+template <const EMemorySubPool MemoryPoolExt = EMemorySubPool::Default>
 struct TWithMiniKQLAlloc {
+    static constexpr EMemorySubPool MemoryPool = MemoryPoolExt;
+
+    static void FreeWithSize(const void* mem, const size_t sz) {
+        NMiniKQL::MKQLFreeWithSize(mem, sz, MemoryPool);
+    }
+
+    static void* AllocWithSize(const size_t sz) {
+        return NMiniKQL::MKQLAllocWithSize(sz, MemoryPool);
+    }
+
     void* operator new(size_t sz) {
-        return MKQLAllocWithSize(sz);
+        return NMiniKQL::MKQLAllocWithSize(sz, MemoryPool);
     }
 
     void* operator new[](size_t sz) {
-        return MKQLAllocWithSize(sz);
+        return NMiniKQL::MKQLAllocWithSize(sz, MemoryPool);
     }
 
     void operator delete(void *mem, std::size_t sz) noexcept {
-        MKQLFreeWithSize(mem, sz);
+        NMiniKQL::MKQLFreeWithSize(mem, sz, MemoryPool);
     }
 
     void operator delete[](void *mem, std::size_t sz) noexcept {
-        MKQLFreeWithSize(mem, sz);
+        NMiniKQL::MKQLFreeWithSize(mem, sz, MemoryPool);
     }
 };
 
 template <typename T, typename... Args>
 T* AllocateOn(TAllocState* state, Args&&... args)
 {
-    void* addr = MKQLAllocFastWithSize(sizeof(T), state);
+    void* addr = MKQLAllocFastWithSize(sizeof(T), state, T::MemoryPool);
     return ::new(addr) T(std::forward<Args>(args)...);
-    static_assert(std::is_base_of<TWithMiniKQLAlloc, T>::value, "Class must inherit TWithMiniKQLAlloc.");
+    static_assert(std::is_base_of<TWithMiniKQLAlloc<T::MemoryPool>, T>::value, "Class must inherit TWithMiniKQLAlloc.");
 }
 
-template <typename Type>
+template <typename Type, enum EMemorySubPool MemoryPool = EMemorySubPool::Default>
 struct TMKQLAllocator
 {
     typedef Type value_type;
@@ -389,15 +413,18 @@ struct TMKQLAllocator
 
     static pointer allocate(size_type n, const void* = nullptr)
     {
-        return static_cast<pointer>(MKQLAllocWithSize(n * sizeof(value_type)));
+        return static_cast<pointer>(MKQLAllocWithSize(n * sizeof(value_type), MemoryPool));
     }
 
     static void deallocate(const_pointer p, size_type n) noexcept
     {
-        return MKQLFreeWithSize(p, n * sizeof(value_type));
+        return MKQLFreeWithSize(p, n * sizeof(value_type), MemoryPool);
     }
 };
 
+using TWithDefaultMiniKQLAlloc = TWithMiniKQLAlloc<EMemorySubPool::Default>;
+using TWithTemporaryMiniKQLAlloc = TWithMiniKQLAlloc<EMemorySubPool::Temporary>;
+
 template <typename T>
 class TPagedList
 {
diff --git a/ydb/library/yql/minikql/mkql_alloc_ut.cpp b/ydb/library/yql/minikql/mkql_alloc_ut.cpp
index 2b4cc8f17ba..fa3d9060d2c 100644
--- a/ydb/library/yql/minikql/mkql_alloc_ut.cpp
+++ b/ydb/library/yql/minikql/mkql_alloc_ut.cpp
@@ -31,24 +31,24 @@ Y_UNIT_TEST_SUITE(TMiniKQLAllocTest) {
 
     Y_UNIT_TEST(TestDeallocated) {
         TScopedAlloc alloc(__LOCATION__);
-        void* p1 = MKQLAllocWithSize(10);
-        void* p2 = MKQLAllocWithSize(20);
+        void* p1 = TWithDefaultMiniKQLAlloc::AllocWithSize(10);
+        void* p2 = TWithDefaultMiniKQLAlloc::AllocWithSize(20);
         UNIT_ASSERT_VALUES_EQUAL(alloc.Ref().GetUsed(), TAlignedPagePool::POOL_PAGE_SIZE);
         UNIT_ASSERT_VALUES_EQUAL(alloc.Ref().GetDeallocatedInPages(), 0);
         UNIT_ASSERT_VALUES_EQUAL(alloc.Ref().GetFreePageCount(), 0);
-        MKQLFreeWithSize(p1, 10);
+        TWithDefaultMiniKQLAlloc::FreeWithSize(p1, 10);
         UNIT_ASSERT_VALUES_EQUAL(alloc.Ref().GetUsed(), TAlignedPagePool::POOL_PAGE_SIZE);
         UNIT_ASSERT_VALUES_EQUAL(alloc.Ref().GetDeallocatedInPages(), 10);
         UNIT_ASSERT_VALUES_EQUAL(alloc.Ref().GetFreePageCount(), 0);
-        MKQLFreeWithSize(p2, 20);
+        TWithDefaultMiniKQLAlloc::FreeWithSize(p2, 20);
         UNIT_ASSERT_VALUES_EQUAL(alloc.Ref().GetUsed(), 0);
         UNIT_ASSERT_VALUES_EQUAL(alloc.Ref().GetDeallocatedInPages(), 0);
         UNIT_ASSERT_VALUES_EQUAL(alloc.Ref().GetFreePageCount(), 1);
-        p1 = MKQLAllocWithSize(10);
+        p1 = TWithDefaultMiniKQLAlloc::AllocWithSize(10);
         UNIT_ASSERT_VALUES_EQUAL(alloc.Ref().GetUsed(), TAlignedPagePool::POOL_PAGE_SIZE);
         UNIT_ASSERT_VALUES_EQUAL(alloc.Ref().GetDeallocatedInPages(), 0);
         UNIT_ASSERT_VALUES_EQUAL(alloc.Ref().GetFreePageCount(), 0);
-        MKQLFreeWithSize(p1, 10);
+        TWithDefaultMiniKQLAlloc::FreeWithSize(p1, 10);
         UNIT_ASSERT_VALUES_EQUAL(alloc.Ref().GetUsed(), 0);
         UNIT_ASSERT_VALUES_EQUAL(alloc.Ref().GetDeallocatedInPages(), 0);
         UNIT_ASSERT_VALUES_EQUAL(alloc.Ref().GetFreePageCount(), 1);
@@ -59,13 +59,13 @@ Y_UNIT_TEST_SUITE(TMiniKQLAllocTest) {
             return;
         }
         TScopedAlloc alloc1(__LOCATION__);
-        void* p1 = MKQLAllocWithSize(10);
-        void* p2 = MKQLAllocWithSize(10);
+        void* p1 = TWithDefaultMiniKQLAlloc::AllocWithSize(10);
+        void* p2 = TWithDefaultMiniKQLAlloc::AllocWithSize(10);
         {
             TScopedAlloc alloc2(__LOCATION__);
-            MKQLFreeWithSize(p1, 10);
+            TWithDefaultMiniKQLAlloc::FreeWithSize(p1, 10);
         }
-        MKQLFreeWithSize(p2, 10);
+        TWithDefaultMiniKQLAlloc::FreeWithSize(p2, 10);
     }
 }
 
diff --git a/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp b/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
index 468da595f11..f90e01991a0 100644
--- a/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
+++ b/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
@@ -374,7 +374,7 @@ public:
         }
 
         MemSize = SizeForFunctionCallInfo(numArgs);
-        Ptr = MKQLAllocWithSize(MemSize);
+        Ptr = TWithDefaultMiniKQLAlloc::AllocWithSize(MemSize);
         auto& callInfo = Ref();
         Zero(callInfo);
         callInfo.flinfo = &CopyFmgrInfo; // client may mutate fn_extra
@@ -389,7 +389,7 @@ public:
 
     ~TFunctionCallInfo() {
         if (Ptr) {
-            MKQLFreeWithSize(Ptr, MemSize);
+            TWithDefaultMiniKQLAlloc::FreeWithSize(Ptr, MemSize);
         }
     }
 
@@ -406,13 +406,13 @@ private:
 class TReturnSetInfo {
 public:
     TReturnSetInfo() {
-        Ptr = MKQLAllocWithSize(sizeof(ReturnSetInfo));
+        Ptr = TWithDefaultMiniKQLAlloc::AllocWithSize(sizeof(ReturnSetInfo));
         Zero(Ref());
         Ref().type = T_ReturnSetInfo;
     }
 
     ~TReturnSetInfo() {
-        MKQLFreeWithSize(Ptr, sizeof(ReturnSetInfo));
+        TWithDefaultMiniKQLAlloc::FreeWithSize(Ptr, sizeof(ReturnSetInfo));
     }
 
     ReturnSetInfo& Ref() {
@@ -863,14 +863,14 @@ public:
             auto lb = ARR_LBOUND(arr);
             auto nitems = ArrayGetNItems(ndim, dims);
 
-            Datum* elems = (Datum*)MKQLAllocWithSize(nitems * sizeof(Datum));
+            Datum* elems = (Datum*)TWithDefaultMiniKQLAlloc::AllocWithSize(nitems * sizeof(Datum));
             Y_DEFER {
-                MKQLFreeWithSize(elems, nitems * sizeof(Datum));
+                TWithDefaultMiniKQLAlloc::FreeWithSize(elems, nitems * sizeof(Datum));
             };
 
-            bool* nulls = (bool*)MKQLAllocWithSize(nitems);
+            bool* nulls = (bool*)TWithDefaultMiniKQLAlloc::AllocWithSize(nitems);
             Y_DEFER {
-                MKQLFreeWithSize(nulls, nitems);
+                TWithDefaultMiniKQLAlloc::FreeWithSize(nulls, nitems);
             };
 
             array_iter iter;
@@ -1271,14 +1271,14 @@ public:
             args.push_back(value);
         }
 
-        Datum* dvalues = (Datum*)MKQLAllocWithSize(nelems * sizeof(Datum));
+        Datum* dvalues = (Datum*)TWithDefaultMiniKQLAlloc::AllocWithSize(nelems * sizeof(Datum));
         Y_DEFER {
-            MKQLFreeWithSize(dvalues, nelems * sizeof(Datum));
+            TWithDefaultMiniKQLAlloc::FreeWithSize(dvalues, nelems * sizeof(Datum));
         };
 
-        bool *dnulls = (bool*)MKQLAllocWithSize(nelems);
+        bool *dnulls = (bool*)TWithDefaultMiniKQLAlloc::AllocWithSize(nelems);
         Y_DEFER {
-            MKQLFreeWithSize(dnulls, nelems);
+            TWithDefaultMiniKQLAlloc::FreeWithSize(dnulls, nelems);
         };
 
         TPAllocScope call;
@@ -2163,9 +2163,9 @@ NUdf::TUnboxedValue ReadSkiffPg(TPgType* type, NCommon::TInputBuf& buf) {
         ui32 size;
         buf.ReadMany((char*)&size, sizeof(size));
         CHECK_STRING_LENGTH_UNSIGNED(size);
-        char* s = (char*)MKQLAllocWithSize(size);
+        char* s = (char*)TWithDefaultMiniKQLAlloc::AllocWithSize(size);
         Y_DEFER {
-            MKQLFreeWithSize(s, size);
+            TWithDefaultMiniKQLAlloc::FreeWithSize(s, size);
         };
 
         buf.ReadMany(s, size);
@@ -2570,15 +2570,15 @@ NUdf::TUnboxedValue DecodePresortPGValue(TPgType* type, TStringBuf& input, TVect
 
 void* PgInitializeContext(const std::string_view& contextType) {
     if (contextType == "Agg") {
-        auto ctx = (AggState*)MKQLAllocWithSize(sizeof(AggState));
+        auto ctx = (AggState*)TWithDefaultMiniKQLAlloc::AllocWithSize(sizeof(AggState));
         Zero(*ctx);
         *(NodeTag*)ctx = T_AggState;
-        ctx->curaggcontext = (ExprContext*)MKQLAllocWithSize(sizeof(ExprContext));
+        ctx->curaggcontext = (ExprContext*)TWithDefaultMiniKQLAlloc::AllocWithSize(sizeof(ExprContext));
         Zero(*ctx->curaggcontext);
         ctx->curaggcontext->ecxt_per_tuple_memory = (MemoryContext)&((TMainContext*)TlsAllocState->MainContext)->Data;
         return ctx;
     } else if (contextType == "WinAgg") {
-        auto ctx = (WindowAggState*)MKQLAllocWithSize(sizeof(WindowAggState));
+        auto ctx = (WindowAggState*)TWithDefaultMiniKQLAlloc::AllocWithSize(sizeof(WindowAggState));
         Zero(*ctx);
         *(NodeTag*)ctx = T_WindowAggState;
         ctx->curaggcontext = (MemoryContext)&((TMainContext*)TlsAllocState->MainContext)->Data;
@@ -2590,10 +2590,10 @@ void* PgInitializeContext(const std::string_view& contextType) {
 
 void PgDestroyContext(const std::string_view& contextType, void* ctx) {
     if (contextType == "Agg") {
-        MKQLFreeWithSize(((AggState*)ctx)->curaggcontext, sizeof(ExprContext));
-        MKQLFreeWithSize(ctx, sizeof(AggState));
+        TWithDefaultMiniKQLAlloc::FreeWithSize(((AggState*)ctx)->curaggcontext, sizeof(ExprContext));
+        TWithDefaultMiniKQLAlloc::FreeWithSize(ctx, sizeof(AggState));
     } else if (contextType == "WinAgg") {
-        MKQLFreeWithSize(ctx, sizeof(WindowAggState));
+        TWithDefaultMiniKQLAlloc::FreeWithSize(ctx, sizeof(WindowAggState));
     } else {
         Y_FAIL("Unsupported context type");
     }
diff --git a/ydb/library/yql/public/udf/service/exception_policy/udf_service.cpp b/ydb/library/yql/public/udf/service/exception_policy/udf_service.cpp
index ff3c42d2aee..f4c5930b58e 100644
--- a/ydb/library/yql/public/udf/service/exception_policy/udf_service.cpp
+++ b/ydb/library/yql/public/udf/service/exception_policy/udf_service.cpp
@@ -23,9 +23,9 @@ extern "C" void UdfUnregisterObject(::NYql::NUdf::TBoxedValue* object) {
 }
 
 extern "C" void* UdfAllocateWithSize(ui64 size) {
-    return ::NKikimr::NMiniKQL::MKQLAllocWithSize(size);
+    return ::NKikimr::NMiniKQL::TWithDefaultMiniKQLAlloc::AllocWithSize(size);
 }
 
 extern "C" void UdfFreeWithSize(const void* mem, ui64 size) {
-    return ::NKikimr::NMiniKQL::MKQLFreeWithSize(mem, size);
+    return ::NKikimr::NMiniKQL::TWithDefaultMiniKQLAlloc::FreeWithSize(mem, size);
 }
diff --git a/ydb/library/yql/public/udf/service/terminate_policy/udf_service.cpp b/ydb/library/yql/public/udf/service/terminate_policy/udf_service.cpp
index a16809b21f3..e795a8f04a9 100644
--- a/ydb/library/yql/public/udf/service/terminate_policy/udf_service.cpp
+++ b/ydb/library/yql/public/udf/service/terminate_policy/udf_service.cpp
@@ -19,9 +19,9 @@ extern "C" void UdfRegisterObject(::NYql::NUdf::TBoxedValue*) {}
 extern "C" void UdfUnregisterObject(::NYql::NUdf::TBoxedValue*) {}
 
 extern "C" void* UdfAllocateWithSize(ui64 size) {
-    return ::NKikimr::NMiniKQL::MKQLAllocWithSize(size);
+    return ::NKikimr::NMiniKQL::TWithDefaultMiniKQLAlloc::AllocWithSize(size);
 }
 
 extern "C" void UdfFreeWithSize(const void* mem, ui64 size) {
-    return ::NKikimr::NMiniKQL::MKQLFreeWithSize(mem, size);
+    return ::NKikimr::NMiniKQL::TWithDefaultMiniKQLAlloc::FreeWithSize(mem, size);
 }