Prefixed vector index construction (#15166)

Author: MBkkt

ydb/core/kqp/opt/logical/kqp_opt_log_indexes.cpp

@@ -150,7 +150,7 @@ bool CanUseVectorIndex(const TIndexDescription& indexDesc, const TExprBase& lamb
     // TODO(mbkkt) We need to account top.Count(), but not clear what to if it's value is runtime?
     auto checkMember = [&] (const TExprBase& expr) {
         auto member = expr.Maybe<TCoMember>();
-        return member && member.Cast().Name().Value() == indexDesc.KeyColumns[0];
+        return member && member.Cast().Name().Value() == indexDesc.KeyColumns.back();
     };
     auto checkUdf = [&] (const TExprBase& expr, bool checkMembers) {
         auto apply = expr.Maybe<TCoApply>();
@@ -452,7 +452,7 @@ TExprBase DoRewriteTopSortOverKMeansTree(
             auto apply = newLambda.Body().Cast<TCoApply>();
             for (auto arg : apply.Args()) {
                 auto oldMember = arg.Maybe<TCoMember>();
-                if (oldMember && oldMember.Cast().Name().Value() == indexDesc.KeyColumns[0]) {
+                if (oldMember && oldMember.Cast().Name().Value() == indexDesc.KeyColumns.back()) {
                     auto newMember = Build<TCoMember>(ctx, pos)
                         .Name().Build(NTableIndex::NTableVectorKmeansTreeIndex::CentroidColumn)
                         .Struct(oldMember.Cast().Struct())

ydb/core/kqp/ut/indexes/kqp_indexes_ut.cpp

@@ -1632,6 +1632,54 @@ message TEvReshuffleKMeansResponse {
     // optional last written primary key
 }
 
+message TEvPrefixKMeansRequest {
+    optional uint64 Id = 1;
+
+    optional uint64 TabletId = 2;
+    optional NKikimrProto.TPathID PathId = 3;
+    
+    optional uint64 SeqNoGeneration = 4;
+    optional uint64 SeqNoRound = 5;
+
+    optional Ydb.Table.VectorIndexSettings Settings = 6;
+
+    optional uint64 Seed = 7;
+
+    optional TEvLocalKMeansRequest.EState Upload = 8;
+
+    optional uint32 K = 9;
+    optional uint32 NeedsRounds = 10;
+
+    // [Child ... Child + (1 + TableSize) * ShardIndex]
+    optional uint64 Child = 11;
+
+    optional string LevelName = 12;
+    optional string PostingName = 13;
+    optional string PrefixName = 14;
+
+    optional string EmbeddingColumn = 15;
+    repeated string DataColumns = 16;
+    optional uint32 PrefixColumns = 17;
+}
+
+message TEvPrefixKMeansResponse {
+    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;
+
+    optional uint64 UploadRows = 8;
+    optional uint64 UploadBytes = 9;
+    optional uint64 ReadRows = 10;
+    optional uint64 ReadBytes = 11;
+}
+
 message TEvCdcStreamScanRequest {
     message TLimits {
         optional uint32 BatchMaxBytes = 1 [default = 512000];

ydb/core/protos/tx_datashard.proto

@@ -349,6 +349,9 @@ namespace TEvDataShard {
         EvForceDataCleanup,
         EvForceDataCleanupResult,
 
+        EvPrefixKMeansRequest,
+        EvPrefixKMeansResponse,
+
         EvEnd
     };
 
@@ -1514,6 +1517,18 @@ namespace TEvDataShard {
                           TEvDataShard::EvLocalKMeansResponse> {
     };
 
+    struct TEvPrefixKMeansRequest
+        : public TEventPB<TEvPrefixKMeansRequest,
+                          NKikimrTxDataShard::TEvPrefixKMeansRequest,
+                          TEvDataShard::EvPrefixKMeansRequest> {
+    };
+
+    struct TEvPrefixKMeansResponse
+        : public TEventPB<TEvPrefixKMeansResponse,
+                          NKikimrTxDataShard::TEvPrefixKMeansResponse,
+                          TEvDataShard::EvPrefixKMeansResponse> {
+    };
+
     struct TEvKqpScan
         : public TEventPB<TEvKqpScan,
                           NKikimrTxDataShard::TEvKqpScan,

ydb/core/tx/datashard/datashard.h

@@ -250,6 +250,7 @@ class TDataShard
     class TTxHandleSafeBuildIndexScan;
     class TTxHandleSafeSampleKScan;
     class TTxHandleSafeLocalKMeansScan;
+    class TTxHandleSafePrefixKMeansScan;
     class TTxHandleSafeReshuffleKMeansScan;
     class TTxHandleSafeStatisticsScan;
 
@@ -1335,6 +1336,8 @@ class TDataShard
     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::TEvPrefixKMeansRequest::TPtr& ev, const TActorContext& ctx);
+    void HandleSafe(TEvDataShard::TEvPrefixKMeansRequest::TPtr& ev, const TActorContext& ctx);
     void Handle(TEvDataShard::TEvCdcStreamScanRequest::TPtr& ev, const TActorContext& ctx);
     void Handle(TEvPrivate::TEvCdcStreamScanRegistered::TPtr& ev, const TActorContext& ctx);
     void Handle(TEvPrivate::TEvCdcStreamScanProgress::TPtr& ev, const TActorContext& ctx);
@@ -3206,6 +3209,7 @@ class TDataShard
             HFunc(TEvDataShard::TEvSampleKRequest, Handle);
             HFunc(TEvDataShard::TEvReshuffleKMeansRequest, Handle);
             HFunc(TEvDataShard::TEvLocalKMeansRequest, Handle);
+            HFunc(TEvDataShard::TEvPrefixKMeansRequest, Handle);
             HFunc(TEvDataShard::TEvCdcStreamScanRequest, Handle);
             HFunc(TEvPrivate::TEvCdcStreamScanRegistered, Handle);
             HFunc(TEvPrivate::TEvCdcStreamScanProgress, Handle);