avoid allocating data on TTypeEnv when making cell vectors (#12552)

3 files changed, 65 insertions, 4 deletions

diff --git a/ydb/core/engine/mkql_keys.cpp b/ydb/core/engine/mkql_keys.cpp
index 83183956df..d92ec03bc8 100644
--- a/ydb/core/engine/mkql_keys.cpp
+++ b/ydb/core/engine/mkql_keys.cpp
@@ -251,8 +251,10 @@ THolder<TKeyDesc> ExtractEraseRow(TCallable& callable, const TTypeEnvironment& e
 #define MAKE_PRIMITIVE_TYPE_CELL(type, layout) \
     case NUdf::TDataType<type>::Id: return MakeCell<layout>(value);
 
-TCell MakeCell(NScheme::TTypeInfo type, const NUdf::TUnboxedValuePod& value,
-    const TTypeEnvironment& env, bool copy,
+
+template<typename TStringBackend>
+TCell MakeCellImpl(NScheme::TTypeInfo type, const NUdf::TUnboxedValuePod& value,
+    const TStringBackend& env, bool copy,
     i32 typmod, TMaybe<TString>* error)
 {
     if (!value)
@@ -300,6 +302,21 @@ TCell MakeCell(NScheme::TTypeInfo type, const NUdf::TUnboxedValuePod& value,
     return TCell(val.Data(), val.Size());
 }
 
+TCell MakeCell(NScheme::TTypeInfo type, const NUdf::TUnboxedValuePod& value,
+    const TTypeEnvironment& env, bool copy,
+    i32 typmod, TMaybe<TString>* error)
+{
+    return MakeCellImpl(type, value, env, copy, typmod, error);
+}
+
+TCell MakeCell(NScheme::TTypeInfo type, const NUdf::TUnboxedValuePod& value,
+    const TStringProviderBackend& env, bool copy,
+    i32 typmod, TMaybe<TString>* error)
+{
+    return MakeCellImpl(type, value, env, copy, typmod, error);
+}
+
+
 #undef MAKE_PRIMITIVE_TYPE_CELL
 
 TReadTarget ExtractFlatReadTarget(TRuntimeNode modeInput) {
diff --git a/ydb/core/engine/mkql_keys.h b/ydb/core/engine/mkql_keys.h
index 60c5c4b1e5..d2a446e970 100644
--- a/ydb/core/engine/mkql_keys.h
+++ b/ydb/core/engine/mkql_keys.h
@@ -52,10 +52,48 @@ TCell MakeCell(const NUdf::TUnboxedValuePod& value) {
     return TCell(reinterpret_cast<const char*>(&v), sizeof(v));
 }
 
+struct TStringProviderBackend {
+    mutable TMemoryPool MemoryPool;
+
+    TStringProviderBackend()
+        : MemoryPool(256)
+    {}
+
+    class TMutableStringData {
+        friend struct TStringProviderBackend;
+
+    private:
+        char* Data_;
+        size_t Size_ = 0;
+
+        explicit TMutableStringData(char* data, size_t size)
+            : Data_(data)
+            , Size_(size)
+        {}
+
+    public:
+        char *Data() const noexcept {
+            return Data_;
+        }
+
+        size_t Size() const noexcept {
+            return Size_;
+        }
+    };
+
+    TMutableStringData NewString(ui32 size) const {
+        return TMutableStringData(reinterpret_cast<char*>(MemoryPool.Allocate(size)), size);
+    }
+};
+
 TCell MakeCell(NScheme::TTypeInfo type, const NUdf::TUnboxedValuePod& value,
     const TTypeEnvironment& env, bool copy = true,
     i32 typmod = -1, TMaybe<TString>* error = {});
 
+TCell MakeCell(NScheme::TTypeInfo type, const NUdf::TUnboxedValuePod& value,
+    const TStringProviderBackend& env, bool copy = true,
+    i32 typmod = -1, TMaybe<TString>* error = {});
+
 void FillKeyTupleValue(const NUdf::TUnboxedValue& row, const TVector<ui32>& rowIndices,
     const TVector<NScheme::TTypeInfo>& rowTypes, TVector<TCell>& cells, const TTypeEnvironment& env);
 
diff --git a/ydb/core/kqp/runtime/kqp_stream_lookup_worker.cpp b/ydb/core/kqp/runtime/kqp_stream_lookup_worker.cpp
index b1d56e60dc..e04690d24b 100644
--- a/ydb/core/kqp/runtime/kqp_stream_lookup_worker.cpp
+++ b/ydb/core/kqp/runtime/kqp_stream_lookup_worker.cpp
@@ -143,12 +143,15 @@ public:
     virtual ~TKqpLookupRows() {}
 
     void AddInputRow(NUdf::TUnboxedValue inputRow) final {
+        NMiniKQL::TStringProviderBackend backend;
         std::vector<TCell> keyCells(LookupKeyColumns.size());
         for (size_t colId = 0; colId < LookupKeyColumns.size(); ++colId) {
             const auto* lookupKeyColumn = LookupKeyColumns[colId];
             YQL_ENSURE(lookupKeyColumn->KeyOrder < static_cast<i64>(keyCells.size()));
+            // when making a cell we don't really need to make a copy of data, because
+            // TOwnedCellVec will make its' own copy.
             keyCells[lookupKeyColumn->KeyOrder] = MakeCell(lookupKeyColumn->PType,
-                inputRow.GetElement(colId), TypeEnv, /* copy */ true);
+                inputRow.GetElement(colId), backend, /* copy */ false);
         }
 
         if (keyCells.size() < KeyColumns.size()) {
@@ -425,13 +428,16 @@ public:
     void AddInputRow(NUdf::TUnboxedValue inputRow) final {
         auto joinKey = inputRow.GetElement(0);
         std::vector<TCell> joinKeyCells(LookupKeyColumns.size());
+        NMiniKQL::TStringProviderBackend backend;
 
         if (joinKey.HasValue()) {
             for (size_t colId = 0; colId < LookupKeyColumns.size(); ++colId) {
                 const auto* joinKeyColumn = LookupKeyColumns[colId];
                 YQL_ENSURE(joinKeyColumn->KeyOrder < static_cast<i64>(joinKeyCells.size()));
+                // when making a cell we don't really need to make a copy of data, because
+                // TOwnedCellVec will make its' own copy.
                 joinKeyCells[joinKeyColumn->KeyOrder] = MakeCell(joinKeyColumn->PType,
-                    joinKey.GetElement(colId), TypeEnv, true);
+                    joinKey.GetElement(colId), backend,  /* copy */ false);
             }
         }