Merge 44f8ee3 into c1fc810

ydb-platform · Sep 25, 2024 · e4126db · e4126db
2 parents c1fc810 + 44f8ee3
commit e4126db
Show file tree

Hide file tree

Showing 13 changed files with 1,637 additions and 336 deletions.
diff --git a/ydb/core/protos/tx_datashard.proto b/ydb/core/protos/tx_datashard.proto
@@ -1522,7 +1522,7 @@ message TEvLocalKMeansRequest {
 
     // id of parent cluster
     optional uint32 Parent = 15;
-    // [Child ... Child + K] ids reserved for our clusters
+    // [Child ... Child + K) ids reserved for this kmeans clusters
     optional uint32 Child = 16;
 
     optional string LevelName = 17;
@@ -1532,7 +1532,7 @@ message TEvLocalKMeansRequest {
     repeated string DataColumns = 20;
 }
 
-message TEvLocalKMeansProgressResponse {
+message TEvLocalKMeansResponse {
     optional uint64 Id = 1;
 
     optional uint64 TabletId = 2;
@@ -1552,6 +1552,53 @@ message TEvLocalKMeansProgressResponse {
     // optional uint32 DoneRounds = 11;
 }
 
+message TEvReshuffleKMeansRequest {
+    optional uint64 Id = 1;
+
+    optional uint64 TabletId = 2;
+    optional NKikimrProto.TPathID PathId = 3;
+
+    optional uint64 SnapshotTxId = 4;
+    optional uint64 SnapshotStep = 5;
+
+    optional uint64 SeqNoGeneration = 6;
+    optional uint64 SeqNoRound = 7;
+
+    optional Ydb.Table.VectorIndexSettings Settings = 8;
+
+    optional TEvLocalKMeansRequest.EState Upload = 9;
+
+    // id of parent cluster
+    optional uint32 Parent = 10;
+    // [Child ... Child + ClustersSize) ids of this kmeans clusters
+    optional uint32 Child = 11;
+    // centroids of clusters
+    repeated string Clusters = 12;
+
+    optional string PostingName = 13;
+
+    optional string EmbeddingColumn = 14;
+    repeated string DataColumns = 15;
+}
+
+message TEvReshuffleKMeansResponse {
+    optional uint64 Id = 1;
+
+    optional uint64 TabletId = 2;
+    optional NKikimrProto.TPathID PathId = 3;
+
+    optional uint64 RequestSeqNoGeneration = 4;
+    optional uint64 RequestSeqNoRound = 5;
+
+    optional NKikimrIndexBuilder.EBuildStatus Status = 6;
+    repeated Ydb.Issue.IssueMessage Issues = 7;
+
+    // TODO(mbkkt) implement slow-path (reliable-path)
+    // optional uint64 RowsDelta = 8;
+    // optional uint64 BytesDelta = 9;
+    // optional last written primary key
+}
+
 message TEvCdcStreamScanRequest {
     message TLimits {
         optional uint32 BatchMaxBytes = 1 [default = 512000];

diff --git a/ydb/core/tx/datashard/buffer_data.h b/ydb/core/tx/datashard/buffer_data.h
@@ -1,3 +1,5 @@
+#pragma once
+
 #include "ydb/core/scheme/scheme_tablecell.h"
 #include "ydb/core/tx/datashard/upload_stats.h"
 #include "ydb/core/tx/tx_proxy/upload_rows.h"

diff --git a/ydb/core/tx/datashard/datashard.h b/ydb/core/tx/datashard/datashard.h
@@ -333,7 +333,10 @@ struct TEvDataShard {
         EvSampleKResponse,
 
         EvLocalKMeansRequest,
-        EvLocalKMeansProgressResponse,
+        EvLocalKMeansResponse,
+
+        EvReshuffleKMeansRequest,
+        EvReshuffleKMeansResponse,
 
         EvEnd
     };
@@ -1457,16 +1460,28 @@ struct TEvDataShard {
                           TEvDataShard::EvSampleKResponse> {
     };
 
+    struct TEvReshuffleKMeansRequest
+        : public TEventPB<TEvReshuffleKMeansRequest,
+                          NKikimrTxDataShard::TEvReshuffleKMeansRequest,
+                          TEvDataShard::EvReshuffleKMeansRequest> {
+    };
+
+    struct TEvReshuffleKMeansResponse
+        : public TEventPB<TEvReshuffleKMeansResponse,
+                          NKikimrTxDataShard::TEvReshuffleKMeansResponse,
+                          TEvDataShard::EvReshuffleKMeansResponse> {
+    };
+
     struct TEvLocalKMeansRequest
         : public TEventPB<TEvLocalKMeansRequest,
                           NKikimrTxDataShard::TEvLocalKMeansRequest,
                           TEvDataShard::EvLocalKMeansRequest> {
     };
 
-    struct TEvLocalKMeansProgressResponse
-        : public TEventPB<TEvLocalKMeansProgressResponse,
-                          NKikimrTxDataShard::TEvLocalKMeansProgressResponse,
-                          TEvDataShard::EvLocalKMeansProgressResponse> {
+    struct TEvLocalKMeansResponse
+        : public TEventPB<TEvLocalKMeansResponse,
+                          NKikimrTxDataShard::TEvLocalKMeansResponse,
+                          TEvDataShard::EvLocalKMeansResponse> {
     };
 
     struct TEvKqpScan

diff --git a/ydb/core/tx/datashard/datashard_impl.h b/ydb/core/tx/datashard/datashard_impl.h
@@ -254,6 +254,7 @@ class TDataShard
     class TTxHandleSafeBuildIndexScan;
     class TTxHandleSafeSampleKScan;
     class TTxHandleSafeLocalKMeansScan;
+    class TTxHandleSafeReshuffleKMeansScan;
     class TTxHandleSafeStatisticsScan;
 
     class TTxMediatorStateRestored;
@@ -1325,6 +1326,8 @@ class TDataShard
     void HandleSafe(TEvDataShard::TEvBuildIndexCreateRequest::TPtr& ev, const TActorContext& ctx);
     void Handle(TEvDataShard::TEvSampleKRequest::TPtr& ev, const TActorContext& ctx);
     void HandleSafe(TEvDataShard::TEvSampleKRequest::TPtr& ev, const TActorContext& ctx);
+    void Handle(TEvDataShard::TEvReshuffleKMeansRequest::TPtr& ev, const TActorContext& ctx);
+    void HandleSafe(TEvDataShard::TEvReshuffleKMeansRequest::TPtr& ev, const TActorContext& ctx);
     void Handle(TEvDataShard::TEvLocalKMeansRequest::TPtr& ev, const TActorContext& ctx);
     void HandleSafe(TEvDataShard::TEvLocalKMeansRequest::TPtr& ev, const TActorContext& ctx);
     void Handle(TEvDataShard::TEvCdcStreamScanRequest::TPtr& ev, const TActorContext& ctx);
@@ -3117,6 +3120,7 @@ class TDataShard
             HFunc(TEvDataShard::TEvDiscardVolatileSnapshotRequest, Handle);
             HFuncTraced(TEvDataShard::TEvBuildIndexCreateRequest, Handle);
             HFunc(TEvDataShard::TEvSampleKRequest, Handle);
+            HFunc(TEvDataShard::TEvReshuffleKMeansRequest, Handle);
             HFunc(TEvDataShard::TEvLocalKMeansRequest, Handle);
             HFunc(TEvDataShard::TEvCdcStreamScanRequest, Handle);
             HFunc(TEvPrivate::TEvCdcStreamScanRegistered, Handle);

diff --git a/ydb/core/tx/datashard/datashard_ut_local_kmeans.cpp b/ydb/core/tx/datashard/datashard_ut_local_kmeans.cpp
@@ -81,7 +81,7 @@ Y_UNIT_TEST_SUITE (TTxDataShardLocalKMeansScan) {
             runtime.SendToPipe(tid, sender, ev.release(), 0, GetPipeConfigWithRetries());
 
             TAutoPtr<IEventHandle> handle;
-            auto reply = runtime.GrabEdgeEventRethrow<TEvDataShard::TEvLocalKMeansProgressResponse>(handle);
+            auto reply = runtime.GrabEdgeEventRethrow<TEvDataShard::TEvLocalKMeansResponse>(handle);
             UNIT_ASSERT_VALUES_EQUAL(reply->Record.GetStatus(), NKikimrIndexBuilder::EBuildStatus::BAD_REQUEST);
         }
     }
@@ -148,7 +148,7 @@ Y_UNIT_TEST_SUITE (TTxDataShardLocalKMeansScan) {
             runtime.SendToPipe(tid, sender, ev2.release(), 0, GetPipeConfigWithRetries());
 
             TAutoPtr<IEventHandle> handle;
-            auto reply = runtime.GrabEdgeEventRethrow<TEvDataShard::TEvLocalKMeansProgressResponse>(handle);
+            auto reply = runtime.GrabEdgeEventRethrow<TEvDataShard::TEvLocalKMeansResponse>(handle);
 
             NYql::TIssues issues;
             NYql::IssuesFromMessage(reply->Record.GetIssues(), issues);