Reduce overhead for TKqpTransaction execution. (KIKIMR-14782)

ref:4a6e0189ebc5ba36b46b2d428158f2158339a235

9 files changed, 118 insertions, 88 deletions

diff --git a/ydb/core/kqp/rm/kqp_rm.cpp b/ydb/core/kqp/rm/kqp_rm.cpp
index 30fc08ff5f8..dd79ef8ea72 100644
--- a/ydb/core/kqp/rm/kqp_rm.cpp
+++ b/ydb/core/kqp/rm/kqp_rm.cpp
@@ -611,7 +611,10 @@ private:
     }
 
     void HandleWork(TEvPrivate::TEvPublishResources::TPtr&) {
-        WbState.PublishScheduledAt.reset();
+        with_lock (Lock) {
+            PublishScheduledAt.reset();
+        }
+
         PublishResourceUsage("batching");
     }
 
@@ -754,8 +757,14 @@ private:
                 if (WbState.LastPublishTime) {
                     str << "Last publish time: " << *WbState.LastPublishTime << Endl;
                 }
-                if (WbState.PublishScheduledAt) {
-                    str << "Next publish time: " << *WbState.PublishScheduledAt << Endl;
+
+                std::optional<TInstant> publishScheduledAt;
+                with_lock (Lock) {
+                    publishScheduledAt = PublishScheduledAt;
+                }
+
+                if (publishScheduledAt) {
+                    str << "Next publish time: " << *publishScheduledAt << Endl;
                 }
 
                 str << Endl << "Transactions:" << Endl;
@@ -804,24 +813,38 @@ private:
     }
 
     void FireResourcesPublishing() {
+        with_lock (Lock) {
+            if (PublishScheduledAt) {
+                return;
+            }
+        }
+
         ActorSystem->Send(SelfId(), new TEvPrivate::TEvSchedulePublishResources);
     }
 
     void PublishResourceUsage(TStringBuf reason) {
-        if (WbState.PublishScheduledAt) {
-            return;
-        }
-
         TDuration publishInterval;
+        std::optional<TInstant> publishScheduledAt;
+
         with_lock (Lock) {
             publishInterval = TDuration::Seconds(Config.GetPublishStatisticsIntervalSec());
+            publishScheduledAt = PublishScheduledAt;
+        }
+
+        if (publishScheduledAt) {
+            return;
         }
 
         auto now = ActorSystem->Timestamp();
         if (publishInterval && WbState.LastPublishTime && now - *WbState.LastPublishTime < publishInterval) {
-            WbState.PublishScheduledAt = *WbState.LastPublishTime + publishInterval;
-            Schedule(*WbState.PublishScheduledAt - now, new TEvPrivate::TEvPublishResources);
-            LOG_D("Schedule publish at " << *WbState.PublishScheduledAt << ", after " << (*WbState.PublishScheduledAt - now));
+            publishScheduledAt = *WbState.LastPublishTime + publishInterval;
+
+            with_lock (Lock) {
+                PublishScheduledAt = publishScheduledAt;
+            }
+
+            Schedule(*publishScheduledAt - now, new TEvPrivate::TEvPublishResources);
+            LOG_D("Schedule publish at " << *publishScheduledAt << ", after " << (*publishScheduledAt - now));
             return;
         }
 
@@ -883,6 +906,9 @@ private:
     std::array<TTxStatesBucket, BucketsCount> Buckets;
     std::atomic<ui64> LastResourceBrokerTaskId = 0;
 
+    // schedule info (guarded by Lock)
+    std::optional<TInstant> PublishScheduledAt;
+
     // Whiteboard specific fields
     struct TWhiteBoardState {
         TString Tenant;
@@ -890,7 +916,6 @@ private:
         ui32 StateStorageGroupId = std::numeric_limits<ui32>::max();
         TActorId BoardPublisherActorId;
         std::optional<TInstant> LastPublishTime;
-        std::optional<TInstant> PublishScheduledAt;
     };
     TWhiteBoardState WbState;
 };
diff --git a/ydb/core/tx/datashard/datashard__engine_host.cpp b/ydb/core/tx/datashard/datashard__engine_host.cpp
index 71329019d41..fa794ce6f1d 100644
--- a/ydb/core/tx/datashard/datashard__engine_host.cpp
+++ b/ydb/core/tx/datashard/datashard__engine_host.cpp
@@ -519,26 +519,28 @@ TEngineBay::TEngineBay(TDataShard * self, TTransactionContext& txc, const TActor
     EngineSettings = MakeHolder<TEngineFlatSettings>(IEngineFlat::EProtocol::V1, AppData(ctx)->FunctionRegistry,
         *TAppData::RandomProvider, *TAppData::TimeProvider, EngineHost.Get(), self->AllocCounters);
 
-    ui64 tabletId = self->TabletID();
-    TraceMessage = Sprintf("Shard %" PRIu64 ", txid %" PRIu64, tabletId, stepTxId.second);
+    auto tabletId = self->TabletID();
+    auto txId = stepTxId.second;
     const TActorSystem* actorSystem = ctx.ExecutorThread.ActorSystem;
-    EngineSettings->LogErrorWriter = [actorSystem, this](const TString& message) {
-        LOG_ERROR_S(*actorSystem, NKikimrServices::MINIKQL_ENGINE, TraceMessage
-            << ", engine error: " << message);
+    EngineSettings->LogErrorWriter = [actorSystem, tabletId, txId](const TString& message) {
+        LOG_ERROR_S(*actorSystem, NKikimrServices::MINIKQL_ENGINE,
+            "Shard %" << tabletId << ", txid %" <<txId << ", engine error: " << message);
     };
 
     if (ctx.LoggerSettings()->Satisfies(NLog::PRI_DEBUG, NKikimrServices::MINIKQL_ENGINE, stepTxId.second)) {
         EngineSettings->BacktraceWriter =
-            [actorSystem, this](const char * operation, ui32 line, const TBackTrace* backtrace)
+            [actorSystem, tabletId, txId](const char * operation, ui32 line, const TBackTrace* backtrace)
             {
-                LOG_DEBUG(*actorSystem, NKikimrServices::MINIKQL_ENGINE, "%s, %s (%" PRIu32 ")\n%s",
-                    TraceMessage.data(), operation, line, backtrace ? backtrace->PrintToString().data() : "");
+                LOG_DEBUG(*actorSystem, NKikimrServices::MINIKQL_ENGINE,
+                    "Shard %" PRIu64 ", txid %, %s (%" PRIu32 ")\n%s",
+                    tabletId, txId, operation, line,
+                    backtrace ? backtrace->PrintToString().data() : "");
             };
     }
 
-    KqpLogFunc = [actorSystem, this](const TStringBuf& message) {
-        LOG_DEBUG_S(*actorSystem, NKikimrServices::KQP_TASKS_RUNNER, TraceMessage
-            << ": " << message);
+    KqpLogFunc = [actorSystem, tabletId, txId](const TStringBuf& message) {
+        LOG_DEBUG_S(*actorSystem, NKikimrServices::KQP_TASKS_RUNNER,
+            "Shard %" << tabletId << ", txid %" << txId << ": " << message);
     };
 
     ComputeCtx = MakeHolder<TKqpDatashardComputeContext>(self, *EngineHost, now);
diff --git a/ydb/core/tx/datashard/datashard__engine_host.h b/ydb/core/tx/datashard/datashard__engine_host.h
index 12adf35243c..12651fdcf94 100644
--- a/ydb/core/tx/datashard/datashard__engine_host.h
+++ b/ydb/core/tx/datashard/datashard__engine_host.h
@@ -84,7 +84,6 @@ public:
 
         Engine.Reset();
         EngineHost.Reset();
-        TraceMessage.clear();
     }
 
     const TValidationInfo& TxInfo() const { return Info; }
@@ -110,7 +109,6 @@ private:
     TValidationInfo Info;
     TEngineHostCounters EngineHostCounters;
     ui64 LockTxId;
-    TString TraceMessage;
     NYql::NDq::TLogFunc KqpLogFunc;
     THolder<NUdf::IApplyContext> KqpApplyCtx;
     THolder<NMiniKQL::TKqpDatashardComputeContext> ComputeCtx;
diff --git a/ydb/core/tx/datashard/datashard_kqp.cpp b/ydb/core/tx/datashard/datashard_kqp.cpp
index 1a5542cee2e..344acc9903a 100644
--- a/ydb/core/tx/datashard/datashard_kqp.cpp
+++ b/ydb/core/tx/datashard/datashard_kqp.cpp
@@ -704,6 +704,7 @@ void KqpFillStats(TDataShard& dataShard, const NKqp::TKqpTasksRunner& tasksRunne
     Y_VERIFY(dataShard.GetUserTables().size() == 1, "TODO: Fix handling of collocated tables");
     auto tableInfo = dataShard.GetUserTables().begin();
 
+    // Directly use StatsMode instead of bool flag, too much is reported for STATS_COLLECTION_BASIC mode.
     bool withProfileStats = statsMode >= NYql::NDqProto::DQ_STATS_MODE_PROFILE;
 
     auto& computeActorStats = *result.Record.MutableComputeActorStats();
@@ -726,24 +727,22 @@ void KqpFillStats(TDataShard& dataShard, const NKqp::TKqpTasksRunner& tasksRunne
         protoTable->SetWriteBytes(taskTableStats.UpdateRowBytes);
         protoTable->SetEraseRows(taskTableStats.NEraseRow);
 
-        { // KQP Extra Stats
+        minFirstRowTimeMs = std::min(minFirstRowTimeMs, protoTask->GetFirstRowTimeMs());
+        maxFinishTimeMs = std::max(maxFinishTimeMs, protoTask->GetFinishTimeMs());
+
+        computeActorStats.SetCpuTimeUs(computeActorStats.GetCpuTimeUs() + protoTask->GetCpuTimeUs());
+
+        if (Y_UNLIKELY(withProfileStats)) {
             NKqpProto::TKqpShardTableExtraStats tableExtraStats;
             tableExtraStats.SetShardId(dataShard.TabletID());
             // tableExtraStats.SetShardCpuTimeUs(...); // TODO: take it from TTxStats
             protoTable->MutableExtra()->PackFrom(tableExtraStats);
         }
-
-        minFirstRowTimeMs = std::min(minFirstRowTimeMs, protoTask->GetFirstRowTimeMs());
-        maxFinishTimeMs = std::max(maxFinishTimeMs, protoTask->GetFinishTimeMs());
-
-        computeActorStats.SetCpuTimeUs(computeActorStats.GetCpuTimeUs() + protoTask->GetCpuTimeUs());
     }
 
     if (maxFinishTimeMs >= minFirstRowTimeMs) {
         computeActorStats.SetDurationUs((maxFinishTimeMs - minFirstRowTimeMs) * 1'000);
     }
-
-    // TODO: fill profile stats
 }
 
 NYql::NDq::TDqTaskRunnerMemoryLimits DefaultKqpDataReqMemoryLimits() {
diff --git a/ydb/core/tx/datashard/datashard_kqp_compute.cpp b/ydb/core/tx/datashard/datashard_kqp_compute.cpp
index fd56a57609c..c7bb8a9152c 100644
--- a/ydb/core/tx/datashard/datashard_kqp_compute.cpp
+++ b/ydb/core/tx/datashard/datashard_kqp_compute.cpp
@@ -98,33 +98,29 @@ ui64 TKqpDatashardComputeContext::GetLocalTableId(const TTableId &tableId) const
     return Shard->GetLocalTableId(tableId);
 }
 
-TVector<std::pair<NScheme::TTypeId, TString>> TKqpDatashardComputeContext::GetKeyColumnsInfo(
-        const TTableId &tableId) const
+const NDataShard::TUserTable::TUserColumn& TKqpDatashardComputeContext::GetKeyColumnInfo(
+    const NDataShard::TUserTable& table, ui32 keyIndex) const
 {
+    MKQL_ENSURE_S(keyIndex <= table.KeyColumnTypes.size());
+    const auto& col = table.Columns.at(table.KeyColumnIds[keyIndex]);
+    MKQL_ENSURE_S(col.IsKey);
+
+    return col;
+}
+
+THashMap<TString, NScheme::TTypeId> TKqpDatashardComputeContext::GetKeyColumnsMap(const TTableId &tableId) const {
     MKQL_ENSURE_S(Shard);
     const NDataShard::TUserTable::TCPtr* tablePtr = Shard->GetUserTables().FindPtr(tableId.PathId.LocalPathId);
     MKQL_ENSURE_S(tablePtr);
     const NDataShard::TUserTable::TCPtr table = *tablePtr;
     MKQL_ENSURE_S(table);
 
-    TVector<std::pair<NScheme::TTypeId, TString>> res;
-    res.reserve(table->KeyColumnTypes.size());
-
+    THashMap<TString, NScheme::TTypeId> columnsMap;
     for (size_t i = 0 ; i < table->KeyColumnTypes.size(); i++) {
         auto col = table->Columns.at(table->KeyColumnIds[i]);
         MKQL_ENSURE_S(col.IsKey);
-        MKQL_ENSURE_S(table->KeyColumnTypes[i] == col.Type);
-        res.push_back({table->KeyColumnTypes[i], col.Name});
-    }
-    return res;
-}
-
-THashMap<TString, NScheme::TTypeId> TKqpDatashardComputeContext::GetKeyColumnsMap(const TTableId &tableId) const {
-    THashMap<TString, NScheme::TTypeId> columnsMap;
+        columnsMap[col.Name] = col.Type;
 
-    auto keyColumns = GetKeyColumnsInfo(tableId);
-    for (const auto& [type, name] : keyColumns) {
-        columnsMap[name] = type;
     }
 
     return columnsMap;
diff --git a/ydb/core/tx/datashard/datashard_kqp_compute.h b/ydb/core/tx/datashard/datashard_kqp_compute.h
index 55a07817fc2..b6dccb27cf5 100644
--- a/ydb/core/tx/datashard/datashard_kqp_compute.h
+++ b/ydb/core/tx/datashard/datashard_kqp_compute.h
@@ -4,12 +4,12 @@
 #include <ydb/core/engine/mkql_engine_flat.h>
 #include <ydb/core/scheme/scheme_tabledefs.h>
 #include <ydb/core/tablet_flat/flat_database.h>
+#include <ydb/core/tx/datashard/datashard_user_table.h>
 
 namespace NKikimr {
     namespace NDataShard {
         class TExecuteKqpScanTxUnit;
         class TDataShard;
-        struct TUserTable;
     }
 }
 
@@ -33,7 +33,8 @@ public:
     void BreakSetLocks() const;
     void SetLockTxId(ui64 lockTxId);
 
-    TVector<std::pair<NScheme::TTypeId, TString>> GetKeyColumnsInfo(const TTableId &tableId) const;
+    const NDataShard::TUserTable::TUserColumn& GetKeyColumnInfo(
+        const NDataShard::TUserTable& table, ui32 keyIndex) const;
     THashMap<TString, NScheme::TTypeId> GetKeyColumnsMap(const TTableId &tableId) const;
 
     void SetHasPersistentChannels(bool value) { PersistentChannels = value; }
diff --git a/ydb/core/tx/datashard/datashard_kqp_delete_rows.cpp b/ydb/core/tx/datashard/datashard_kqp_delete_rows.cpp
index 6a6c3d8231b..f03e13ddf39 100644
--- a/ydb/core/tx/datashard/datashard_kqp_delete_rows.cpp
+++ b/ydb/core/tx/datashard/datashard_kqp_delete_rows.cpp
@@ -130,13 +130,12 @@ IComputationNode* WrapKqpDeleteRows(TCallable& callable, const TComputationNodeF
     auto rowsNode = callable.GetInput(1);
 
     auto tableId = NKqp::ParseTableId(tableNode);
-    auto localTableId = computeCtx.GetLocalTableId(tableId);
-    MKQL_ENSURE_S(localTableId);
-    auto tableKeyTypes = computeCtx.GetKeyColumnsInfo(tableId);
+    auto tableInfo = computeCtx.GetTable(tableId);
+    MKQL_ENSURE(tableInfo, "Table not found: " << tableId.PathId.ToString());
 
     auto rowType = AS_TYPE(TStructType, AS_TYPE(TStreamType, rowsNode.GetStaticType())->GetItemType());
-    MKQL_ENSURE_S(tableKeyTypes.size() == rowType->GetMembersCount(), "Table key column count mismatch"
-        << ", expected: " << tableKeyTypes.size()
+    MKQL_ENSURE_S(tableInfo->KeyColumnIds.size() == rowType->GetMembersCount(), "Table key column count mismatch"
+        << ", expected: " << tableInfo->KeyColumnIds.size()
         << ", actual: " << rowType->GetMembersCount());
 
     THashMap<TString, ui32> inputIndex;
@@ -153,13 +152,15 @@ IComputationNode* WrapKqpDeleteRows(TCallable& callable, const TComputationNodeF
         rowTypes[i] = typeId;
     }
 
-    TVector<ui32> keyIndices(tableKeyTypes.size());
-    for (ui32 i = 0; i < tableKeyTypes.size(); i++) {
-        auto it = inputIndex.find(tableKeyTypes[i].second);
+    TVector<ui32> keyIndices(tableInfo->KeyColumnIds.size());
+    for (ui32 i = 0; i < keyIndices.size(); i++) {
+        auto& columnInfo = computeCtx.GetKeyColumnInfo(*tableInfo, i);
 
-        MKQL_ENSURE_S(rowTypes[it->second] == tableKeyTypes[i].first, "Key type mismatch"
-            << ", column: " << tableKeyTypes[i].second
-            << ", expected: " << tableKeyTypes[i].first
+        auto it = inputIndex.find(columnInfo.Name);
+
+        MKQL_ENSURE_S(rowTypes[it->second] == columnInfo.Type, "Key type mismatch"
+            << ", column: " << columnInfo.Name
+            << ", expected: " << columnInfo.Type
             << ", actual: " << rowTypes[it->second]);
 
         keyIndices[i] = it->second;
diff --git a/ydb/core/tx/datashard/datashard_kqp_read_table.cpp b/ydb/core/tx/datashard/datashard_kqp_read_table.cpp
index f3d5848c4c5..278b5792324 100644
--- a/ydb/core/tx/datashard/datashard_kqp_read_table.cpp
+++ b/ydb/core/tx/datashard/datashard_kqp_read_table.cpp
@@ -16,7 +16,7 @@ using namespace NUdf;
 
 namespace {
 
-void ValidateKeyType(const TType* keyType, const std::pair<NScheme::TTypeId, TString>& keyColumn) {
+void ValidateKeyType(const TType* keyType, NScheme::TTypeId expectedType) {
     auto type = keyType;
 
     if (type->IsOptional()) {
@@ -25,25 +25,31 @@ void ValidateKeyType(const TType* keyType, const std::pair<NScheme::TTypeId, TSt
 
     auto dataType = AS_TYPE(TDataType, type)->GetSchemeType();
 
-    MKQL_ENSURE_S(dataType == keyColumn.first);
+    MKQL_ENSURE_S(dataType == expectedType);
 }
 
-void ValidateKeyTuple(const TTupleType* tupleType, const TVector<std::pair<NScheme::TTypeId, TString>>& keyColumns) {
+void ValidateKeyTuple(const TTupleType* tupleType, const NDataShard::TUserTable& tableInfo,
+    const TKqpDatashardComputeContext& computeCtx)
+{
     MKQL_ENSURE_S(tupleType);
-    MKQL_ENSURE_S(tupleType->GetElementsCount() <= keyColumns.size());
+    MKQL_ENSURE_S(tupleType->GetElementsCount() <= tableInfo.KeyColumnIds.size());
 
     for (ui32 i = 0; i < tupleType->GetElementsCount(); ++i) {
-        ValidateKeyType(tupleType->GetElementType(i), keyColumns[i]);
+        auto& columnInfo = computeCtx.GetKeyColumnInfo(tableInfo, i);
+        ValidateKeyType(tupleType->GetElementType(i), columnInfo.Type);
     }
 }
 
-void ValidateRangeBound(const TTupleType* tupleType, const TVector<std::pair<NScheme::TTypeId, TString>>& keyColumns) {
+void ValidateRangeBound(const TTupleType* tupleType, const NDataShard::TUserTable& tableInfo,
+    const TKqpDatashardComputeContext& computeCtx)
+{
     MKQL_ENSURE_S(tupleType);
-    MKQL_ENSURE_S(tupleType->GetElementsCount() == keyColumns.size() + 1);
+    MKQL_ENSURE_S(tupleType->GetElementsCount() == tableInfo.KeyColumnIds.size() + 1);
 
-    for (ui32 i = 0; i < keyColumns.size(); ++i) {
+    for (ui32 i = 0; i < tupleType->GetElementsCount(); ++i) {
+        auto& columnInfo = computeCtx.GetKeyColumnInfo(tableInfo, i);
         auto elementType = tupleType->GetElementType(i);
-        ValidateKeyType(AS_TYPE(TOptionalType, elementType)->GetItemType(), keyColumns[i]);
+        ValidateKeyType(AS_TYPE(TOptionalType, elementType)->GetItemType(), columnInfo.Type);
     }
 }
 
@@ -382,11 +388,13 @@ IComputationNode* WrapKqpWideReadTableRanges(TCallable& callable, const TComputa
     auto parseResult = ParseWideReadTableRanges(callable);
     auto rangesNode = LocateNode(ctx.NodeLocator, *parseResult.Ranges);
 
-    auto keyColumns = computeCtx.GetKeyColumnsInfo(parseResult.TableId);
+    auto tableInfo = computeCtx.GetTable(parseResult.TableId);
+    MKQL_ENSURE(tableInfo, "Table not found: " << parseResult.TableId.PathId.ToString());
+
     auto keyRangesType = ParseKeyRangesType(parseResult.Ranges->GetType());
     if (keyRangesType) {
-        ValidateRangeBound(keyRangesType->From, keyColumns);
-        ValidateRangeBound(keyRangesType->To, keyColumns);
+        ValidateRangeBound(keyRangesType->From, *tableInfo, computeCtx);
+        ValidateRangeBound(keyRangesType->To, *tableInfo, computeCtx);
     }
 
     IComputationNode* itemsLimit = nullptr;
@@ -408,9 +416,11 @@ IComputationNode* WrapKqpWideReadTable(TCallable& callable, const TComputationNo
     auto fromNode = LocateNode(ctx.NodeLocator, *parseResult.FromTuple);
     auto toNode = LocateNode(ctx.NodeLocator, *parseResult.ToTuple);
 
-    auto keyColumns = computeCtx.GetKeyColumnsInfo(parseResult.TableId);
-    ValidateKeyTuple(parseResult.FromTuple->GetType(), keyColumns);
-    ValidateKeyTuple(parseResult.ToTuple->GetType(), keyColumns);
+    auto tableInfo = computeCtx.GetTable(parseResult.TableId);
+    MKQL_ENSURE(tableInfo, "Table not found: " << parseResult.TableId.PathId.ToString());
+
+    ValidateKeyTuple(parseResult.FromTuple->GetType(), *tableInfo, computeCtx);
+    ValidateKeyTuple(parseResult.ToTuple->GetType(), *tableInfo, computeCtx);
 
     IComputationNode* itemsLimit = nullptr;
     if (parseResult.ItemsLimit) {
diff --git a/ydb/core/tx/datashard/datashard_kqp_upsert_rows.cpp b/ydb/core/tx/datashard/datashard_kqp_upsert_rows.cpp
index a2dfa61bbfc..8e6e3ba023f 100644
--- a/ydb/core/tx/datashard/datashard_kqp_upsert_rows.cpp
+++ b/ydb/core/tx/datashard/datashard_kqp_upsert_rows.cpp
@@ -167,31 +167,30 @@ IComputationNode* WrapKqpUpsertRows(TCallable& callable, const TComputationNodeF
     auto upsertColumnsNode = callable.GetInput(2);
 
     auto tableId = NKqp::ParseTableId(tableNode);
-    auto localTableId = computeCtx.GetLocalTableId(tableId);
-    MKQL_ENSURE_S(localTableId);
-    auto tableKeyTypes = computeCtx.GetKeyColumnsInfo(tableId);
-
     auto tableInfo = computeCtx.GetTable(tableId);
     MKQL_ENSURE(tableInfo, "Table not found: " << tableId.PathId.ToString());
 
     auto rowType = AS_TYPE(TStructType, AS_TYPE(TStreamType, rowsNode.GetStaticType())->GetItemType());
 
-    MKQL_ENSURE_S(tableKeyTypes.size() <= rowType->GetMembersCount(), "not enough columns in the runtime node");
+    MKQL_ENSURE_S(tableInfo->KeyColumnIds.size() <= rowType->GetMembersCount(),
+        "not enough columns in the runtime node");
 
-    THashMap<TString, ui32> inputIndex;
+    THashMap<TStringBuf, ui32> inputIndex;
     TVector<NUdf::TDataTypeId> rowTypes(rowType->GetMembersCount());
     for (ui32 i = 0; i < rowTypes.size(); ++i) {
         const auto& name = rowType->GetMemberName(i);
-        MKQL_ENSURE_S(inputIndex.emplace(TString(name), i).second);
+        MKQL_ENSURE_S(inputIndex.emplace(name, i).second);
         rowTypes[i] = NKqp::UnwrapDataTypeFromStruct(*rowType, i);
     }
 
-    TVector<ui32> keyIndices(tableKeyTypes.size());
-    for (ui32 i = 0; i < tableKeyTypes.size(); i++) {
-        auto it = inputIndex.find(tableKeyTypes[i].second);
+    TVector<ui32> keyIndices(tableInfo->KeyColumnIds.size());
+    for (ui32 i = 0; i < keyIndices.size(); i++) {
+        auto& columnInfo = computeCtx.GetKeyColumnInfo(*tableInfo, i);
+
+        auto it = inputIndex.find(columnInfo.Name);
         MKQL_ENSURE_S(it != inputIndex.end());
         auto typeId = NKqp::UnwrapDataTypeFromStruct(*rowType, it->second);
-        MKQL_ENSURE_S(typeId == tableKeyTypes[i].first, "row key type missmatch with table key type");
+        MKQL_ENSURE_S(typeId == columnInfo.Type, "row key type missmatch with table key type");
         keyIndices[i] = it->second;
     }
 
@@ -221,7 +220,6 @@ IComputationNode* WrapKqpUpsertRows(TCallable& callable, const TComputationNodeF
 
         MKQL_ENSURE_S(rowTypes[upsertColumn.RowIndex] == tableColumn->Type,
             "upsert column type missmatch, column: " << tableColumn->Name);
-
     }
 
     return new TKqpUpsertRowsWrapper(ctx.Mutables, computeCtx, tableId,