Use vector index recompute in scheme shard (#19154) (#19854)

Author: vitalif

ydb/core/base/kmeans_clusters.cpp

@@ -125,43 +125,35 @@ class TClusters: public IClusters {
     using TSum = TMetric::TSum;
     using TEmbedding = TVector<TSum>;
 
-    ui32 InitK = 0;
-    ui32 K = 0;
     const ui32 Dimensions = 0;
+    const ui32 MaxRounds = 0;
+    const ui8 TypeByte = 0;
 
     TVector<TString> Clusters;
     TVector<ui64> ClusterSizes;
-
-    struct TAggregatedCluster {
-        TEmbedding Cluster;
-        ui64 Size = 0;
-    };
-    TVector<TAggregatedCluster> AggregatedClusters;
+    TVector<TEmbedding> NextClusters;
+    TVector<ui64> NextClusterSizes;
 
     ui32 Round = 0;
-    ui32 MaxRounds = 0;
 
 public:
-    TClusters(ui32 dimensions)
+    TClusters(ui32 dimensions, ui32 maxRounds, ui8 typeByte)
         : Dimensions(dimensions)
+        , MaxRounds(maxRounds)
+        , TypeByte(typeByte)
     {
     }
 
-    void Init(ui32 k, ui32 maxRounds) override {
-        InitK = k;
-        K = k;
-        MaxRounds = maxRounds;
-    }
-
-    ui32 GetK() const override {
-        return K;
+    void SetRound(ui32 round) override {
+        Round = round;
     }
 
     TString Debug() const override {
-        if (!MaxRounds) {
-            return TStringBuilder() << "K: " << K;
+        auto sb = TStringBuilder() << "K: " << Clusters.size();
+        if (MaxRounds) {
+            sb << " Round: " << Round << " / " << MaxRounds;
         }
-        return TStringBuilder() << "K: " << K << " Round: " << Round << " / " << MaxRounds;
+        return sb;
     }
 
     const TVector<TString>& GetClusters() const override {
@@ -172,11 +164,19 @@ class TClusters: public IClusters {
         return ClusterSizes;
     }
 
+    const TVector<ui64>& GetNextClusterSizes() const override {
+        return NextClusterSizes;
+    }
+
+    virtual void SetClusterSize(ui32 num, ui64 size) override {
+        ClusterSizes.at(num) = size;
+    }
+
     void Clear() override {
-        K = InitK;
         Clusters.clear();
         ClusterSizes.clear();
-        AggregatedClusters.clear();
+        NextClusterSizes.clear();
+        NextClusters.clear();
         Round = 0;
     }
 
@@ -189,40 +189,37 @@ class TClusters: public IClusters {
                 return false;
             }
         }
-        Clusters = newClusters;
-        K = newClusters.size();
-        return true;
-    }
-
-    void InitAggregatedClusters() override {
-        AggregatedClusters.resize(K);
-        ClusterSizes.resize(K, 0);
-        for (auto& aggregate : AggregatedClusters) {
-            aggregate.Cluster.resize(Dimensions, 0);
+        Clusters = std::move(newClusters);
+        ClusterSizes.clear();
+        ClusterSizes.resize(Clusters.size());
+        NextClusterSizes.clear();
+        NextClusterSizes.resize(Clusters.size());
+        NextClusters.clear();
+        NextClusters.resize(Clusters.size());
+        for (auto& aggregate : NextClusters) {
+            aggregate.resize(Dimensions, 0);
         }
-        Round = 1;
+        return true;
     }
 
     bool RecomputeClusters() override {
-        Y_ENSURE(K >= 1);
         ui64 vectorCount = 0;
         ui64 reassignedCount = 0;
-        for (size_t i = 0; auto& aggregate : AggregatedClusters) {
-            vectorCount += aggregate.Size;
+        for (size_t i = 0; auto& aggregate : NextClusters) {
+            auto newSize = NextClusterSizes[i];
+            vectorCount += newSize;
 
-            auto& clusterSize = ClusterSizes[i];
-            reassignedCount += clusterSize < aggregate.Size ? aggregate.Size - clusterSize : 0;
-            clusterSize = aggregate.Size;
+            auto clusterSize = ClusterSizes[i];
+            reassignedCount += clusterSize < newSize ? newSize - clusterSize : 0;
 
-            if (aggregate.Size != 0) {
-                this->Fill(Clusters[i], aggregate.Cluster.data(), aggregate.Size);
-                Y_ENSURE(aggregate.Size == 0);
+            if (newSize != 0) {
+                this->Fill(Clusters[i], aggregate.data(), newSize);
             }
             ++i;
         }
-        Y_ENSURE(vectorCount >= K);
+
         Y_ENSURE(reassignedCount <= vectorCount);
-        if (K == 1) {
+        if (Clusters.size() == 1) {
             return true;
         }
 
@@ -232,7 +229,6 @@ class TClusters: public IClusters {
             last = changes < MinVectorsNeedsReassigned;
         }
         if (!last) {
-            ++Round;
             return false;
         }
         return true;
@@ -251,6 +247,25 @@ class TClusters: public IClusters {
         Clusters.erase(Clusters.begin() + w, Clusters.end());
     }
 
+    bool NextRound() override {
+        bool isLast = RecomputeClusters();
+        ClusterSizes = std::move(NextClusterSizes);
+        RemoveEmptyClusters();
+        if (isLast) {
+            NextClusters.clear();
+            return true;
+        }
+        ++Round;
+        NextClusterSizes.clear();
+        NextClusterSizes.resize(Clusters.size());
+        NextClusters.clear();
+        NextClusters.resize(Clusters.size());
+        for (auto& aggregate : NextClusters) {
+            aggregate.resize(Dimensions, 0);
+        }
+        return false;
+    }
+
     std::optional<ui32> FindCluster(TArrayRef<const TCell> row, ui32 embeddingPos) override {
         Y_ENSURE(embeddingPos < row.size());
         const auto embedding = row.at(embeddingPos).AsRef();
@@ -271,16 +286,17 @@ class TClusters: public IClusters {
         return closest;
     }
 
-    void AggregateToCluster(ui32 pos, const char* embedding) override {
-        auto& aggregate = AggregatedClusters[pos];
-        auto* coords = aggregate.Cluster.data();
-        for (auto coord : this->GetCoords(embedding)) {
-            *coords++ += coord;
+    void AggregateToCluster(ui32 pos, const TArrayRef<const char>& embedding, ui64 weight) override {
+        auto& aggregate = NextClusters.at(pos);
+        auto* coords = aggregate.data();
+        Y_ENSURE(IsExpectedSize(embedding));
+        for (auto coord : this->GetCoords(embedding.data())) {
+            *coords++ += (TSum)coord * weight;
         }
-        ++aggregate.Size;
+        NextClusterSizes.at(pos) += weight;
     }
 
-    bool IsExpectedSize(TArrayRef<const char> data) override {
+    bool IsExpectedSize(const TArrayRef<const char>& data) override {
         return data.size() == 1 + sizeof(TCoord) * Dimensions;
     }
 
@@ -295,36 +311,37 @@ class TClusters: public IClusters {
 
     void Fill(TString& d, TSum* embedding, ui64& c) {
         Y_ENSURE(c > 0);
-        const auto count = static_cast<TSum>(std::exchange(c, 0));
+        const auto count = static_cast<TSum>(c);
         auto data = GetData(d.MutRef().data());
         for (auto& coord : data) {
             coord = *embedding / count;
-            *embedding++ = 0;
+            embedding++;
         }
     }
 };
 
-std::unique_ptr<IClusters> CreateClusters(const Ydb::Table::VectorIndexSettings& settings, TString& error) {
+std::unique_ptr<IClusters> CreateClusters(const Ydb::Table::VectorIndexSettings& settings, ui32 maxRounds, TString& error) {
     if (settings.vector_dimension() < 1) {
         error = "Dimension of vector should be at least one";
         return nullptr;
     }
 
+    const ui8 typeVal = (ui8)settings.vector_type();
     const ui32 dim = settings.vector_dimension();
 
     auto handleMetric = [&]<typename T>() -> std::unique_ptr<IClusters> {
         switch (settings.metric()) {
             case Ydb::Table::VectorIndexSettings::SIMILARITY_INNER_PRODUCT:
-                return std::make_unique<TClusters<TMaxInnerProductSimilarity<T>>>(dim);
+                return std::make_unique<TClusters<TMaxInnerProductSimilarity<T>>>(dim, maxRounds, typeVal);
             case Ydb::Table::VectorIndexSettings::SIMILARITY_COSINE:
             case Ydb::Table::VectorIndexSettings::DISTANCE_COSINE:
                 // We don't need to have separate implementation for distance,
                 // because clusters will be same as for similarity
-                return std::make_unique<TClusters<TCosineSimilarity<T>>>(dim);
+                return std::make_unique<TClusters<TCosineSimilarity<T>>>(dim, maxRounds, typeVal);
             case Ydb::Table::VectorIndexSettings::DISTANCE_MANHATTAN:
-                return std::make_unique<TClusters<TL1Distance<T>>>(dim);
+                return std::make_unique<TClusters<TL1Distance<T>>>(dim, maxRounds, typeVal);
             case Ydb::Table::VectorIndexSettings::DISTANCE_EUCLIDEAN:
-                return std::make_unique<TClusters<TL2Distance<T>>>(dim);
+                return std::make_unique<TClusters<TL2Distance<T>>>(dim, maxRounds, typeVal);
             default:
                 error = "Wrong similarity";
                 break;

ydb/core/base/kmeans_clusters.h

@@ -10,33 +10,35 @@ class IClusters {
 public:
     virtual ~IClusters() = default;
 
-    virtual void Init(ui32 k, ui32 maxRounds) = 0;
-
-    virtual ui32 GetK() const = 0;
+    virtual void SetRound(ui32 round) = 0;
 
     virtual TString Debug() const = 0;
 
     virtual const TVector<TString>& GetClusters() const = 0;
 
     virtual const TVector<ui64>& GetClusterSizes() const = 0;
 
+    virtual const TVector<ui64>& GetNextClusterSizes() const = 0;
+
+    virtual void SetClusterSize(ui32 num, ui64 size) = 0;
+
     virtual void Clear() = 0;
 
     virtual bool SetClusters(TVector<TString> && newClusters) = 0;
 
-    virtual void InitAggregatedClusters() = 0;
-
     virtual bool RecomputeClusters() = 0;
 
+    virtual bool NextRound() = 0;
+
     virtual void RemoveEmptyClusters() = 0;
 
     virtual std::optional<ui32> FindCluster(TArrayRef<const TCell> row, ui32 embeddingPos) = 0;
 
-    virtual void AggregateToCluster(ui32 pos, const char* embedding) = 0;
+    virtual void AggregateToCluster(ui32 pos, const TArrayRef<const char>& embedding, ui64 weight = 1) = 0;
 
-    virtual bool IsExpectedSize(TArrayRef<const char> data) = 0;
+    virtual bool IsExpectedSize(const TArrayRef<const char>& data) = 0;
 };
 
-std::unique_ptr<IClusters> CreateClusters(const Ydb::Table::VectorIndexSettings& settings, TString& error);
+std::unique_ptr<IClusters> CreateClusters(const Ydb::Table::VectorIndexSettings& settings, ui32 maxRounds, TString& error);
 
 }

ydb/core/base/table_vector_index.h

@@ -21,4 +21,6 @@ inline constexpr const char* BuildSuffix1 = "1build";
 // Prefix table
 inline constexpr const char* PrefixTable = "indexImplPrefixTable";
 
+inline constexpr const int DefaultKMeansRounds = 3;
+
 }

ydb/core/tx/datashard/build_index/local_kmeans.cpp

@@ -132,7 +132,6 @@ class TLocalKMeansScan: public TActor<TLocalKMeansScan>, public NTable::IScan {
         , Clusters(std::move(clusters))
     {
         LOG_I("Create " << Debug());
-        Clusters->Init(request.GetK(), request.GetNeedsRounds());
 
         const auto& embedding = request.GetEmbeddingColumn();
         const auto& data = request.GetDataColumns();
@@ -231,7 +230,7 @@ class TLocalKMeansScan: public TActor<TLocalKMeansScan>, public NTable::IScan {
         if (PrefixColumns && !Prefix) {
             Prefix = TSerializedCellVec{key.subspan(0, PrefixColumns)};
             auto newParent = key.at(0).template AsValue<ui64>();
-            Child += (newParent - Parent) * Clusters->GetK();
+            Child += (newParent - Parent) * K;
             Parent = newParent;
         }
 
@@ -380,13 +379,11 @@ class TLocalKMeansScan: public TActor<TLocalKMeansScan>, public NTable::IScan {
             }
             bool ok = Clusters->SetClusters(std::move(rows));
             Y_ENSURE(ok);
-            Clusters->InitAggregatedClusters();
             return false; // do KMEANS
         }
 
         if (State == EState::KMEANS) {
-            if (Clusters->RecomputeClusters()) {
-                Clusters->RemoveEmptyClusters();
+            if (Clusters->NextRound()) {
                 FormLevelRows();
                 State = UploadState;
                 return false; // do UPLOAD_*
@@ -444,7 +441,7 @@ class TLocalKMeansScan: public TActor<TLocalKMeansScan>, public NTable::IScan {
     void FeedKMeans(TArrayRef<const TCell> row) noexcept
     {
         if (auto pos = Clusters->FindCluster(row, EmbeddingPos); pos) {
-            Clusters->AggregateToCluster(*pos, row.at(EmbeddingPos).Data());
+            Clusters->AggregateToCluster(*pos, row.at(EmbeddingPos).AsRef());
         }
     }
 
@@ -661,7 +658,7 @@ void TDataShard::HandleSafe(TEvDataShard::TEvLocalKMeansRequest::TPtr& ev, const
 
         // 3. Validating vector index settings
         TString error;
-        auto clusters = NKikimr::NKMeans::CreateClusters(request.GetSettings(), error);
+        auto clusters = NKikimr::NKMeans::CreateClusters(request.GetSettings(), request.GetNeedsRounds(), error);
         if (!clusters) {
             badRequest(error);
             auto sent = trySendBadRequest();

ydb/core/tx/datashard/build_index/prefix_kmeans.cpp

@@ -134,7 +134,6 @@ class TPrefixKMeansScan: public TActor<TPrefixKMeansScan>, public NTable::IScan
         , Clusters(std::move(clusters))
     {
         LOG_I("Create " << Debug());
-        Clusters->Init(request.GetK(), request.GetNeedsRounds());
 
         const auto& embedding = request.GetEmbeddingColumn();
         TVector<TString> data{request.GetDataColumns().begin(), request.GetDataColumns().end()};
@@ -355,7 +354,7 @@ class TPrefixKMeansScan: public TActor<TPrefixKMeansScan>, public NTable::IScan
     }
 
     void StartNewPrefix() {
-        Parent = Child + Clusters->GetK();
+        Parent = Child + K;
         Child = Parent + 1;
         State = EState::SAMPLE;
         Lead.To(Prefix.GetCells(), NTable::ESeek::Upper); // seek to (prefix, inf)
@@ -426,13 +425,11 @@ class TPrefixKMeansScan: public TActor<TPrefixKMeansScan>, public NTable::IScan
             }
             bool ok = Clusters->SetClusters(std::move(rows));
             Y_ENSURE(ok);
-            Clusters->InitAggregatedClusters();
             return false; // do KMEANS
         }
 
         if (State == EState::KMEANS) {
-            if (Clusters->RecomputeClusters()) {
-                Clusters->RemoveEmptyClusters();
+            if (Clusters->NextRound()) {
                 FormLevelRows();
                 State = UploadState;
                 return false; // do UPLOAD_*
@@ -484,7 +481,7 @@ class TPrefixKMeansScan: public TActor<TPrefixKMeansScan>, public NTable::IScan
     void FeedKMeans(TArrayRef<const TCell> row)
     {
         if (auto pos = Clusters->FindCluster(row, EmbeddingPos); pos) {
-            Clusters->AggregateToCluster(*pos, row.at(EmbeddingPos).Data());
+            Clusters->AggregateToCluster(*pos, row.at(EmbeddingPos).AsRef());
         }
     }
 
@@ -650,7 +647,7 @@ void TDataShard::HandleSafe(TEvDataShard::TEvPrefixKMeansRequest::TPtr& ev, cons
 
         // 3. Validating vector index settings
         TString error;
-        auto clusters = NKikimr::NKMeans::CreateClusters(request.GetSettings(), error);
+        auto clusters = NKikimr::NKMeans::CreateClusters(request.GetSettings(), request.GetNeedsRounds(), error);
         if (!clusters) {
             badRequest(error);
             auto sent = trySendBadRequest();