KIKIMR-14972: O(1) for TColumnShard::UpdateInsertTableCounters()

ref:17cb374ce73842bbf04ab94d83d5f6a62c00c758

5 files changed, 62 insertions, 59 deletions

diff --git a/ydb/core/tx/columnshard/columnshard.cpp b/ydb/core/tx/columnshard/columnshard.cpp
index 53ff945f0b..742130267e 100644
--- a/ydb/core/tx/columnshard/columnshard.cpp
+++ b/ydb/core/tx/columnshard/columnshard.cpp
@@ -128,27 +128,18 @@ void TColumnShard::UpdateBlobMangerCounters() {
     IncCounter(COUNTER_SMALL_BLOB_DELETE_BYTES, counters.SmallBlobsBytesDeleted);
 }
 
-void TColumnShard::UpdateInsertTableCounters(bool updateCommitted) {
-    using TCounters = NOlap::TInsertTable::TCounters;
-
-    TCounters prepared = InsertTable->GetCountersPrepared();
-    TCounters committed;
+void TColumnShard::UpdateInsertTableCounters() {
+    auto& prepared = InsertTable->GetCountersPrepared();
+    auto& committed = InsertTable->GetCountersCommitted();
 
     SetCounter(COUNTER_PREPARED_RECORDS, prepared.Rows);
     SetCounter(COUNTER_PREPARED_BYTES, prepared.Bytes);
+    SetCounter(COUNTER_COMMITTED_RECORDS, committed.Rows);
+    SetCounter(COUNTER_COMMITTED_BYTES, committed.Bytes);
 
-    if (updateCommitted) {
-        committed = InsertTable->GetCountersCommitted();
-        SetCounter(COUNTER_COMMITTED_RECORDS, committed.Rows);
-        SetCounter(COUNTER_COMMITTED_BYTES, committed.Bytes);
-
-        LOG_S_DEBUG("InsertTable. Prepared: " << prepared.Bytes << " in " << prepared.Rows
-            << " records, committed: " << committed.Bytes << " in " << committed.Rows
-            << " records at tablet " << TabletID());
-    } else {
-        LOG_S_DEBUG("InsertTable. Prepared: " << prepared.Bytes << " in " << prepared.Rows
-            << " records at tablet " << TabletID());
-    }
+    LOG_S_DEBUG("InsertTable. Prepared: " << prepared.Bytes << " in " << prepared.Rows
+        << " records, committed: " << committed.Bytes << " in " << committed.Rows
+        << " records at tablet " << TabletID());
 }
 
 void TColumnShard::UpdateIndexCounters() {
diff --git a/ydb/core/tx/columnshard/columnshard__write.cpp b/ydb/core/tx/columnshard/columnshard__write.cpp
index e641e24746..af19dcca40 100644
--- a/ydb/core/tx/columnshard/columnshard__write.cpp
+++ b/ydb/core/tx/columnshard/columnshard__write.cpp
@@ -82,8 +82,7 @@ bool TTxWrite::Execute(TTransactionContext& txc, const TActorContext&) {
             Y_VERIFY(logoBlobId.BlobSize() == data.size());
             NBlobCache::AddRangeToCache(NBlobCache::TBlobRange(logoBlobId, 0, data.size()), data);
 
-            bool committedChanged = !allAborted.empty();
-            Self->UpdateInsertTableCounters(committedChanged);
+            Self->UpdateInsertTableCounters();
 
             ui64 blobsWritten = Ev->Get()->BlobBatch.GetBlobCount();
             ui64 bytesWritten = Ev->Get()->BlobBatch.GetTotalSize();
diff --git a/ydb/core/tx/columnshard/columnshard_impl.h b/ydb/core/tx/columnshard/columnshard_impl.h
index 6d62b60aab..42c261080f 100644
--- a/ydb/core/tx/columnshard/columnshard_impl.h
+++ b/ydb/core/tx/columnshard/columnshard_impl.h
@@ -443,7 +443,7 @@ private:
     std::unique_ptr<TEvPrivate::TEvWriteIndex> SetupCleanup();
 
     void UpdateBlobMangerCounters();
-    void UpdateInsertTableCounters(bool updateCommitted = true);
+    void UpdateInsertTableCounters();
     void UpdateIndexCounters();
     void UpdateResourceMetrics(const TUsage& usage);
 
diff --git a/ydb/core/tx/columnshard/engines/insert_table.cpp b/ydb/core/tx/columnshard/engines/insert_table.cpp
index 88295ba805..6b205df394 100644
--- a/ydb/core/tx/columnshard/engines/insert_table.cpp
+++ b/ydb/core/tx/columnshard/engines/insert_table.cpp
@@ -13,7 +13,11 @@ bool TInsertTable::Insert(IDbWrapper& dbTable, TInsertedData&& data) {
     }
 
     dbTable.Insert(data);
-    Inserted.emplace(writeId, std::move(data));
+    ui32 dataSize = data.BlobSize();
+    if (Inserted.emplace(writeId, std::move(data)).second) {
+        StatsPrepared.Rows = Inserted.size();
+        StatsPrepared.Bytes += dataSize;
+    }
     return true;
 }
 
@@ -39,8 +43,15 @@ TInsertTable::TCounters TInsertTable::Commit(IDbWrapper& dbTable, ui64 planStep,
         data->Commit(planStep, txId);
         dbTable.Commit(*data);
 
-        CommittedByPathId[data->PathId].emplace(std::move(*data));
-        Inserted.erase(writeId);
+        ui32 dataSize = data->BlobSize();
+        if (CommittedByPathId[data->PathId].emplace(std::move(*data)).second) {
+            ++StatsCommitted.Rows;
+            StatsCommitted.Bytes += dataSize;
+        }
+        if (Inserted.erase(writeId)) {
+            StatsPrepared.Rows = Inserted.size();
+            StatsPrepared.Bytes -= dataSize;
+        }
     }
 
     return counters;
@@ -56,8 +67,12 @@ void TInsertTable::Abort(IDbWrapper& dbTable, ui64 metaShard, const THashSet<TWr
             dbTable.EraseInserted(*data);
             dbTable.Abort(*data);
 
+            ui32 dataSize = data->BlobSize();
             Aborted.emplace(writeId, std::move(*data));
-            Inserted.erase(writeId);
+            if (Inserted.erase(writeId)) {
+                StatsPrepared.Rows = Inserted.size();
+                StatsPrepared.Bytes -= dataSize;
+            }
         }
     }
 }
@@ -104,7 +119,10 @@ THashSet<TWriteId> TInsertTable::DropPath(IDbWrapper& dbTable, ui64 pathId) {
     TSet<TInsertedData> committed = std::move(CommittedByPathId[pathId]);
     CommittedByPathId.erase(pathId);
 
+    StatsCommitted.Rows -= committed.size();
     for (auto& data : committed) {
+        StatsCommitted.Bytes -= data.BlobSize();
+
         dbTable.EraseCommitted(data);
 
         TInsertedData copy = data;
@@ -118,23 +136,16 @@ THashSet<TWriteId> TInsertTable::DropPath(IDbWrapper& dbTable, ui64 pathId) {
     return toAbort;
 }
 
-void TInsertTable::EraseInserted(IDbWrapper& dbTable, const TInsertedData& data) {
-    TWriteId writeId{data.WriteTxId};
-    if (!Inserted.count(writeId)) {
-        return;
-    }
-
-    dbTable.EraseInserted(data);
-    Inserted.erase(writeId);
-}
-
 void TInsertTable::EraseCommitted(IDbWrapper& dbTable, const TInsertedData& data) {
     if (!CommittedByPathId.count(data.PathId)) {
         return;
     }
 
     dbTable.EraseCommitted(data);
-    CommittedByPathId[data.PathId].erase(data);
+    if (CommittedByPathId[data.PathId].erase(data)) {
+        --StatsCommitted.Rows;
+        StatsCommitted.Bytes -= data.BlobSize();
+    }
 }
 
 void TInsertTable::EraseAborted(IDbWrapper& dbTable, const TInsertedData& data) {
@@ -152,7 +163,28 @@ bool TInsertTable::Load(IDbWrapper& dbTable, const TInstant& loadTime) {
     CommittedByPathId.clear();
     Aborted.clear();
 
-    return dbTable.Load(Inserted, CommittedByPathId, Aborted, loadTime);
+    if (!dbTable.Load(Inserted, CommittedByPathId, Aborted, loadTime)) {
+        return false;
+    }
+
+    // update stats
+
+    StatsPrepared = {};
+    StatsCommitted = {};
+
+    StatsPrepared.Rows = Inserted.size();
+    for (auto& [_, data] : Inserted) {
+        StatsPrepared.Bytes += data.BlobSize();
+    }
+
+    for (auto& [_, set] : CommittedByPathId) {
+        StatsCommitted.Rows += set.size();
+        for (auto& data : set) {
+            StatsCommitted.Bytes += data.BlobSize();
+        }
+    }
+
+    return true;
 }
 
 std::vector<TCommittedBlob> TInsertTable::Read(ui64 pathId, ui64 plan, ui64 txId) const {
@@ -181,24 +213,4 @@ void TInsertTable::SetOverloaded(ui64 pathId, bool overload) {
     }
 }
 
-TInsertTable::TCounters TInsertTable::GetCountersPrepared() const {
-    TCounters prepared;
-    prepared.Rows = Inserted.size();
-    for (auto& [_, data] : Inserted) {
-        prepared.Bytes += data.BlobSize();
-    }
-    return prepared;
-}
-
-TInsertTable::TCounters TInsertTable::GetCountersCommitted() const {
-    TCounters committed;
-    for (auto& [_, set] : CommittedByPathId) {
-        committed.Rows += set.size();
-        for (auto& data : set) {
-            committed.Bytes += data.BlobSize();
-        }
-    }
-    return committed;
-}
-
 }
diff --git a/ydb/core/tx/columnshard/engines/insert_table.h b/ydb/core/tx/columnshard/engines/insert_table.h
index ab25ddd721..64922c289f 100644
--- a/ydb/core/tx/columnshard/engines/insert_table.h
+++ b/ydb/core/tx/columnshard/engines/insert_table.h
@@ -117,13 +117,12 @@ public:
     void Abort(IDbWrapper& dbTable, ui64 metaShard, const THashSet<TWriteId>& writeIds);
     THashSet<TWriteId> AbortOld(IDbWrapper& dbTable, const TInstant& now);
     THashSet<TWriteId> DropPath(IDbWrapper& dbTable, ui64 pathId);
-    void EraseInserted(IDbWrapper& dbTable, const TInsertedData& key);
     void EraseCommitted(IDbWrapper& dbTable, const TInsertedData& key);
     void EraseAborted(IDbWrapper& dbTable, const TInsertedData& key);
     std::vector<TCommittedBlob> Read(ui64 pathId, ui64 plan, ui64 txId) const;
     bool Load(IDbWrapper& dbTable, const TInstant& loadTime);
-    TCounters GetCountersPrepared() const;
-    TCounters GetCountersCommitted() const;
+    const TCounters& GetCountersPrepared() const { return StatsPrepared; }
+    const TCounters& GetCountersCommitted() const { return StatsCommitted; }
 
     size_t InsertedSize() const { return Inserted.size(); }
     const THashMap<ui64, TSet<TInsertedData>>& GetCommitted() const { return CommittedByPathId; }
@@ -137,6 +136,8 @@ private:
     THashMap<TWriteId, TInsertedData> Aborted;
     THashSet<ui64> PathsOverloaded;
     TInstant LastCleanup;
+    TCounters StatsPrepared;
+    TCounters StatsCommitted;
 };
 
 }