Fix vector_index_build reboots test (#18900) (#18953)

vitalif · azevaykin · commit 30de5dbb44f8 · 2025-06-05T12:00:31.000+03:00
diff --git a/ydb/core/tx/schemeshard/ut_helpers/helpers.cpp b/ydb/core/tx/schemeshard/ut_helpers/helpers.cpp
@@ -2627,4 +2627,31 @@ namespace NSchemeShardUT_Private {
 
         return result;
     }
+
+    void WriteVectorTableRows(TTestActorRuntime& runtime, ui64 schemeShardId, ui64 txId, const TString & tablePath, bool withValue, ui32 shard, ui32 min, ui32 max) {
+        TVector<TCell> cells;
+        ui8 str[6] = { 0 };
+        str[4] = (ui8)Ydb::Table::VectorIndexSettings::VECTOR_TYPE_UINT8;
+        for (ui32 key = min; key < max; ++key) {
+            str[0] = ((key+106)* 7) % 256;
+            str[1] = ((key+106)*17) % 256;
+            str[2] = ((key+106)*37) % 256;
+            str[3] = ((key+106)*47) % 256;
+            cells.emplace_back(TCell::Make(key));
+            cells.emplace_back(TCell((const char*)str, 5));
+            if (withValue) {
+                // optionally use the same value for an additional covered string column
+                cells.emplace_back(TCell((const char*)str, 5));
+            }
+        }
+        std::vector<ui32> columnIds{1, 2};
+        if (withValue) {
+            columnIds.push_back(3);
+        }
+        TSerializedCellMatrix matrix(cells, max-min, withValue ? 3 : 2);
+        WriteOp(runtime, schemeShardId, txId, tablePath,
+            shard, NKikimrDataEvents::TEvWrite::TOperation::OPERATION_UPSERT,
+            columnIds, std::move(matrix), true);
+    };
+
 }
diff --git a/ydb/core/tx/schemeshard/ut_helpers/helpers.h b/ydb/core/tx/schemeshard/ut_helpers/helpers.h
@@ -653,5 +653,7 @@ namespace NSchemeShardUT_Private {
     NKikimrMiniKQL::TResult ReadTable(TTestActorRuntime& runtime, ui64 tabletId,
         const TString& table, const TVector<TString>& pk, const TVector<TString>& columns, const TString& rangeFlags = "");
 
+    void WriteVectorTableRows(TTestActorRuntime& runtime, ui64 schemeShardId, ui64 txId, const TString & tablePath,
+        bool withValue, ui32 shard, ui32 min, ui32 max);
 
 } //NSchemeShardUT_Private
diff --git a/ydb/core/tx/schemeshard/ut_helpers/test_env.cpp b/ydb/core/tx/schemeshard/ut_helpers/test_env.cpp
@@ -1173,7 +1173,9 @@ bool NSchemeShardUT_Private::TTestWithReboots::PassUserRequests(TTestActorRuntim
            event->Type == TEvIndexBuilder::EvCreateRequest ||
            event->Type == TEvIndexBuilder::EvGetRequest ||
            event->Type == TEvIndexBuilder::EvCancelRequest ||
-           event->Type == TEvIndexBuilder::EvForgetRequest
+           event->Type == TEvIndexBuilder::EvForgetRequest ||
+           // without it, ut_vector_index_build_reboots test hangs on GetRequest on the very first reboot
+           event->Type == TEvTablet::EvCommitResult
         ;
 }
 
diff --git a/ydb/core/tx/schemeshard/ut_helpers/test_env.h b/ydb/core/tx/schemeshard/ut_helpers/test_env.h
@@ -154,6 +154,10 @@ namespace NSchemeShardUT_Private {
 
     ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
     // A wrapper to run test scenarios with reboots of schemeshard, hive and coordinator
+    // The idea is to run the same test scenario multiple times and on each run restart a tablet **once**
+    // on receiving a non-filtered event. A given tablet is restarted when it receives an event for which
+    // PassUserRequests() doesn't return true. On the first run, it is restarted on the first such event,
+    // on the second run - on the second event and so on.
     class TTestWithReboots {
     protected:
         struct TDatashardLogBatchingSwitch {
diff --git a/ydb/core/tx/schemeshard/ut_helpers/test_with_reboots.h b/ydb/core/tx/schemeshard/ut_helpers/test_with_reboots.h
@@ -6,6 +6,9 @@ namespace NSchemeShardUT_Private {
 
 class TTestWithTabletReboots: public TTestWithReboots {
 public:
+    explicit TTestWithTabletReboots(bool killOnCommit = false)
+        : TTestWithReboots(killOnCommit)
+    {}
     void Run(std::function<void(TTestActorRuntime& runtime, bool& activeZone)> testScenario) {
         TDatashardLogBatchingSwitch logBatchingSwitch(false /* without batching */);
         RunWithTabletReboots(testScenario);
@@ -14,6 +17,9 @@ class TTestWithTabletReboots: public TTestWithReboots {
 
 class TTestWithPipeResets: public TTestWithReboots {
 public:
+    explicit TTestWithPipeResets(bool killOnCommit = false)
+        : TTestWithReboots(killOnCommit)
+    {}
     void Run(std::function<void(TTestActorRuntime& runtime, bool& activeZone)> testScenario) {
         TDatashardLogBatchingSwitch logBatchingSwitch(false /* without batching */);
         RunWithPipeResets(testScenario);
diff --git a/ydb/core/tx/schemeshard/ut_index_build/ut_vector_index_build.cpp b/ydb/core/tx/schemeshard/ut_index_build/ut_vector_index_build.cpp
@@ -90,27 +90,9 @@ Y_UNIT_TEST_SUITE (VectorIndexBuildTest) {
         env.TestWaitNotification(runtime, txId, tenantSchemeShard);
 
         // Write data directly into shards
-        auto fillRows = [&](const TString & tablePath, ui32 shard, ui32 min, ui32 max) {
-            TVector<TCell> cells;
-            ui8 str[6] = { 0 };
-            str[4] = (ui8)Ydb::Table::VectorIndexSettings::VECTOR_TYPE_UINT8;
-            for (ui32 key = min; key < max; ++key) {
-                str[0] = ((key+106)* 7) % 256;
-                str[1] = ((key+106)*17) % 256;
-                str[2] = ((key+106)*37) % 256;
-                str[3] = ((key+106)*47) % 256;
-                cells.emplace_back(TCell::Make(key));
-                cells.emplace_back(TCell((const char*)str, 5));
-            }
-            std::vector<ui32> columnIds{1, 2};
-            TSerializedCellMatrix matrix(cells, max-min, 2);
-            WriteOp(runtime, tenantSchemeShard, ++txId, tablePath,
-                shard, NKikimrDataEvents::TEvWrite::TOperation::OPERATION_UPSERT,
-                columnIds, std::move(matrix), true);
-        };
-        fillRows("/MyRoot/ServerLessDB/Table", 0, 0, 50);
-        fillRows("/MyRoot/ServerLessDB/Table", 1, 50, 150);
-        fillRows("/MyRoot/ServerLessDB/Table", 2, 150, 200);
+        WriteVectorTableRows(runtime, tenantSchemeShard, ++txId, "/MyRoot/ServerLessDB/Table", false, 0, 0, 50);
+        WriteVectorTableRows(runtime, tenantSchemeShard, ++txId, "/MyRoot/ServerLessDB/Table", false, 1, 50, 150);
+        WriteVectorTableRows(runtime, tenantSchemeShard, ++txId, "/MyRoot/ServerLessDB/Table", false, 2, 150, 200);
 
         runtime.SetLogPriority(NKikimrServices::TX_DATASHARD, NLog::PRI_TRACE);
         runtime.SetLogPriority(NKikimrServices::BUILD_INDEX, NLog::PRI_TRACE);
@@ -293,7 +275,7 @@ Y_UNIT_TEST_SUITE (VectorIndexBuildTest) {
         )");
         env.TestWaitNotification(runtime, txId, tenantSchemeShard);
 
-        fillRows("/MyRoot/CommonDB/Table", 0, 100, 300);
+        WriteVectorTableRows(runtime, tenantSchemeShard, ++txId, "/MyRoot/CommonDB/Table", false, 0, 100, 300);
 
         TVector<TString> billRecords;
         observerHolder = runtime.AddObserver<NMetering::TEvMetering::TEvWriteMeteringJson>([&](NMetering::TEvMetering::TEvWriteMeteringJson::TPtr& event) {
diff --git a/ydb/core/tx/schemeshard/ut_vector_index_build_reboots/ut_vector_index_build_reboots.cpp b/ydb/core/tx/schemeshard/ut_vector_index_build_reboots/ut_vector_index_build_reboots.cpp
@@ -8,107 +8,30 @@ using namespace NSchemeShard;
 using namespace NSchemeShardUT_Private;
 
 
-static void WriteRows(TTestActorRuntime& runtime, ui64 tabletId, ui32 key, ui32 index) {
-    TString writeQuery = Sprintf(R"(
-        (
-            (let keyNull   '( '('key   (Null) ) ) )
-            (let row0   '( '('embedding (String '%s ) )  '('value (String 'aaaa) ) ) )
-            (let key0   '( '('key   (Uint32 '%u ) ) ) )
-            (let row0   '( '('embedding (String '%s ) )  '('value (String 'aaaa) ) ) )
-            (let key1   '( '('key   (Uint32 '%u ) ) ) )
-            (let row1   '( '('embedding (String '%s ) )  '('value (String 'aaaa) ) ) )
-            (let key2   '( '('key   (Uint32 '%u ) ) ) )
-            (let row2   '( '('embedding (String '%s ) )  '('value (String 'aaaa) ) ) )
-            (let key3   '( '('key   (Uint32 '%u ) ) ) )
-            (let row3   '( '('embedding (String '%s ) )  '('value (String 'aaaa) ) ) )
-            (let key4   '( '('key   (Uint32 '%u ) ) ) )
-            (let row4   '( '('embedding (String '%s ) )  '('value (String 'aaaa) ) ) )
-            (let key5   '( '('key   (Uint32 '%u ) ) ) )
-            (let row5   '( '('embedding (String '%s ) )  '('value (String 'aaaa) ) ) )
-            (let key6   '( '('key   (Uint32 '%u ) ) ) )
-            (let row6   '( '('embedding (String '%s ) )  '('value (String 'aaaa) ) ) )
-            (let key7   '( '('key   (Uint32 '%u ) ) ) )
-            (let row7   '( '('embedding (String '%s ) )  '('value (String 'aaaa) ) ) )
-            (let key8   '( '('key   (Uint32 '%u ) ) ) )
-            (let row8   '( '('embedding (String '%s ) )  '('value (String 'aaaa) ) ) )
-            (let key9   '( '('key   (Uint32 '%u ) ) ) )
-            (let row9   '( '('embedding (String '%s ) )  '('value (String 'aaaa) ) ) )
-
-            (return (AsList
-                        (UpdateRow '__user__Table keyNull row0)
-                        (UpdateRow '__user__Table key0 row0)
-                        (UpdateRow '__user__Table key1 row1)
-                        (UpdateRow '__user__Table key2 row2)
-                        (UpdateRow '__user__Table key3 row3)
-                        (UpdateRow '__user__Table key4 row4)
-                        (UpdateRow '__user__Table key5 row5)
-                        (UpdateRow '__user__Table key6 row6)
-                        (UpdateRow '__user__Table key7 row7)
-                        (UpdateRow '__user__Table key8 row8)
-                        (UpdateRow '__user__Table key9 row9)
-                     )
-            )
-        )
-    )",
-                       std::to_string(1000*index + 0).c_str(),
-         1000*key + 0, std::to_string(1000*index + 0).c_str(),
-         1000*key + 1, std::to_string(1000*index + 1).c_str(),
-         1000*key + 2, std::to_string(1000*index + 2).c_str(),
-         1000*key + 3, std::to_string(1000*index + 3).c_str(),
-         1000*key + 4, std::to_string(1000*index + 4).c_str(),
-         1000*key + 5, std::to_string(1000*index + 5).c_str(),
-         1000*key + 6, std::to_string(1000*index + 6).c_str(),
-         1000*key + 7, std::to_string(1000*index + 7).c_str(),
-         1000*key + 8, std::to_string(1000*index + 8).c_str(),
-         1000*key + 9, std::to_string(1000*index + 9).c_str());
-
-    NKikimrMiniKQL::TResult result;
-    TString err;
-    NKikimrProto::EReplyStatus status = LocalMiniKQL(runtime, tabletId, writeQuery, result, err);
-    UNIT_ASSERT_VALUES_EQUAL(err, "");
-    UNIT_ASSERT_VALUES_EQUAL(status, NKikimrProto::EReplyStatus::OK);;
-}
-
 Y_UNIT_TEST_SUITE(VectorIndexBuildTestReboots) {
     Y_UNIT_TEST_WITH_REBOOTS(BaseCase) {
-        T t;
+        // Without killOnCommit, the schemeshard doesn't get rebooted on TEvDataShard::Ev***KMeansResponse's,
+        // and thus the vector index build process is never interrupted at all because there are no other
+        // events to reboot on.
+        T t(true /*killOnCommit*/);
         t.Run([&](TTestActorRuntime& runtime, bool& activeZone) {
             {
                 TInactiveZone inactive(activeZone);
 
                 TestCreateTable(runtime, ++t.TxId, "/MyRoot", R"(
-                  Name: "dir/Table"
-                  Columns { Name: "key"       Type: "Uint32" }
-                  Columns { Name: "embedding" Type: "String" }
-                  Columns { Name: "value"     Type: "String" }
-                  KeyColumnNames: ["key"]
-                  UniformPartitionsCount: 2
-                 )");
+                    Name: "dir/Table"
+                    Columns { Name: "key"       Type: "Uint32" }
+                    Columns { Name: "embedding" Type: "String" }
+                    Columns { Name: "value"     Type: "String" }
+                    KeyColumnNames: ["key"]
+                    SplitBoundary { KeyPrefix { Tuple { Optional { Uint32: 50 } } } }
+                    SplitBoundary { KeyPrefix { Tuple { Optional { Uint32: 150 } } } }
+                )");
                 t.TestEnv->TestWaitNotification(runtime, t.TxId);
 
-                for (ui32 delta = 0; delta < 2; ++delta) {
-                    WriteRows(runtime, TTestTxConfig::FakeHiveTablets, 1 + delta, 100 + delta);
-                }
-
-                // Check src shard has the lead key as null
-                {
-                    NKikimrMiniKQL::TResult result;
-                    TString err;
-                    ui32 status = LocalMiniKQL(runtime, TTestTxConfig::FakeHiveTablets, R"(
-                    (
-                        (let range '('('key (Null) (Void))))
-                        (let columns '('key 'embedding))
-                        (let result (SelectRange '__user__Table range columns '()))
-                        (return (AsList (SetResult 'Result result)))
-                    )
-                    )", result, err);
-
-                    UNIT_ASSERT_VALUES_EQUAL_C(status, static_cast<ui32>(NKikimrProto::OK), err);
-                    UNIT_ASSERT_VALUES_EQUAL(err, "");
-
-                    //                                   V -- here the null key
-                    NKqp::CompareYson(R"([[[[[["101000"];#];[["100000"];["1000"]];[["100001"];["1001"]];[["100002"];["1002"]];[["100003"];["1003"]];[["100004"];["1004"]];[["100005"];["1005"]];[["100006"];["1006"]];[["100007"];["1007"]];[["100008"];["1008"]];[["100009"];["1009"]];[["101000"];["2000"]];[["101001"];["2001"]];[["101002"];["2002"]];[["101003"];["2003"]];[["101004"];["2004"]];[["101005"];["2005"]];[["101006"];["2006"]];[["101007"];["2007"]];[["101008"];["2008"]];[["101009"];["2009"]]];%false]]])", result);
-                }
+                WriteVectorTableRows(runtime, TTestTxConfig::SchemeShard, ++t.TxId, "/MyRoot/dir/Table", true, 0, 0, 50);
+                WriteVectorTableRows(runtime, TTestTxConfig::SchemeShard, ++t.TxId, "/MyRoot/dir/Table", true, 1, 50, 150);
+                WriteVectorTableRows(runtime, TTestTxConfig::SchemeShard, ++t.TxId, "/MyRoot/dir/Table", true, 2, 150, 200);
             }
 
             AsyncBuildVectorIndex(runtime,  ++t.TxId, TTestTxConfig::SchemeShard, "/MyRoot", "/MyRoot/dir/Table", "index1", "embedding", {"value"});

Original file line number	Diff line number	Diff line change
`@@ -1173,7 +1173,9 @@ bool NSchemeShardUT_Private::TTestWithReboots::PassUserRequests(TTestActorRuntim`
`1173`	`1173`	`event->Type == TEvIndexBuilder::EvCreateRequest \|\|`
`1174`	`1174`	`event->Type == TEvIndexBuilder::EvGetRequest \|\|`
`1175`	`1175`	`event->Type == TEvIndexBuilder::EvCancelRequest \|\|`
`1176`		`- event->Type == TEvIndexBuilder::EvForgetRequest`
	`1176`	`+ event->Type == TEvIndexBuilder::EvForgetRequest \|\|`
	`1177`	`+ // without it, ut_vector_index_build_reboots test hangs on GetRequest on the very first reboot`
	`1178`	`+ event->Type == TEvTablet::EvCommitResult`
`1177`	`1179`	`;`
`1178`	`1180`	`}`
`1179`	`1181`