Add reshuffle kmeans actor scan (#9725)

Author: MBkkt

ydb/core/protos/tx_datashard.proto

@@ -1524,7 +1524,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;
@@ -1534,7 +1534,7 @@ message TEvLocalKMeansRequest {
     repeated string DataColumns = 20;
 }
 
-message TEvLocalKMeansProgressResponse {
+message TEvLocalKMeansResponse {
     optional uint64 Id = 1;
 
     optional uint64 TabletId = 2;
@@ -1554,6 +1554,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];

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"

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

ydb/core/tx/datashard/datashard_impl.h

@@ -255,6 +255,7 @@ class TDataShard
     class TTxHandleSafeBuildIndexScan;
     class TTxHandleSafeSampleKScan;
     class TTxHandleSafeLocalKMeansScan;
+    class TTxHandleSafeReshuffleKMeansScan;
     class TTxHandleSafeStatisticsScan;
 
     class TTxMediatorStateRestored;
@@ -1327,6 +1328,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);
@@ -3132,6 +3135,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);