Skip to content

Commit 318bca5

Browse files
MBkktMBkkt
committed
WIP
1 parent eb88635 commit 318bca5

7 files changed

+199
-10
lines changed

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

+4
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,10 @@ void TSchemeShard::Handle(TEvDataShard::TEvReshuffleKMeansResponse::TPtr& ev, co
3636
Execute(CreateTxReply(ev), ctx);
3737
}
3838

39+
void TSchemeShard::Handle(TEvDataShard::TEvLocalKMeansResponse::TPtr& ev, const TActorContext& ctx) {
40+
Execute(CreateTxReply(ev), ctx);
41+
}
42+
3943
void TSchemeShard::Handle(TEvIndexBuilder::TEvUploadSampleKResponse::TPtr& ev, const TActorContext& ctx) {
4044
Execute(CreateTxReply(ev), ctx);
4145
}

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

+187-7
Original file line numberDiff line numberDiff line change
@@ -584,6 +584,48 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
584584
ToTabletSend.emplace_back(shardId, ui64(BuildId), std::move(ev));
585585
}
586586

587+
void SendKMeansLocalRequest(TShardIdx shardIdx, TIndexBuildInfo& buildInfo) {
588+
Y_ASSERT(buildInfo.IsBuildVectorIndex());
589+
auto ev = MakeHolder<TEvDataShard::TEvLocalKMeansRequest>();
590+
ev->Record.SetId(ui64(BuildId));
591+
592+
PathIdFromPathId(buildInfo.TablePathId, ev->Record.MutablePathId());
593+
*ev->Record.MutableSettings() = std::get<1>(buildInfo.SpecializedIndexDescription).GetSettings();
594+
ev->Record.SetK(buildInfo.KMeans.K);
595+
ev->Record.SetUpload(buildInfo.KMeans.GetUpload());
596+
ev->Record.SetState(NKikimrTxDataShard::TEvLocalKMeansRequest::SAMPLE);
597+
598+
ev->Record.SetDoneRounds(0);
599+
ev->Record.SetNeedsRounds(3); // TODO(mbkkt) should be configurable
600+
601+
ev->Record.SetParent(buildInfo.KMeans.Parent);
602+
ev->Record.SetChild(buildInfo.KMeans.ChildBegin);
603+
604+
auto path = TPath::Init(buildInfo.TablePathId, Self).Dive(buildInfo.IndexName);
605+
if (buildInfo.KMeans.GetUpload() == NKikimrTxDataShard::TEvLocalKMeansRequest::UPLOAD_MAIN_TO_POSTING
606+
|| buildInfo.KMeans.GetUpload() == NKikimrTxDataShard::TEvLocalKMeansRequest::UPLOAD_BUILD_TO_POSTING) {
607+
path.Dive(NTableIndex::NTableVectorKmeansTreeIndex::PostingTable);
608+
} else if (buildInfo.KMeans.Level % 2 == 0) {
609+
path.Dive(NTableIndex::NTableVectorKmeansTreeIndex::BuildPostingTableSuffix0);
610+
} else {
611+
path.Dive(NTableIndex::NTableVectorKmeansTreeIndex::BuildPostingTableSuffix1);
612+
}
613+
ev->Record.SetPostingName(path.PathString());
614+
path.Rise().Dive(NTableIndex::NTableVectorKmeansTreeIndex::LevelTable);
615+
ev->Record.SetLevelName(path.PathString());
616+
617+
ev->Record.SetEmbeddingColumn(buildInfo.IndexColumns[0]);
618+
*ev->Record.MutableDataColumns() = {
619+
buildInfo.DataColumns.begin(), buildInfo.DataColumns.end()
620+
};
621+
622+
auto shardId = CommonFillRecord(ev->Record, shardIdx, buildInfo);
623+
ev->Record.SetSeed(ui64(shardId));
624+
LOG_D("TTxBuildProgress: TEvLocalKMeansRequest: " << ev->Record.ShortDebugString());
625+
626+
ToTabletSend.emplace_back(shardId, ui64(BuildId), std::move(ev));
627+
}
628+
587629
void SendBuildIndexRequest(TShardIdx shardIdx, TIndexBuildInfo& buildInfo) {
588630
auto ev = MakeHolder<TEvDataShard::TEvBuildIndexCreateRequest>();
589631
ev->Record.SetBuildIndexId(ui64(BuildId));
@@ -719,10 +761,9 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
719761
Y_ABORT("Unreachable");
720762
}
721763
}
722-
// TODO(mbkkt) enable detection of Local case
723-
// if (buildInfo.ToUploadShards.size() == 1 && buildInfo.DoneShardsSize == 0) {
724-
// buildInfo.KMeans.State = TIndexBuildInfo::TKMeans::Local;
725-
// }
764+
if (buildInfo.DoneShardsSize == 0 && buildInfo.ToUploadShards.size() == 1) {
765+
buildInfo.KMeans.State = TIndexBuildInfo::TKMeans::Local;
766+
}
726767
}
727768

728769
bool SendKMeansSample(TIndexBuildInfo& buildInfo) {
@@ -739,6 +780,10 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
739780
return SendToShards(buildInfo, [&](TShardIdx shardIdx) { SendKMeansReshuffleRequest(shardIdx, buildInfo); });
740781
}
741782

783+
bool SendKMeansLocal(TIndexBuildInfo& buildInfo) {
784+
return SendToShards(buildInfo, [&](TShardIdx shardIdx) { SendKMeansLocalRequest(shardIdx, buildInfo); });
785+
}
786+
742787
bool SendVectorIndex(TIndexBuildInfo& buildInfo) {
743788
switch (buildInfo.KMeans.State) {
744789
case TIndexBuildInfo::TKMeans::Sample:
@@ -748,9 +793,8 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
748793
// return SendKMeansRecompute(buildInfo);
749794
case TIndexBuildInfo::TKMeans::Reshuffle:
750795
return SendKMeansReshuffle(buildInfo);
751-
// TODO(mbkkt)
752-
// case TIndexBuildInfo::TKMeans::Local:
753-
// return SendKMeansLocal(buildInfo);
796+
case TIndexBuildInfo::TKMeans::Local:
797+
return SendKMeansLocal(buildInfo);
754798
}
755799
return true;
756800
}
@@ -1344,6 +1388,138 @@ struct TSchemeShard::TIndexBuilder::TTxReplySampleK: public TSchemeShard::TIndex
13441388
}
13451389
};
13461390

1391+
struct TSchemeShard::TIndexBuilder::TTxReplyLocalKMeans: public TSchemeShard::TIndexBuilder::TTxReply {
1392+
private:
1393+
TEvDataShard::TEvLocalKMeansResponse::TPtr Local;
1394+
1395+
public:
1396+
explicit TTxReplyLocalKMeans(TSelf* self, TEvDataShard::TEvLocalKMeansResponse::TPtr& local)
1397+
: TTxReply(self)
1398+
, Local(local)
1399+
{
1400+
}
1401+
1402+
bool DoExecute(TTransactionContext& txc, const TActorContext& ctx) override {
1403+
auto& record = Local->Get()->Record;
1404+
1405+
LOG_I("TTxReply : TEvLocalKMeansResponse, Id# " << record.GetId());
1406+
1407+
const auto buildId = TIndexBuildId(record.GetId());
1408+
const auto* buildInfoPtr = Self->IndexBuilds.FindPtr(buildId);
1409+
if (!buildInfoPtr) {
1410+
return true;
1411+
}
1412+
auto& buildInfo = *buildInfoPtr->Get();
1413+
LOG_D("TTxReply : TEvLocalKMeansResponse"
1414+
<< ", TIndexBuildInfo: " << buildInfo
1415+
<< ", record: " << record.ShortDebugString());
1416+
1417+
TTabletId shardId = TTabletId(record.GetTabletId());
1418+
if (!Self->TabletIdToShardIdx.contains(shardId)) {
1419+
return true;
1420+
}
1421+
1422+
TShardIdx shardIdx = Self->TabletIdToShardIdx.at(shardId);
1423+
if (!buildInfo.Shards.contains(shardIdx)) {
1424+
return true;
1425+
}
1426+
1427+
switch (const auto state = buildInfo.State; state) {
1428+
case TIndexBuildInfo::EState::Filling:
1429+
{
1430+
TIndexBuildInfo::TShardStatus& shardStatus = buildInfo.Shards.at(shardIdx);
1431+
1432+
auto actualSeqNo = std::pair<ui64, ui64>(Self->Generation(), shardStatus.SeqNoRound);
1433+
auto recordSeqNo = std::pair<ui64, ui64>(record.GetRequestSeqNoGeneration(), record.GetRequestSeqNoRound());
1434+
1435+
if (actualSeqNo != recordSeqNo) {
1436+
LOG_D("TTxReply : TEvLocalKMeansResponse"
1437+
<< " ignore progress message by seqNo"
1438+
<< ", TIndexBuildInfo: " << buildInfo
1439+
<< ", actual seqNo for the shard " << shardId << " (" << shardIdx << ") is: " << Self->Generation() << ":" << shardStatus.SeqNoRound
1440+
<< ", record: " << record.ShortDebugString());
1441+
Y_ABORT_UNLESS(actualSeqNo > recordSeqNo);
1442+
return true;
1443+
}
1444+
1445+
// TODO(mbkkt) add billing
1446+
1447+
NYql::TIssues issues;
1448+
NYql::IssuesFromMessage(record.GetIssues(), issues);
1449+
shardStatus.DebugMessage = issues.ToString();
1450+
1451+
NIceDb::TNiceDb db(txc.DB);
1452+
Self->PersistBuildIndexUploadProgress(db, buildInfo, shardIdx);
1453+
shardStatus.Status = record.GetStatus();
1454+
1455+
switch (shardStatus.Status) {
1456+
case NKikimrIndexBuilder::EBuildStatus::DONE:
1457+
if (buildInfo.InProgressShards.erase(shardIdx)) {
1458+
++buildInfo.DoneShardsSize;
1459+
1460+
Self->IndexBuildPipes.Close(buildId, shardId, ctx);
1461+
1462+
Progress(buildId);
1463+
}
1464+
break;
1465+
1466+
case NKikimrIndexBuilder::EBuildStatus::ABORTED:
1467+
// datashard gracefully rebooted, reschedule shard
1468+
if (buildInfo.InProgressShards.erase(shardIdx)) {
1469+
buildInfo.ToUploadShards.emplace_front(shardIdx);
1470+
1471+
Self->IndexBuildPipes.Close(buildId, shardId, ctx);
1472+
1473+
Progress(buildId);
1474+
}
1475+
break;
1476+
1477+
case NKikimrIndexBuilder::EBuildStatus::BUILD_ERROR:
1478+
case NKikimrIndexBuilder::EBuildStatus::BAD_REQUEST:
1479+
buildInfo.Issue += TStringBuilder()
1480+
<< "One of the shards report "<< shardStatus.Status
1481+
<< " at Filling stage, process has to be canceled"
1482+
<< ", shardId: " << shardId
1483+
<< ", shardIdx: " << shardIdx;
1484+
Self->PersistBuildIndexIssue(db, buildInfo);
1485+
ChangeState(buildInfo.Id, TIndexBuildInfo::EState::Rejection_Applying);
1486+
1487+
Progress(buildId);
1488+
break;
1489+
1490+
case NKikimrIndexBuilder::EBuildStatus::INVALID:
1491+
case NKikimrIndexBuilder::EBuildStatus::ACCEPTED:
1492+
case NKikimrIndexBuilder::EBuildStatus::IN_PROGRESS:
1493+
Y_ABORT("Unreachable");
1494+
}
1495+
break;
1496+
}
1497+
case TIndexBuildInfo::EState::AlterMainTable:
1498+
case TIndexBuildInfo::EState::Invalid:
1499+
case TIndexBuildInfo::EState::Locking:
1500+
case TIndexBuildInfo::EState::GatheringStatistics:
1501+
case TIndexBuildInfo::EState::Initiating:
1502+
case TIndexBuildInfo::EState::DropBuild:
1503+
case TIndexBuildInfo::EState::CreateBuild:
1504+
case TIndexBuildInfo::EState::Applying:
1505+
case TIndexBuildInfo::EState::Unlocking:
1506+
case TIndexBuildInfo::EState::Done:
1507+
Y_FAIL_S("Unreachable " << Name(state));
1508+
case TIndexBuildInfo::EState::Cancellation_Applying:
1509+
case TIndexBuildInfo::EState::Cancellation_Unlocking:
1510+
case TIndexBuildInfo::EState::Cancelled:
1511+
case TIndexBuildInfo::EState::Rejection_Applying:
1512+
case TIndexBuildInfo::EState::Rejection_Unlocking:
1513+
case TIndexBuildInfo::EState::Rejected:
1514+
LOG_D("TTxReply : TEvLocalKMeansResponse"
1515+
<< " superfluous message " << record.ShortDebugString());
1516+
break;
1517+
}
1518+
1519+
return true;
1520+
}
1521+
};
1522+
13471523
struct TSchemeShard::TIndexBuilder::TTxReplyReshuffleKMeans: public TSchemeShard::TIndexBuilder::TTxReply {
13481524
private:
13491525
TEvDataShard::TEvReshuffleKMeansResponse::TPtr Reshuffle;
@@ -2118,6 +2294,10 @@ ITransaction* TSchemeShard::CreateTxReply(TEvDataShard::TEvReshuffleKMeansRespon
21182294
return new TIndexBuilder::TTxReplyReshuffleKMeans(this, reshuffle);
21192295
}
21202296

2297+
ITransaction* TSchemeShard::CreateTxReply(TEvDataShard::TEvLocalKMeansResponse::TPtr& local) {
2298+
return new TIndexBuilder::TTxReplyLocalKMeans(this, local);
2299+
}
2300+
21212301
ITransaction* TSchemeShard::CreateTxReply(TEvIndexBuilder::TEvUploadSampleKResponse::TPtr& upload) {
21222302
return new TIndexBuilder::TTxReplyUpload(this, upload);
21232303
}

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

+1
Original file line numberDiff line numberDiff line change
@@ -4752,6 +4752,7 @@ void TSchemeShard::StateWork(STFUNC_SIG) {
47524752
HFuncTraced(TEvPrivate::TEvIndexBuildingMakeABill, Handle);
47534753
HFuncTraced(TEvDataShard::TEvSampleKResponse, Handle);
47544754
HFuncTraced(TEvDataShard::TEvReshuffleKMeansResponse, Handle);
4755+
HFuncTraced(TEvDataShard::TEvLocalKMeansResponse, Handle);
47554756
HFuncTraced(TEvIndexBuilder::TEvUploadSampleKResponse, Handle);
47564757
// } // NIndexBuilder
47574758

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

+3
Original file line numberDiff line numberDiff line change
@@ -1344,6 +1344,7 @@ class TSchemeShard
13441344
struct TTxReplyRetry;
13451345
struct TTxReplySampleK;
13461346
struct TTxReplyReshuffleKMeans;
1347+
struct TTxReplyLocalKMeans;
13471348
struct TTxReplyUpload;
13481349

13491350
struct TTxPipeReset;
@@ -1362,6 +1363,7 @@ class TSchemeShard
13621363
NTabletFlatExecutor::ITransaction* CreateTxReply(TEvDataShard::TEvBuildIndexProgressResponse::TPtr& progress);
13631364
NTabletFlatExecutor::ITransaction* CreateTxReply(TEvDataShard::TEvSampleKResponse::TPtr& sampleK);
13641365
NTabletFlatExecutor::ITransaction* CreateTxReply(TEvDataShard::TEvReshuffleKMeansResponse::TPtr& reshuffle);
1366+
NTabletFlatExecutor::ITransaction* CreateTxReply(TEvDataShard::TEvLocalKMeansResponse::TPtr& local);
13651367
NTabletFlatExecutor::ITransaction* CreateTxReply(TEvIndexBuilder::TEvUploadSampleKResponse::TPtr& upload);
13661368
NTabletFlatExecutor::ITransaction* CreatePipeRetry(TIndexBuildId indexBuildId, TTabletId tabletId);
13671369
NTabletFlatExecutor::ITransaction* CreateTxBilling(TEvPrivate::TEvIndexBuildingMakeABill::TPtr& ev);
@@ -1375,6 +1377,7 @@ class TSchemeShard
13751377
void Handle(TEvDataShard::TEvBuildIndexProgressResponse::TPtr& ev, const TActorContext& ctx);
13761378
void Handle(TEvDataShard::TEvSampleKResponse::TPtr& ev, const TActorContext& ctx);
13771379
void Handle(TEvDataShard::TEvReshuffleKMeansResponse::TPtr& ev, const TActorContext& ctx);
1380+
void Handle(TEvDataShard::TEvLocalKMeansResponse::TPtr& ev, const TActorContext& ctx);
13781381
void Handle(TEvIndexBuilder::TEvUploadSampleKResponse::TPtr& ev, const TActorContext& ctx);
13791382

13801383
void Handle(TEvPrivate::TEvIndexBuildingMakeABill::TPtr& ev, const TActorContext& ctx);

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -2997,7 +2997,7 @@ struct TIndexBuildInfo: public TSimpleRefCount<TIndexBuildInfo> {
29972997
Sample = 0,
29982998
// Recompute,
29992999
Reshuffle,
3000-
// Local,
3000+
Local,
30013001
};
30023002
ui32 Level = 0;
30033003

‎ydb/core/tx/schemeshard/ut_helpers/helpers.h

+1
Original file line numberDiff line numberDiff line change
@@ -371,6 +371,7 @@ namespace NSchemeShardUT_Private {
371371
TVector<TString> IndexColumns;
372372
TVector<TString> DataColumns;
373373
TVector<NYdb::NTable::TGlobalIndexSettings> GlobalIndexSettings = {};
374+
std::variant<std::monostate, NKikimrSchemeOp::TVectorIndexKmeansTreeDescription> SpecializedIndexDescription;
374375
};
375376

376377
std::unique_ptr<TEvIndexBuilder::TEvCreateRequest> CreateBuildColumnRequest(ui64 id, const TString& dbName, const TString& src, const TString& columnName, const Ydb::TypedValue& literal);

‎ydb/core/tx/schemeshard/ut_index_build/ut_vector_index_build.cpp

+2-2
Original file line numberDiff line numberDiff line change
@@ -96,8 +96,8 @@ Y_UNIT_TEST_SUITE (VectorIndexBuildTest) {
9696
UNIT_ASSERT_VALUES_EQUAL(err, "");
9797
UNIT_ASSERT_VALUES_EQUAL(status, NKikimrProto::EReplyStatus::OK);
9898
};
99-
for (ui32 key = 0; key < 200; ++key) {
100-
fnWriteRow(TTestTxConfig::FakeHiveTablets + 6, key, std::to_string(key), "Table");
99+
for (ui32 key = 0; key < 128; ++key) {
100+
fnWriteRow(TTestTxConfig::FakeHiveTablets + 6, key, std::to_string(key) + "\3", "Table");
101101
}
102102

103103
runtime.SetLogPriority(NKikimrServices::TX_DATASHARD, NLog::PRI_TRACE);

0 commit comments

Comments
 (0)