[mkql] Use type memoization in compiler

YQL-19355
commit_hash:ff1684e65f8ec72be0a1407be5d31bb3fb0465a7

3 files changed, 6 insertions, 10 deletions

diff --git a/yql/essentials/core/arrow_kernels/request/request.cpp b/yql/essentials/core/arrow_kernels/request/request.cpp
index 6249012269..b51adda710 100644
--- a/yql/essentials/core/arrow_kernels/request/request.cpp
+++ b/yql/essentials/core/arrow_kernels/request/request.cpp
@@ -1,5 +1,4 @@
 #include "request.h"
-#include <yql/essentials/providers/common/mkql/yql_type_mkql.h>
 #include <yql/essentials/providers/common/mkql/yql_provider_mkql.h>
 #include <yql/essentials/minikql/mkql_node_cast.h>
 #include <yql/essentials/minikql/mkql_node_serialization.h>
@@ -246,18 +245,13 @@ TRuntimeNode TKernelRequestBuilder::MakeArg(const TTypeAnnotationNode* type) {
 }
 
 TBlockType* TKernelRequestBuilder::MakeType(const TTypeAnnotationNode* type) {
-    auto [it, inserted] = CachedTypes_.emplace(type, nullptr);
-    if (!inserted) {
-        return it->second;
-    }
-
     TStringStream err;
-    const auto ret = NCommon::BuildType(*type, Pb_, err);
+    const auto ret = NCommon::BuildType(*type, Pb_, TypesMemoization_, err);
     if (!ret) {
         ythrow yexception() << err.Str();
     }
 
-    return it->second = AS_TYPE(TBlockType, ret);
+    return AS_TYPE(TBlockType, ret);
 }
 
 }
diff --git a/yql/essentials/core/arrow_kernels/request/request.h b/yql/essentials/core/arrow_kernels/request/request.h
index bfed8d5dd7..9dad9b6830 100644
--- a/yql/essentials/core/arrow_kernels/request/request.h
+++ b/yql/essentials/core/arrow_kernels/request/request.h
@@ -1,6 +1,7 @@
 #pragma once
 #include <yql/essentials/ast/yql_expr.h>
 #include <yql/essentials/minikql/mkql_program_builder.h>
+#include <yql/essentials/providers/common/mkql/yql_type_mkql.h>
 
 #include <unordered_map>
 
@@ -63,7 +64,7 @@ private:
     NKikimr::NMiniKQL::TProgramBuilder Pb_;
     std::vector<NKikimr::NMiniKQL::TRuntimeNode> Items_;
     std::vector<NKikimr::NMiniKQL::TRuntimeNode> ArgsItems_;
-    std::unordered_map<const TTypeAnnotationNode*, NKikimr::NMiniKQL::TBlockType*> CachedTypes_;
+    NCommon::TMemoizedTypesMap TypesMemoization_;
     std::unordered_map<const TTypeAnnotationNode*, NKikimr::NMiniKQL::TRuntimeNode> CachedArgs_;
 };
 
diff --git a/yql/essentials/core/services/yql_eval_params.cpp b/yql/essentials/core/services/yql_eval_params.cpp
index 8002becd8e..23b716d0ec 100644
--- a/yql/essentials/core/services/yql_eval_params.cpp
+++ b/yql/essentials/core/services/yql_eval_params.cpp
@@ -27,6 +27,7 @@ bool BuildParameterValuesAsNodes(const THashMap<TStringBuf, const TTypeAnnotatio
     TTypeEnvironment env(alloc);
     TMemoryUsageInfo memInfo("Parameters");
     THolderFactory holderFactory(alloc.Ref(), memInfo);
+    NCommon::TMemoizedTypesMap typesMemoization;
     bool isOk = true;
     auto& paramDataMap = paramData.AsMap();
     for (auto& p : paramTypes) {
@@ -34,7 +35,7 @@ bool BuildParameterValuesAsNodes(const THashMap<TStringBuf, const TTypeAnnotatio
 
         TStringStream err;
         TProgramBuilder pgmBuilder(env, functionRegistry);
-        TType* mkqlType = NCommon::BuildType(*p.second, pgmBuilder, err);
+        TType* mkqlType = NCommon::BuildType(*p.second, pgmBuilder, typesMemoization, err);
         if (!mkqlType) {
             ctx.AddError(TIssue({}, TStringBuilder() << "Failed to process type for parameter: " << name << ", reason: " << err.Str()));
             isOk = false;