YQL-20312: Use TCheckpointHashMap in peephole opt

commit_hash:e325770a55a4c7f78286224c83b771457407b371

1 files changed, 739 insertions, 0 deletions

diff --git a/yql/essentials/utils/checkpoint_map_ut.cpp b/yql/essentials/utils/checkpoint_map_ut.cpp
new file mode 100644
index 00000000000..ae347973629
--- /dev/null
+++ b/yql/essentials/utils/checkpoint_map_ut.cpp
@@ -0,0 +1,739 @@
+#include "checkpoint_map.h"
+
+#include <library/cpp/testing/unittest/registar.h>
+#include <util/random/random.h>
+
+namespace NYql {
+
+Y_UNIT_TEST_SUITE(CheckpointHashMapTests) {
+Y_UNIT_TEST(BasicOperations) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    // Check empty map
+    UNIT_ASSERT(map.empty());
+    UNIT_ASSERT_VALUES_EQUAL(0, map.size());
+
+    // Add elements
+    map.Set(1, "one");
+    map.Set(2, "two");
+
+    UNIT_ASSERT(!map.empty());
+    UNIT_ASSERT_VALUES_EQUAL(2, map.size());
+
+    // Check values using find
+    auto it1 = map.find(1);
+    UNIT_ASSERT(it1 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("one", it1->second);
+
+    auto it2 = map.find(2);
+    UNIT_ASSERT(it2 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("two", it2->second);
+}
+
+Y_UNIT_TEST(CheckpointAndRollback) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    // Add initial data
+    map.Set(1, "one");
+    map.Set(2, "two");
+
+    // Make checkpoint
+    map.Checkpoint();
+    UNIT_ASSERT_VALUES_EQUAL(0, map.GetCommandCount());
+
+    // Modify data after checkpoint
+    map.Set(1, "ONE");
+    map.Set(3, "three");
+    map.erase(2);
+
+    UNIT_ASSERT_VALUES_EQUAL(3, map.GetCommandCount());
+
+    auto it1 = map.find(1);
+    UNIT_ASSERT(it1 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("ONE", it1->second);
+
+    auto it3 = map.find(3);
+    UNIT_ASSERT(it3 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("three", it3->second);
+
+    UNIT_ASSERT_VALUES_EQUAL(2, map.size());
+    UNIT_ASSERT(map.find(2) == map.end());
+
+    // Rollback
+    map.Rollback();
+
+    UNIT_ASSERT_VALUES_EQUAL(0, map.GetCommandCount());
+
+    it1 = map.find(1);
+    UNIT_ASSERT(it1 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("one", it1->second);
+
+    auto it2 = map.find(2);
+    UNIT_ASSERT(it2 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("two", it2->second);
+
+    UNIT_ASSERT_VALUES_EQUAL(2, map.size());
+    UNIT_ASSERT(map.find(3) == map.end());
+}
+
+Y_UNIT_TEST(SetOperations) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    map.Checkpoint();
+
+    // Test Set
+    map.Set(1, "one");
+    auto it1 = map.find(1);
+    UNIT_ASSERT(it1 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("one", it1->second);
+
+    // Overwrite existing key
+    map.Set(1, "ONE");
+    it1 = map.find(1);
+    UNIT_ASSERT(it1 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("ONE", it1->second);
+
+    // Test Set with another key
+    map.Set(2, "two");
+    auto it2 = map.find(2);
+    UNIT_ASSERT(it2 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("two", it2->second);
+
+    UNIT_ASSERT_VALUES_EQUAL(2, map.size());
+
+    // Rollback
+    map.Rollback();
+    UNIT_ASSERT(map.empty());
+}
+
+Y_UNIT_TEST(EraseOperations) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    // Prepare data
+    map.Set(1, "one");
+    map.Set(2, "two");
+    map.Set(3, "three");
+
+    map.Checkpoint();
+
+    // Erase by key
+    UNIT_ASSERT_VALUES_EQUAL(1, map.erase(2));
+    UNIT_ASSERT_VALUES_EQUAL(0, map.erase(10)); // non-existent key
+
+    // Erase by key
+    UNIT_ASSERT_VALUES_EQUAL(1, map.erase(3));
+
+    UNIT_ASSERT_VALUES_EQUAL(1, map.size());
+    auto it1 = map.find(1);
+    UNIT_ASSERT(it1 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("one", it1->second);
+
+    // Rollback
+    map.Rollback();
+
+    UNIT_ASSERT_VALUES_EQUAL(3, map.size());
+    it1 = map.find(1);
+    UNIT_ASSERT(it1 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("one", it1->second);
+
+    auto it2 = map.find(2);
+    UNIT_ASSERT(it2 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("two", it2->second);
+
+    auto it3 = map.find(3);
+    UNIT_ASSERT(it3 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("three", it3->second);
+}
+
+Y_UNIT_TEST(SearchOperations) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    map.Set(1, "one");
+    map.Set(2, "two");
+
+    // Test find
+    auto it1 = map.find(1);
+    UNIT_ASSERT(it1 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("one", it1->second);
+
+    auto it2 = map.find(10);
+    UNIT_ASSERT(it2 == map.end());
+
+    // Test count
+    UNIT_ASSERT_VALUES_EQUAL(1, map.count(1));
+    UNIT_ASSERT_VALUES_EQUAL(0, map.count(10));
+
+    // Test contains
+    UNIT_ASSERT(map.contains(1));
+    UNIT_ASSERT(!map.contains(10));
+}
+
+Y_UNIT_TEST(Iterators) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    map.Set(1, "one");
+    map.Set(2, "two");
+    map.Set(3, "three");
+
+    // Test iterators
+    ui32 count = 0;
+    for (auto it = map.cbegin(); it != map.cend(); ++it) {
+        count++;
+        UNIT_ASSERT(it->first >= 1 && it->first <= 3);
+    }
+    UNIT_ASSERT_VALUES_EQUAL(3, count);
+
+    // Test const iterators
+    const auto& constMap = map;
+    count = 0;
+    for (auto it = constMap.cbegin(); it != constMap.cend(); ++it) {
+        count++;
+    }
+    UNIT_ASSERT_VALUES_EQUAL(3, count);
+
+    // Test range-based for
+    count = 0;
+    for (const auto& pair : map) {
+        count++;
+        UNIT_ASSERT(pair.first >= 1 && pair.first <= 3);
+    }
+    UNIT_ASSERT_VALUES_EQUAL(3, count);
+}
+
+Y_UNIT_TEST(ComparisonOperators) {
+    TCheckpointHashMap<ui32, TString> map1;
+    TCheckpointHashMap<ui32, TString> map2;
+
+    // Empty maps are equal
+    UNIT_ASSERT(map1 == map2);
+    UNIT_ASSERT(!(map1 != map2));
+
+    // Add data to first map
+    map1.Set(1, "one");
+    UNIT_ASSERT(!(map1 == map2));
+    UNIT_ASSERT(map1 != map2);
+
+    // Add same data to second map
+    map2.Set(1, "one");
+    UNIT_ASSERT(map1 == map2);
+    UNIT_ASSERT(!(map1 != map2));
+}
+
+Y_UNIT_TEST(MultipleCheckpointsAndRollbacks) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    // State 1
+    map.Set(1, "one");
+    map.Checkpoint();
+
+    // State 2
+    map.Set(2, "two");
+    map.Checkpoint();
+
+    // State 3
+    map.Set(3, "three");
+    map.Set(1, "ONE");
+
+    UNIT_ASSERT_VALUES_EQUAL(3, map.size());
+    auto it1 = map.find(1);
+    UNIT_ASSERT(it1 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("ONE", it1->second);
+
+    // Rollback to state 2
+    map.Rollback();
+    UNIT_ASSERT_VALUES_EQUAL(2, map.size());
+    it1 = map.find(1);
+    UNIT_ASSERT(it1 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("one", it1->second);
+
+    auto it2 = map.find(2);
+    UNIT_ASSERT(it2 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("two", it2->second);
+    UNIT_ASSERT(map.find(3) == map.end());
+
+    // New changes
+    map.Set(4, "four");
+    map.erase(1);
+
+    // Verify state after changes
+    UNIT_ASSERT_VALUES_EQUAL(2, map.size());
+    it2 = map.find(2);
+    UNIT_ASSERT(it2 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("two", it2->second);
+
+    auto it4 = map.find(4);
+    UNIT_ASSERT(it4 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("four", it4->second);
+    UNIT_ASSERT(map.find(1) == map.end());
+
+    // Rollback to state 1
+    map.Rollback();
+    UNIT_ASSERT_VALUES_EQUAL(1, map.size());
+    it1 = map.find(1);
+    UNIT_ASSERT(it1 != map.end());
+    UNIT_ASSERT_VALUES_EQUAL("one", it1->second);
+    UNIT_ASSERT(map.find(2) == map.end());
+    UNIT_ASSERT(map.find(4) == map.end());
+}
+
+Y_UNIT_TEST(ComplexScenario) {
+    TCheckpointHashMap<TString, ui32> map;
+
+    // Initial state
+    map.Set("a", 1);
+    map.Set("b", 2);
+    map.Set("c", 3);
+    map.Checkpoint();
+
+    // Complex changes
+    map.Set("a", 10); // update
+    map.Set("d", 4);  // new element
+    map.erase("b");   // deletion
+    map.Set("e", 5);  // another new element
+    map.Set("c", 30); // another update
+
+    UNIT_ASSERT_VALUES_EQUAL(4, map.size());
+    auto ita = map.find("a");
+    UNIT_ASSERT(ita != map.end());
+    UNIT_ASSERT_VALUES_EQUAL(10, ita->second);
+
+    auto itc = map.find("c");
+    UNIT_ASSERT(itc != map.end());
+    UNIT_ASSERT_VALUES_EQUAL(30, itc->second);
+
+    auto itd = map.find("d");
+    UNIT_ASSERT(itd != map.end());
+    UNIT_ASSERT_VALUES_EQUAL(4, itd->second);
+
+    auto ite = map.find("e");
+    UNIT_ASSERT(ite != map.end());
+    UNIT_ASSERT_VALUES_EQUAL(5, ite->second);
+    UNIT_ASSERT(map.find("b") == map.end());
+
+    // Rollback
+    map.Rollback();
+
+    UNIT_ASSERT_VALUES_EQUAL(3, map.size());
+    ita = map.find("a");
+    UNIT_ASSERT(ita != map.end());
+    UNIT_ASSERT_VALUES_EQUAL(1, ita->second);
+
+    auto itb = map.find("b");
+    UNIT_ASSERT(itb != map.end());
+    UNIT_ASSERT_VALUES_EQUAL(2, itb->second);
+
+    itc = map.find("c");
+    UNIT_ASSERT(itc != map.end());
+    UNIT_ASSERT_VALUES_EQUAL(3, itc->second);
+    UNIT_ASSERT(map.find("d") == map.end());
+    UNIT_ASSERT(map.find("e") == map.end());
+}
+
+Y_UNIT_TEST(EmptyMapOperations) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    map.Checkpoint();
+
+    // Operations on empty map
+    UNIT_ASSERT_VALUES_EQUAL(0, map.erase(1));
+    UNIT_ASSERT(map.find(1) == map.end());
+    UNIT_ASSERT_VALUES_EQUAL(0, map.count(1));
+    UNIT_ASSERT(!map.contains(1));
+
+    map.Rollback();
+    UNIT_ASSERT(map.empty());
+}
+
+Y_UNIT_TEST(RollbackWithNoChanges) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    map.Checkpoint();
+    map.Checkpoint();
+    map.Set(1, "string");
+    map.Checkpoint();
+    UNIT_ASSERT_VALUES_EQUAL(1, map.size());
+    map.Rollback();
+    UNIT_ASSERT_VALUES_EQUAL(1, map.size());
+    map.Rollback();
+    UNIT_ASSERT_VALUES_EQUAL(0, map.size());
+    map.Rollback();
+    UNIT_ASSERT(map.empty());
+}
+
+Y_UNIT_TEST(RandomOperationsTest) {
+    // Test constants
+    constexpr ui32 initialSize = 20;
+    constexpr ui32 numOperations = 100;
+    constexpr ui32 keyRange = 50;
+    constexpr ui32 valueRange = 1000;
+    constexpr ui32 initialValueMultiplier = 10;
+
+    enum class EOperation {
+        Insert = 0,
+        Erase = 1,
+        Update = 2,
+        Count = 3
+    };
+
+    TCheckpointHashMap<ui32, ui32> map;
+    THashMap<ui32, ui32> referenceMap;
+
+    // Fill initial data
+    for (ui32 i = 0; i < initialSize; ++i) {
+        ui32 value = i * initialValueMultiplier;
+        map.Set(i, value);
+        referenceMap[i] = value;
+    }
+
+    // Save reference state before checkpoint
+    auto savedState = referenceMap;
+    map.Checkpoint();
+
+    // Perform random operations
+    for (ui32 op = 0; op < numOperations; ++op) {
+        auto operation = static_cast<EOperation>(RandomNumber<ui32>() % static_cast<ui32>(EOperation::Count));
+        ui32 key = RandomNumber<ui32>() % keyRange;
+
+        switch (operation) {
+            case EOperation::Insert: {
+                ui32 value = RandomNumber<ui32>() % valueRange;
+                map.Set(key, value);
+                referenceMap[key] = value;
+                break;
+            }
+            case EOperation::Erase: {
+                map.erase(key);
+                referenceMap.erase(key);
+                break;
+            }
+            case EOperation::Update: {
+                auto it = referenceMap.find(key);
+                if (it != referenceMap.end()) {
+                    ui32 newValue = RandomNumber<ui32>() % valueRange;
+                    map.Set(key, newValue);
+                    referenceMap[key] = newValue;
+                }
+                break;
+            }
+            case EOperation::Count:
+                // Should not reach here
+                break;
+        }
+    }
+
+    // Verify current state matches reference
+    UNIT_ASSERT_VALUES_EQUAL(map.size(), referenceMap.size());
+    for (const auto& pair : referenceMap) {
+        UNIT_ASSERT(map.contains(pair.first));
+        auto it = map.find(pair.first);
+        UNIT_ASSERT(it != map.end());
+        UNIT_ASSERT_VALUES_EQUAL(it->second, pair.second);
+    }
+
+    // Rollback and verify state matches saved state
+    map.Rollback();
+
+    UNIT_ASSERT_VALUES_EQUAL(map.size(), savedState.size());
+    for (const auto& pair : savedState) {
+        UNIT_ASSERT(map.contains(pair.first));
+        auto it = map.find(pair.first);
+        UNIT_ASSERT(it != map.end());
+        UNIT_ASSERT_VALUES_EQUAL(it->second, pair.second);
+    }
+
+    // Verify no extra elements exist
+    for (const auto& pair : map) {
+        UNIT_ASSERT(savedState.find(pair.first) != savedState.end());
+    }
+}
+
+Y_UNIT_TEST(MultipleCheckpointsAndRollbacksTest) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    // Initial state
+    map.Set(1, "initial1");
+    map.Set(2, "initial2");
+
+    // First checkpoint
+    map.Checkpoint();
+    map.Set(1, "checkpoint1_1");
+    map.Set(3, "checkpoint1_3");
+
+    // Second checkpoint
+    map.Checkpoint();
+    map.Set(2, "checkpoint2_2");
+    map.Set(4, "checkpoint2_4");
+
+    // Third checkpoint
+    map.Checkpoint();
+    map.Set(1, "checkpoint3_1");
+    map.erase(3);
+
+    // Verify current state
+    UNIT_ASSERT_VALUES_EQUAL(3, map.size());
+    UNIT_ASSERT_VALUES_EQUAL("checkpoint3_1", map.Get(1));
+    UNIT_ASSERT_VALUES_EQUAL("checkpoint2_2", map.Get(2));
+    UNIT_ASSERT_VALUES_EQUAL("checkpoint2_4", map.Get(4));
+    UNIT_ASSERT(map.find(3) == map.end());
+
+    // First rollback (undo checkpoint 3)
+    map.Rollback();
+    UNIT_ASSERT_VALUES_EQUAL(4, map.size());
+    UNIT_ASSERT_VALUES_EQUAL("checkpoint1_1", map.Get(1));
+    UNIT_ASSERT_VALUES_EQUAL("checkpoint2_2", map.Get(2));
+    UNIT_ASSERT_VALUES_EQUAL("checkpoint1_3", map.Get(3));
+    UNIT_ASSERT_VALUES_EQUAL("checkpoint2_4", map.Get(4));
+
+    // Second rollback (undo checkpoint 2)
+    map.Rollback();
+    UNIT_ASSERT_VALUES_EQUAL(3, map.size());
+    UNIT_ASSERT_VALUES_EQUAL("checkpoint1_1", map.Get(1));
+    UNIT_ASSERT_VALUES_EQUAL("initial2", map.Get(2));
+    UNIT_ASSERT_VALUES_EQUAL("checkpoint1_3", map.Get(3));
+    UNIT_ASSERT(map.find(4) == map.end());
+
+    // Third rollback (undo checkpoint 1)
+    map.Rollback();
+    UNIT_ASSERT_VALUES_EQUAL(2, map.size());
+    UNIT_ASSERT_VALUES_EQUAL("initial1", map.Get(1));
+    UNIT_ASSERT_VALUES_EQUAL("initial2", map.Get(2));
+    UNIT_ASSERT(map.find(3) == map.end());
+    UNIT_ASSERT(map.find(4) == map.end());
+
+    // Test multiple checkpoints without rollback
+    map.Checkpoint();
+    map.Set(5, "test5");
+    map.Checkpoint();
+    map.Set(6, "test6");
+    map.Checkpoint();
+    map.Set(7, "test7");
+
+    UNIT_ASSERT_VALUES_EQUAL(5, map.size());
+
+    // Multiple rollbacks
+    map.Rollback(); // Remove 7
+    UNIT_ASSERT_VALUES_EQUAL(4, map.size());
+    UNIT_ASSERT(map.find(7) == map.end());
+
+    map.Rollback(); // Remove 6
+    UNIT_ASSERT_VALUES_EQUAL(3, map.size());
+    UNIT_ASSERT(map.find(6) == map.end());
+
+    map.Rollback(); // Remove 5
+    UNIT_ASSERT_VALUES_EQUAL(2, map.size());
+    UNIT_ASSERT(map.find(5) == map.end());
+}
+
+Y_UNIT_TEST(RollbackWithoutCheckpointTest) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    // Try to rollback without checkpoint - should throw
+    UNIT_ASSERT_EXCEPTION_CONTAINS(map.Rollback(), yexception, "Cannot rollback without checkpoint");
+
+    // Add some data and try rollback again - should still throw
+    map.Set(1, "test");
+    UNIT_ASSERT_EXCEPTION_CONTAINS(map.Rollback(), yexception, "Cannot rollback without checkpoint");
+
+    // Now make checkpoint and rollback should work
+    map.Checkpoint();
+    map.Set(2, "test2");
+    map.Rollback(); // Should not throw
+
+    // After rollback, another rollback should throw
+    UNIT_ASSERT_EXCEPTION_CONTAINS(map.Rollback(), yexception, "Cannot rollback without checkpoint");
+}
+
+Y_UNIT_TEST(GetUnderlyingMapTest) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    map.Set(1, "one");
+    map.Set(2, "two");
+
+    const auto& underlying = map.GetUnderlyingMap();
+    UNIT_ASSERT_VALUES_EQUAL(2, underlying.size());
+    UNIT_ASSERT(underlying.find(1) != underlying.end());
+    UNIT_ASSERT(underlying.find(2) != underlying.end());
+    UNIT_ASSERT_VALUES_EQUAL("one", underlying.at(1));
+    UNIT_ASSERT_VALUES_EQUAL("two", underlying.at(2));
+}
+
+Y_UNIT_TEST(GetCommandCountTest) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    // No checkpoint - command count should be 0
+    UNIT_ASSERT_VALUES_EQUAL(0, map.GetCommandCount());
+
+    map.Set(1, "one");
+    UNIT_ASSERT_VALUES_EQUAL(0, map.GetCommandCount()); // Still no checkpoint
+
+    map.Checkpoint();
+    UNIT_ASSERT_VALUES_EQUAL(0, map.GetCommandCount()); // Just after checkpoint
+
+    map.Set(2, "two");
+    UNIT_ASSERT_VALUES_EQUAL(1, map.GetCommandCount());
+
+    map.Set(3, "three");
+    UNIT_ASSERT_VALUES_EQUAL(2, map.GetCommandCount());
+
+    map.erase(1);
+    UNIT_ASSERT_VALUES_EQUAL(3, map.GetCommandCount());
+
+    map.Rollback();
+    UNIT_ASSERT_VALUES_EQUAL(0, map.GetCommandCount()); // After rollback
+}
+
+Y_UNIT_TEST(OperationsWithoutCheckpointTest) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    // Operations without checkpoint should not record commands
+    map.Set(1, "one");
+    map.Set(2, "two");
+    map.erase(1);
+    map.Set(3, "three");
+
+    UNIT_ASSERT_VALUES_EQUAL(0, map.GetCommandCount());
+    UNIT_ASSERT_VALUES_EQUAL(2, map.size());
+    UNIT_ASSERT_VALUES_EQUAL("two", map.Get(2));
+    UNIT_ASSERT_VALUES_EQUAL("three", map.Get(3));
+
+    // After checkpoint, operations should record commands
+    map.Checkpoint();
+    UNIT_ASSERT_VALUES_EQUAL(0, map.GetCommandCount());
+
+    map.Set(4, "four");
+    UNIT_ASSERT_VALUES_EQUAL(1, map.GetCommandCount());
+
+    map.Set(2, "TWO");
+    UNIT_ASSERT_VALUES_EQUAL(2, map.GetCommandCount());
+
+    // Rollback should work
+    map.Rollback();
+    UNIT_ASSERT_VALUES_EQUAL(0, map.GetCommandCount());
+    UNIT_ASSERT_VALUES_EQUAL(2, map.size());
+    UNIT_ASSERT_VALUES_EQUAL("two", map.Get(2));
+    UNIT_ASSERT(map.find(4) == map.end());
+}
+
+Y_UNIT_TEST(RollbackThenModifyWithoutCheckpointTest) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    // Initial state
+    map.Set(1, "one");
+    map.Set(2, "two");
+    map.Set(3, "three");
+
+    // Make checkpoint
+    map.Checkpoint();
+    map.Set(1, "ONE");
+    map.Set(4, "four");
+    map.erase(2);
+
+    UNIT_ASSERT_VALUES_EQUAL(3, map.size());
+    UNIT_ASSERT_VALUES_EQUAL("ONE", map.Get(1));
+    UNIT_ASSERT_VALUES_EQUAL("three", map.Get(3));
+    UNIT_ASSERT_VALUES_EQUAL("four", map.Get(4));
+    UNIT_ASSERT(map.find(2) == map.end());
+
+    // Rollback
+    map.Rollback();
+    UNIT_ASSERT_VALUES_EQUAL(3, map.size());
+    UNIT_ASSERT_VALUES_EQUAL("one", map.Get(1));
+    UNIT_ASSERT_VALUES_EQUAL("two", map.Get(2));
+    UNIT_ASSERT_VALUES_EQUAL("three", map.Get(3));
+    UNIT_ASSERT(map.find(4) == map.end());
+
+    // Make changes AFTER rollback but WITHOUT new checkpoint
+    map.Set(2, "TWO");
+    map.Set(5, "five");
+    map.erase(3);
+
+    UNIT_ASSERT_VALUES_EQUAL(3, map.size());
+    UNIT_ASSERT_VALUES_EQUAL("one", map.Get(1));
+    UNIT_ASSERT_VALUES_EQUAL("TWO", map.Get(2));
+    UNIT_ASSERT_VALUES_EQUAL("five", map.Get(5));
+    UNIT_ASSERT(map.find(3) == map.end());
+
+    // Try to rollback again - should throw because no active checkpoint
+    UNIT_ASSERT_EXCEPTION_CONTAINS(map.Rollback(), yexception, "Cannot rollback without checkpoint");
+
+    // State should remain unchanged after failed rollback
+    UNIT_ASSERT_VALUES_EQUAL(3, map.size());
+    UNIT_ASSERT_VALUES_EQUAL("one", map.Get(1));
+    UNIT_ASSERT_VALUES_EQUAL("TWO", map.Get(2));
+    UNIT_ASSERT_VALUES_EQUAL("five", map.Get(5));
+    UNIT_ASSERT(map.find(3) == map.end());
+}
+
+Y_UNIT_TEST(RollbackThenCheckpointThenModifyThenRollbackTest) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    // Initial state
+    map.Set(1, "one");
+    map.Set(2, "two");
+
+    // First checkpoint
+    map.Checkpoint();
+    map.Set(1, "ONE");
+    map.Set(3, "three");
+
+    // First rollback
+    map.Rollback();
+    UNIT_ASSERT_VALUES_EQUAL(2, map.size());
+    UNIT_ASSERT_VALUES_EQUAL("one", map.Get(1));
+    UNIT_ASSERT_VALUES_EQUAL("two", map.Get(2));
+    UNIT_ASSERT(map.find(3) == map.end());
+
+    // Second checkpoint AFTER rollback
+    map.Checkpoint();
+    map.Set(2, "TWO");
+    map.Set(4, "four");
+    map.erase(1);
+
+    UNIT_ASSERT_VALUES_EQUAL(2, map.size());
+    UNIT_ASSERT_VALUES_EQUAL("TWO", map.Get(2));
+    UNIT_ASSERT_VALUES_EQUAL("four", map.Get(4));
+    UNIT_ASSERT(map.find(1) == map.end());
+
+    // Second rollback
+    map.Rollback();
+    UNIT_ASSERT_VALUES_EQUAL(2, map.size());
+    UNIT_ASSERT_VALUES_EQUAL("one", map.Get(1));
+    UNIT_ASSERT_VALUES_EQUAL("two", map.Get(2));
+    UNIT_ASSERT(map.find(3) == map.end());
+    UNIT_ASSERT(map.find(4) == map.end());
+}
+
+Y_UNIT_TEST(CheckpointGuardTest) {
+    TCheckpointHashMap<ui32, TString> map;
+
+    // Initial state
+    map.Set(1, "one");
+    map.Set(2, "two");
+
+    {
+        // Guard automatically creates checkpoint
+        TCheckpointGuard<ui32, TString> guard(map);
+
+        // Make changes
+        map.Set(1, "ONE");
+        map.Set(3, "three");
+        map.erase(2);
+
+        UNIT_ASSERT_VALUES_EQUAL(2, map.size());
+        UNIT_ASSERT_VALUES_EQUAL("ONE", map.Get(1));
+        UNIT_ASSERT_VALUES_EQUAL("three", map.Get(3));
+        UNIT_ASSERT(map.find(2) == map.end());
+
+        // Guard destructor will automatically rollback
+    }
+
+    // After guard destruction, changes should be rolled back
+    UNIT_ASSERT_VALUES_EQUAL(2, map.size());
+    UNIT_ASSERT_VALUES_EQUAL("one", map.Get(1));
+    UNIT_ASSERT_VALUES_EQUAL("two", map.Get(2));
+    UNIT_ASSERT(map.find(3) == map.end());
+}
+
+} // Y_UNIT_TEST_SUITE(CheckpointHashMapTests)
+
+} // namespace NYql