Refactor FillVectorIndex state machine (#18579) (#19128)

vitalif · kunga · commit 637353f516f1 · 2025-06-09T14:37:31.000+03:00
diff --git a/ydb/core/tx/schemeshard/schemeshard_build_index__progress.cpp b/ydb/core/tx/schemeshard/schemeshard_build_index__progress.cpp
@@ -764,7 +764,6 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
             buildInfo.Sample.Rows, buildInfo.KMeans.Parent, buildInfo.KMeans.Child);
 
         TActivationContext::AsActorContext().MakeFor(Self->SelfId()).Register(actor);
-        buildInfo.Sample.State = TIndexBuildInfo::TSample::EState::Upload;
 
         LOG_N("TTxBuildProgress: TUploadSampleK: " << buildInfo);
     }
@@ -807,6 +806,10 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
         }
     }
 
+    bool NoShardsAdded(TIndexBuildInfo& buildInfo) {
+        return buildInfo.DoneShards.empty() && buildInfo.InProgressShards.empty() && buildInfo.ToUploadShards.empty();
+    }
+
     void AddAllShards(TIndexBuildInfo& buildInfo) {
         ToTabletSend.clear();
         Self->IndexBuildPipes.CloseAll(BuildId, Self->ActorContext());
@@ -816,10 +819,37 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
         }
     }
 
+    void AddGlobalShardsForCurrentParent(TIndexBuildInfo& buildInfo) {
+        Y_ENSURE(NoShardsAdded(buildInfo));
+        if (buildInfo.KMeans.Parent == 0) {
+            AddAllShards(buildInfo);
+            return;
+        }
+        auto it = buildInfo.Cluster2Shards.lower_bound(buildInfo.KMeans.Parent);
+        Y_ENSURE(it != buildInfo.Cluster2Shards.end());
+        if (it->second.Shards.size() > 1) {
+            for (const auto& idx : it->second.Shards) {
+                const auto& status = buildInfo.Shards.at(idx);
+                AddShard(buildInfo, idx, status);
+            }
+        }
+    }
+
+    void AddLocalClusters(TIndexBuildInfo& buildInfo) {
+        Y_ENSURE(NoShardsAdded(buildInfo));
+        for (const auto& [to, state] : buildInfo.Cluster2Shards) {
+            if (state.Shards.size() == 1) {
+                const auto* status = buildInfo.Shards.FindPtr(state.Shards[0]);
+                Y_ENSURE(status);
+                AddShard(buildInfo, state.Shards[0], *status);
+            }
+        }
+    }
+
     bool FillSecondaryIndex(TIndexBuildInfo& buildInfo) {
         LOG_D("FillSecondaryIndex Start");
         
-        if (buildInfo.DoneShards.empty() && buildInfo.ToUploadShards.empty() && buildInfo.InProgressShards.empty()) {
+        if (NoShardsAdded(buildInfo)) {
             AddAllShards(buildInfo);
         }
         auto done = SendToShards(buildInfo, [&](TShardIdx shardIdx) { SendBuildSecondaryIndexRequest(shardIdx, buildInfo); }) &&
@@ -833,65 +863,21 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
     }
 
     bool FillPrefixKMeans(TIndexBuildInfo& buildInfo) {
-        if (buildInfo.DoneShards.empty() && buildInfo.ToUploadShards.empty() && buildInfo.InProgressShards.empty()) {
+        if (NoShardsAdded(buildInfo)) {
             AddAllShards(buildInfo);
         }
         return SendToShards(buildInfo, [&](TShardIdx shardIdx) { SendPrefixKMeansRequest(shardIdx, buildInfo); }) &&
                buildInfo.DoneShards.size() == buildInfo.Shards.size();
     }
 
     bool FillLocalKMeans(TIndexBuildInfo& buildInfo) {
-        if (buildInfo.DoneShards.empty() && buildInfo.ToUploadShards.empty() && buildInfo.InProgressShards.empty()) {
+        if (NoShardsAdded(buildInfo)) {
             AddAllShards(buildInfo);
         }
         return SendToShards(buildInfo, [&](TShardIdx shardIdx) { SendKMeansLocalRequest(shardIdx, buildInfo); }) &&
                buildInfo.DoneShards.size() == buildInfo.Shards.size();
     }
 
-    bool InitSingleKMeans(TIndexBuildInfo& buildInfo) {
-        if (!buildInfo.DoneShards.empty() || !buildInfo.InProgressShards.empty() || !buildInfo.ToUploadShards.empty()) {
-            return false;
-        }
-        if (buildInfo.KMeans.State == TIndexBuildInfo::TKMeans::MultiLocal) {
-            InitMultiKMeans(buildInfo);
-            return false;
-        }
-        if (buildInfo.KMeans.Parent == 0) {
-            AddAllShards(buildInfo);
-        } else {
-            auto it = buildInfo.Cluster2Shards.lower_bound(buildInfo.KMeans.Parent);
-            Y_ENSURE(it != buildInfo.Cluster2Shards.end());
-            if (it->second.Local == InvalidShardIdx) {
-                for (const auto& idx : it->second.Global) {
-                    const auto& status = buildInfo.Shards.at(idx);
-                    AddShard(buildInfo, idx, status);
-                }
-            }
-        }
-        if (buildInfo.DoneShards.size() + buildInfo.ToUploadShards.size() <= 1) {
-            buildInfo.KMeans.State = TIndexBuildInfo::TKMeans::Local;
-        }
-        return true;
-    }
-
-    bool InitMultiKMeans(TIndexBuildInfo& buildInfo) {
-        if (buildInfo.Cluster2Shards.empty()) {
-            return false;
-        }
-        Y_ENSURE(buildInfo.KMeans.Parent != 0);
-        for (const auto& [to, state] : buildInfo.Cluster2Shards) {
-            if (const auto& [from, local, global] = state; local != InvalidShardIdx) {
-                if (const auto* status = buildInfo.Shards.FindPtr(local)) {
-                    AddShard(buildInfo, local, *status);
-                }
-            }
-        }
-        buildInfo.KMeans.State = TIndexBuildInfo::TKMeans::MultiLocal;
-        buildInfo.Cluster2Shards.clear();
-        Y_ENSURE(buildInfo.InProgressShards.empty());
-        return !buildInfo.ToUploadShards.empty();
-    }
-
     bool SendKMeansSample(TIndexBuildInfo& buildInfo) {
         if (buildInfo.Sample.MaxProbability == 0) {
             buildInfo.ToUploadShards.clear();
@@ -910,22 +896,6 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
         return SendToShards(buildInfo, [&](TShardIdx shardIdx) { SendKMeansLocalRequest(shardIdx, buildInfo); });
     }
 
-    bool SendVectorIndex(TIndexBuildInfo& buildInfo) {
-        switch (buildInfo.KMeans.State) {
-            case TIndexBuildInfo::TKMeans::Sample:
-                return SendKMeansSample(buildInfo);
-            // TODO(mbkkt)
-            // case TIndexBuildInfo::TKMeans::Recompute:
-            //     return SendKMeansRecompute(buildInfo);
-            case TIndexBuildInfo::TKMeans::Reshuffle:
-                return SendKMeansReshuffle(buildInfo);
-            case TIndexBuildInfo::TKMeans::Local:
-            case TIndexBuildInfo::TKMeans::MultiLocal:
-                return SendKMeansLocal(buildInfo);
-        }
-        return true;
-    }
-
     void ClearDoneShards(TTransactionContext& txc, TIndexBuildInfo& buildInfo) {
         if (buildInfo.DoneShards.empty()) {
             return;
@@ -966,6 +936,7 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
             // it's approximate but upper bound, so it's ok
             buildInfo.KMeans.TableSize = std::max<ui64>(1, buildInfo.Processed.GetUploadRows());
             buildInfo.KMeans.PrefixIndexDone(doneShards);
+            buildInfo.KMeans.State = TIndexBuildInfo::TKMeans::MultiLocal;
             LOG_D("FillPrefixedVectorIndex PrefixIndexDone " << buildInfo.DebugString());
 
             PersistKMeansState(txc, buildInfo);
@@ -1006,60 +977,93 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
     }
 
     bool FillVectorIndex(TTransactionContext& txc, TIndexBuildInfo& buildInfo) {
-        // FIXME: Very non-intuitive state machine, rework it by adding an explicit vector index fill state
         LOG_D("FillVectorIndex Start " << buildInfo.DebugString());
 
-        if (buildInfo.Sample.State == TIndexBuildInfo::TSample::EState::Upload) {
-            return false;
-        }
-        if (InitSingleKMeans(buildInfo)) {
-            LOG_D("FillVectorIndex SingleKMeans " << buildInfo.DebugString());
-        }
-        if (!SendVectorIndex(buildInfo)) {
-            return false;
-        }
-
-        if (buildInfo.KMeans.State == TIndexBuildInfo::TKMeans::Sample &&
-            !buildInfo.Sample.Rows.empty()) {
-            if (buildInfo.Sample.State == TIndexBuildInfo::TSample::EState::Collect) {
-                LOG_D("FillVectorIndex SendUploadSampleKRequest " << buildInfo.DebugString());
-                SendUploadSampleKRequest(buildInfo);
-                return false;
+        // (Sample -> Reshuffle)* -> MultiLocal -> NextLevel
+        if (buildInfo.KMeans.State == TIndexBuildInfo::TKMeans::Sample) {
+            return FillVectorIndexSamples(txc, buildInfo);
+        } else if (buildInfo.KMeans.State == TIndexBuildInfo::TKMeans::Reshuffle) {
+            if (NoShardsAdded(buildInfo)) {
+                AddGlobalShardsForCurrentParent(buildInfo);
             }
-        }
-
-        LOG_D("FillVectorIndex DoneLevel " << buildInfo.DebugString());
-        ClearDoneShards(txc, buildInfo);
-
-        if (!buildInfo.Sample.Rows.empty()) {
-            if (buildInfo.KMeans.State == TIndexBuildInfo::TKMeans::Sample) {
-                buildInfo.KMeans.State = TIndexBuildInfo::TKMeans::Reshuffle;
-                LOG_D("FillVectorIndex NextState " << buildInfo.DebugString());
-                PersistKMeansState(txc, buildInfo);
-                Progress(BuildId);
+            if (!SendKMeansReshuffle(buildInfo)) {
                 return false;
             }
+            ClearDoneShards(txc, buildInfo);
             buildInfo.Sample.Clear();
             NIceDb::TNiceDb db{txc.DB};
             Self->PersistBuildIndexSampleForget(db, buildInfo);
-            LOG_D("FillVectorIndex DoneState " << buildInfo.DebugString());
+            return FillVectorIndexNextParent(txc, buildInfo);
+        } else if (buildInfo.KMeans.State == TIndexBuildInfo::TKMeans::MultiLocal) {
+            if (!SendKMeansLocal(buildInfo)) {
+                return false;
+            }
+            ClearDoneShards(txc, buildInfo);
+            return FillVectorIndexNextParent(txc, buildInfo);
         }
+        Y_ENSURE(false);
+    }
 
-        if (buildInfo.KMeans.NextParent()) {
-            LOG_D("FillVectorIndex NextParent " << buildInfo.DebugString());
+    bool FillVectorIndexSamples(TTransactionContext& txc, TIndexBuildInfo& buildInfo) {
+        if (buildInfo.Sample.State == TIndexBuildInfo::TSample::EState::Collect) {
+            if (NoShardsAdded(buildInfo)) {
+                AddGlobalShardsForCurrentParent(buildInfo);
+                if (!buildInfo.DoneShards.size() && !buildInfo.ToUploadShards.size()) {
+                    // No "global" shards to handle - parent only has 1 shard,
+                    // it will be handled during the MultiLocal phase
+                    return FillVectorIndexNextParent(txc, buildInfo);
+                }
+                // Otherwise, we collect samples
+                LOG_D("FillVectorIndex Samples " << buildInfo.DebugString());
+            }
+            if (!SendKMeansSample(buildInfo)) {
+                return false;
+            }
+            ClearDoneShards(txc, buildInfo);
+            if (buildInfo.Sample.Rows.empty()) {
+                // No samples => no data for this cluster
+                return FillVectorIndexNextParent(txc, buildInfo);
+            }
+            LOG_D("FillVectorIndex SendUploadSampleKRequest " << buildInfo.DebugString());
+            SendUploadSampleKRequest(buildInfo);
+            buildInfo.Sample.State = TIndexBuildInfo::TSample::EState::Upload;
+            return false;
+        } else if (buildInfo.Sample.State == TIndexBuildInfo::TSample::EState::Upload) {
+            // Just wait until samples are uploaded (saved)
+            return false;
+        } else if (buildInfo.Sample.State == TIndexBuildInfo::TSample::EState::Done) {
+            buildInfo.KMeans.State = TIndexBuildInfo::TKMeans::Reshuffle;
+            LOG_D("FillVectorIndex NextState " << buildInfo.DebugString());
             PersistKMeansState(txc, buildInfo);
             Progress(BuildId);
             return false;
         }
+        Y_ENSURE(false);
+    }
 
-        if (InitMultiKMeans(buildInfo)) {
-            LOG_D("FillVectorIndex MultiKMeans " << buildInfo.DebugString());
+    bool FillVectorIndexNextParent(TTransactionContext& txc, TIndexBuildInfo& buildInfo) {
+        if (buildInfo.KMeans.NextParent()) {
+            buildInfo.KMeans.State = TIndexBuildInfo::TKMeans::Sample;
+            LOG_D("FillVectorIndex NextParent " << buildInfo.DebugString());
             PersistKMeansState(txc, buildInfo);
             Progress(BuildId);
             return false;
         }
 
+        if (!buildInfo.Cluster2Shards.empty()) {
+            AddLocalClusters(buildInfo);
+            buildInfo.Cluster2Shards.clear();
+            if (!buildInfo.ToUploadShards.empty()) {
+                LOG_D("FillVectorIndex MultiKMeans " << buildInfo.DebugString());
+                buildInfo.KMeans.State = TIndexBuildInfo::TKMeans::MultiLocal;
+                PersistKMeansState(txc, buildInfo);
+                Progress(BuildId);
+                return false;
+            }
+        }
+
         if (buildInfo.KMeans.NextLevel()) {
+            buildInfo.KMeans.State = TIndexBuildInfo::TKMeans::Sample;
             LOG_D("FillVectorIndex NextLevel " << buildInfo.DebugString());
             PersistKMeansState(txc, buildInfo);
             NIceDb::TNiceDb db{txc.DB};
@@ -1070,6 +1074,7 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
             Progress(BuildId);
             return false;
         }
+
         LOG_D("FillVectorIndex Done " << buildInfo.DebugString());
         return true;
     }
diff --git a/ydb/core/tx/schemeshard/schemeshard_info_types.cpp b/ydb/core/tx/schemeshard/schemeshard_info_types.cpp
@@ -2225,14 +2225,14 @@ void TIndexBuildInfo::AddParent(const TSerializedTableRange& range, TShardIdx sh
     const auto [parentFrom, parentTo] = KMeans.RangeToBorders(range);
     // TODO(mbkkt) We can make it more granular
 
-    // if new range is not intersect with other ranges, it's local
+    // the new range does not intersect with other ranges, just add it with 1 shard
     auto itFrom = Cluster2Shards.lower_bound(parentFrom);
     if (itFrom == Cluster2Shards.end() || parentTo < itFrom->second.From) {
-        Cluster2Shards.emplace_hint(itFrom, parentTo, TClusterShards{.From = parentFrom, .Local = shard});
+        Cluster2Shards.emplace_hint(itFrom, parentTo, TClusterShards{.From = parentFrom, .Shards = {shard}});
         return;
     }
 
-    // otherwise, this range is global and we need to merge all intersecting ranges
+    // otherwise, this range has multiple shards and we need to merge all intersecting ranges
     auto itTo = parentTo < itFrom->first ? itFrom : Cluster2Shards.lower_bound(parentTo);
     if (itTo == Cluster2Shards.end()) {
         itTo = Cluster2Shards.rbegin().base();
@@ -2244,25 +2244,16 @@ void TIndexBuildInfo::AddParent(const TSerializedTableRange& range, TShardIdx sh
         itTo = Cluster2Shards.insert(Cluster2Shards.end(), std::move(node));
         itFrom = needsToReplaceFrom ? itTo : itFrom;
     }
-    auto& [toFrom, toLocal, toGlobal] = itTo->second;
+    auto& [toFrom, toShards] = itTo->second;
 
     toFrom = std::min(toFrom, parentFrom);
-    if (toLocal != InvalidShardIdx) {
-        toGlobal.emplace_back(toLocal);
-        toLocal = InvalidShardIdx;
-    }
-    toGlobal.emplace_back(shard);
+    toShards.emplace_back(shard);
 
     while (itFrom != itTo) {
-        const auto& [fromFrom, fromLocal, fromGlobal] = itFrom->second;
+        const auto& [fromFrom, fromShards] = itFrom->second;
         toFrom = std::min(toFrom, fromFrom);
-        if (fromLocal != InvalidShardIdx) {
-            Y_ASSERT(fromGlobal.empty());
-            toGlobal.emplace_back(fromLocal);
-        } else {
-            Y_ASSERT(!fromGlobal.empty());
-            toGlobal.insert(toGlobal.end(), fromGlobal.begin(), fromGlobal.end());
-        }
+        Y_ASSERT(!fromShards.empty());
+        toShards.insert(toShards.end(), fromShards.begin(), fromShards.end());
         itFrom = Cluster2Shards.erase(itFrom);
     }
 }
diff --git a/ydb/core/tx/schemeshard/schemeshard_info_types.h b/ydb/core/tx/schemeshard/schemeshard_info_types.h
@@ -3091,7 +3091,6 @@ struct TIndexBuildInfo: public TSimpleRefCount<TIndexBuildInfo> {
             Sample = 0,
             // Recompute,
             Reshuffle,
-            Local,
             MultiLocal,
         };
         ui32 Level = 1;
@@ -3141,7 +3140,6 @@ struct TIndexBuildInfo: public TSimpleRefCount<TIndexBuildInfo> {
             if (!NeedsAnotherParent()) {
                 return false;
             }
-            State = Sample;
             ++Parent;
             Child += K;
             return true;
@@ -3151,14 +3149,12 @@ struct TIndexBuildInfo: public TSimpleRefCount<TIndexBuildInfo> {
             if (!NeedsAnotherLevel()) {
                 return false;
             }
-            State = Sample;
             NextLevel(ChildCount());
             return true;
         }
 
         void PrefixIndexDone(ui64 shards) {
             Y_ENSURE(NeedsAnotherLevel());
-            State = MultiLocal;
             // There's two worst cases, but in both one shard contains TableSize rows
             // 1. all rows have unique prefix (*), in such case we need 1 id for each row (parent, id in prefix table)
             // 2. all unique prefixes have size K, so we have TableSize/K parents + TableSize childs
@@ -3442,10 +3438,9 @@ struct TIndexBuildInfo: public TSimpleRefCount<TIndexBuildInfo> {
 
     struct TClusterShards {
         NTableIndex::TClusterId From = std::numeric_limits<NTableIndex::TClusterId>::max();
-        TShardIdx Local = InvalidShardIdx;
-        std::vector<TShardIdx> Global;
+        std::vector<TShardIdx> Shards;
     };
-    TMap<NTableIndex::TClusterId, TClusterShards> Cluster2Shards;
+    TMap<NTableIndex::TClusterId, TClusterShards> Cluster2Shards; // To => { From, Shards }
 
     void AddParent(const TSerializedTableRange& range, TShardIdx shard);
 
