Skip to content

Commit 4eea2c6

Browse files
kungakunga
authored
Fixed vector index Seed for tests (#20647)
1 parent 1e497ce commit 4eea2c6

File tree

3 files changed

+46
-16
lines changed

3 files changed

+46
-16
lines changed

‎ydb/core/kqp/ut/cost/kqp_cost_ut.cpp

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
#include <format>
44
#include <ydb/core/kqp/counters/kqp_counters.h>
5+
#include <ydb/core/tx/schemeshard/schemeshard_impl.h>
56
#include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/proto/accessor.h>
67

78
#include <library/cpp/json/json_reader.h>
@@ -312,6 +313,8 @@ Y_UNIT_TEST_SUITE(KqpCost) {
312313

313314
kikimr.GetTestServer().GetRuntime()->GetAppData().FeatureFlags.SetEnableVectorIndex(true);
314315
kikimr.GetTestServer().GetRuntime()->GetAppData().FeatureFlags.SetEnableAccessToIndexImplTables(true);
316+
kikimr.GetTestServer().GetRuntime()->SetLogPriority(NKikimrServices::BUILD_INDEX, NActors::NLog::PRI_INFO);
317+
NSchemeShard::gVectorIndexSeed = 1337;
315318

316319
{ // 1. CREATE TABLE
317320
auto result = session.ExecuteSchemeQuery(R"(
@@ -382,6 +385,19 @@ Y_UNIT_TEST_SUITE(KqpCost) {
382385
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToOneLineString());
383386
}
384387

388+
auto debugPringTable = [&](const TString& name) {
389+
auto query = Q_(Sprintf(R"(
390+
SELECT * FROM `%s`
391+
)", name.c_str()));
392+
393+
auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx(), GetDataQuerySettings()).ExtractValueSync();
394+
395+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToOneLineString());
396+
397+
Cerr << name << ":" << Endl;
398+
Cerr << NYdb::FormatResultSetYson(result.GetResultSet(0)) << Endl;
399+
};
400+
385401
auto checkSelect = [&](auto query, TMap<TString, ui64> expectedReadsByTable) {
386402
auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx(), GetDataQuerySettings()).ExtractValueSync();
387403

@@ -390,12 +406,16 @@ Y_UNIT_TEST_SUITE(KqpCost) {
390406
auto stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
391407

392408
TMap<TString, ui64> readsByTable;
393-
for(const auto& queryPhase : stats.query_phases()) {
409+
for (const auto& queryPhase : stats.query_phases()) {
394410
for(const auto& tableAccess: queryPhase.table_access()) {
395411
readsByTable[tableAccess.name()] += tableAccess.reads().rows();
396412
}
397413
}
398414

415+
for (const auto& kvp : readsByTable) {
416+
debugPringTable(kvp.first);
417+
}
418+
399419
UNIT_ASSERT_VALUES_EQUAL_C(expectedReadsByTable, readsByTable, query);
400420
};
401421

@@ -450,8 +470,8 @@ Y_UNIT_TEST_SUITE(KqpCost) {
450470
LIMIT 10;
451471
)"), {
452472
{"/Root/Vectors/vector_idx/indexImplLevelTable", 10},
453-
{"/Root/Vectors/vector_idx/indexImplPostingTable", 32},
454-
{"/Root/Vectors", 32}
473+
{"/Root/Vectors/vector_idx/indexImplPostingTable", 24},
474+
{"/Root/Vectors", 24}
455475
});
456476
}
457477

@@ -485,7 +505,7 @@ Y_UNIT_TEST_SUITE(KqpCost) {
485505
LIMIT 10;
486506
)"), {
487507
{"/Root/Vectors/vector_idx_covered/indexImplLevelTable", 10},
488-
{"/Root/Vectors/vector_idx_covered/indexImplPostingTable", 32},
508+
{"/Root/Vectors/vector_idx_covered/indexImplPostingTable", 24},
489509
});
490510
}
491511

@@ -499,8 +519,8 @@ Y_UNIT_TEST_SUITE(KqpCost) {
499519
)"), {
500520
{"/Root/Vectors/vector_idx_prefixed/indexImplPrefixTable", 1},
501521
{"/Root/Vectors/vector_idx_prefixed/indexImplLevelTable", 4},
502-
{"/Root/Vectors/vector_idx_prefixed/indexImplPostingTable", 2}, // about rows / 10 / clusters^levels = 100 / 10 / 2^3 = 1.25
503-
{"/Root/Vectors", 2} // same as posting
522+
{"/Root/Vectors/vector_idx_prefixed/indexImplPostingTable", 4}, // about rows / 10 / clusters^levels = 100 / 10 / 2^2 = 2.5
523+
{"/Root/Vectors", 4} // same as posting
504524
});
505525
}
506526
}

‎ydb/core/tx/schemeshard/schemeshard_build_index__progress.cpp

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222
namespace NKikimr {
2323
namespace NSchemeShard {
2424

25+
ui64 gVectorIndexSeed = 0;
26+
2527
// return count, parts, step
2628
static std::tuple<NTableIndex::TClusterId, NTableIndex::TClusterId, NTableIndex::TClusterId> ComputeKMeansBoundaries(const NSchemeShard::TTableInfo& tableInfo, const TIndexBuildInfo& buildInfo) {
2729
const auto& kmeans = buildInfo.KMeans;
@@ -512,7 +514,7 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
512514
TMap<TTabletId, THolder<IEventBase>> ToTabletSend;
513515

514516
template <bool WithSnapshot = true, typename TRequest>
515-
TTabletId CommonFillScanRequest(TRequest& request, TShardIdx shardIdx, TIndexBuildInfo& buildInfo) {
517+
TTabletId FillScanRequestCommon(TRequest& request, TShardIdx shardIdx, TIndexBuildInfo& buildInfo) {
516518
TTabletId shardId = Self->ShardInfos.at(shardIdx).TabletID;
517519
request.SetTabletId(ui64(shardId));
518520

@@ -543,6 +545,13 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
543545
return shardId;
544546
}
545547

548+
template <typename TRequest>
549+
void FillScanRequestSeed(TRequest& request) {
550+
request.SetSeed(gVectorIndexSeed
551+
? gVectorIndexSeed
552+
: request.GetTabletId());
553+
}
554+
546555
void SendSampleKRequest(TShardIdx shardIdx, TIndexBuildInfo& buildInfo) {
547556
Y_ENSURE(buildInfo.IsBuildVectorIndex());
548557
auto ev = MakeHolder<TEvDataShard::TEvSampleKRequest>();
@@ -566,8 +575,8 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
566575

567576
ev->Record.AddColumns(buildInfo.IndexColumns.back());
568577

569-
auto shardId = CommonFillScanRequest(ev->Record, shardIdx, buildInfo);
570-
ev->Record.SetSeed(ui64(shardId));
578+
auto shardId = FillScanRequestCommon(ev->Record, shardIdx, buildInfo);
579+
FillScanRequestSeed(ev->Record);
571580
LOG_N("TTxBuildProgress: TEvSampleKRequest: " << ev->Record.ShortDebugString());
572581

573582
ToTabletSend.emplace(shardId, std::move(ev));
@@ -606,7 +615,7 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
606615
buildInfo.DataColumns.begin(), buildInfo.DataColumns.end()
607616
};
608617

609-
auto shardId = CommonFillScanRequest(ev->Record, shardIdx, buildInfo);
618+
auto shardId = FillScanRequestCommon(ev->Record, shardIdx, buildInfo);
610619
LOG_N("TTxBuildProgress: TEvReshuffleKMeansRequest: " << ToShortDebugString(ev->Record));
611620

612621
ToTabletSend.emplace(shardId, std::move(ev));
@@ -638,7 +647,7 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
638647

639648
ev->Record.SetEmbeddingColumn(buildInfo.IndexColumns.back());
640649

641-
auto shardId = CommonFillScanRequest(ev->Record, shardIdx, buildInfo);
650+
auto shardId = FillScanRequestCommon(ev->Record, shardIdx, buildInfo);
642651
LOG_N("TTxBuildProgress: TEvRecomputeKMeansRequest: " << ToShortDebugString(ev->Record));
643652

644653
ToTabletSend.emplace(shardId, std::move(ev));
@@ -686,8 +695,8 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
686695
buildInfo.DataColumns.begin(), buildInfo.DataColumns.end()
687696
};
688697

689-
auto shardId = CommonFillScanRequest(ev->Record, shardIdx, buildInfo);
690-
ev->Record.SetSeed(ui64(shardId));
698+
auto shardId = FillScanRequestCommon(ev->Record, shardIdx, buildInfo);
699+
FillScanRequestSeed(ev->Record);
691700
LOG_N("TTxBuildProgress: TEvLocalKMeansRequest: " << ev->Record.ShortDebugString());
692701

693702
ToTabletSend.emplace(shardId, std::move(ev));
@@ -731,8 +740,8 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
731740
ev->Record.AddSourcePrimaryKeyColumns(tableInfo.Columns.at(keyPos).Name);
732741
}
733742

734-
auto shardId = CommonFillScanRequest<false>(ev->Record, shardIdx, buildInfo);
735-
ev->Record.SetSeed(ui64(shardId));
743+
auto shardId = FillScanRequestCommon<false>(ev->Record, shardIdx, buildInfo);
744+
FillScanRequestSeed(ev->Record);
736745
LOG_N("TTxBuildProgress: TEvPrefixKMeansRequest: " << ev->Record.ShortDebugString());
737746

738747
ToTabletSend.emplace(shardId, std::move(ev));
@@ -780,7 +789,7 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
780789

781790
ev->Record.SetTargetName(buildInfo.TargetName);
782791

783-
auto shardId = CommonFillScanRequest(ev->Record, shardIdx, buildInfo);
792+
auto shardId = FillScanRequestCommon(ev->Record, shardIdx, buildInfo);
784793

785794
LOG_N("TTxBuildProgress: TEvBuildIndexCreateRequest: " << ev->Record.ShortDebugString());
786795

‎ydb/core/tx/schemeshard/schemeshard_impl.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@ namespace NKikimr {
8282
namespace NSchemeShard {
8383

8484
extern const ui64 NEW_TABLE_ALTER_VERSION;
85+
extern ui64 gVectorIndexSeed; // for tests only
8586

8687
class TDataErasureManager;
8788

0 commit comments

Comments
 (0)