Support assigning vectors to multiple clusters to improve vector search quality and performance (#24390) (#30180)

Author: vitalif

ydb/core/base/kmeans_clusters.cpp

@@ -1,5 +1,7 @@
 #include "kmeans_clusters.h"
 
+#include <ydb/public/api/protos/ydb_table.pb.h>
+
 #include <library/cpp/dot_product/dot_product.h>
 #include <library/cpp/l1_distance/l1_distance.h>
 #include <library/cpp/l2_distance/l2_distance.h>
@@ -86,6 +88,14 @@ namespace {
         ValidateSettingInRange(name, result, minValue, maxValue, error);
         return result;
     }
+
+    double ParseDouble(const TString& name, const TString& value, TString& error) {
+        double result = 0;
+        if (!TryFromString(value, result)) {
+            error = TStringBuilder() << "Invalid " << name << ": " << value;
+        }
+        return result;
+    }
 }
 
 // TODO(mbkkt) maybe compute floating sum in double? Needs benchmark
@@ -575,6 +585,18 @@ bool ValidateSettings(const Ydb::Table::KMeansTreeSettings& settings, TString& e
         return false;
     }
 
+    if (settings.has_overlap_clusters() &&
+        settings.overlap_clusters() > settings.clusters()) {
+        error = TStringBuilder() << "overlap_clusters should be less than or equal to clusters";
+        return false;
+    }
+
+    if (settings.has_overlap_ratio() &&
+        settings.overlap_ratio() < 0) {
+        error = TStringBuilder() << "overlap_ratio should be >= 0";
+        return false;
+    }
+
     ui64 clustersPowLevels = 1;
     for (ui64 i = 0; i < settings.levels(); ++i) {
         clustersPowLevels *= settings.clusters();
@@ -650,6 +672,10 @@ bool FillSetting(Ydb::Table::KMeansTreeSettings& settings, const TString& name,
         settings.set_clusters(ParseUInt32(name, value, MinClusters, MaxClusters, error));
     } else if (nameLower =="levels") {
         settings.set_levels(ParseUInt32(name, value, MinLevels, MaxLevels, error));
+    } else if (nameLower == "overlap_clusters") {
+        settings.set_overlap_clusters(ParseUInt32(name, value, MinClusters, MaxClusters, error));
+    } else if (nameLower == "overlap_ratio") {
+        settings.set_overlap_ratio(ParseDouble(name, value, error));
     } else {
         error = TStringBuilder() << "Unknown index setting: " << name;
         return false;
@@ -683,4 +709,27 @@ void FilterOverlapRows(TVector<TSerializedCellVec>& rows, size_t distancePos, ui
     }
 }
 
+void FilterOverlapRows(TVector<std::pair<NTableIndex::NKMeans::TClusterId, double>>& rowClusters, ui32 overlapClusters, double overlapRatio) {
+    if (rowClusters.size() <= 1) {
+        return;
+    }
+    std::sort(rowClusters.begin(), rowClusters.end(),
+        [&](const std::pair<NTableIndex::NKMeans::TClusterId, double>& a,
+            const std::pair<NTableIndex::NKMeans::TClusterId, double>& b) {
+            return a.second < b.second;
+        });
+    if (rowClusters.size() > overlapClusters) {
+        rowClusters.resize(overlapClusters);
+    }
+    if (overlapRatio > 0) {
+        double thresh = (rowClusters[0].second < 0 ? rowClusters[0].second/overlapRatio : rowClusters[0].second*overlapRatio);
+        for (size_t i = 1; i < rowClusters.size(); i++) {
+            if (rowClusters[i].second > thresh) {
+                rowClusters.resize(i);
+                break;
+            }
+        }
+    }
+}
+
 }

ydb/core/base/kmeans_clusters.h

@@ -2,7 +2,14 @@
 
 #include <ydb/core/scheme/scheme_tablecell.h>
 
-#include <ydb/public/api/protos/ydb_table.pb.h>
+namespace Ydb::Table {
+    class VectorIndexSettings;
+    class KMeansTreeSettings;
+}
+
+namespace NKikimr::NTableIndex::NKMeans {
+    using TClusterId = ui64;
+}
 
 namespace NKikimr::NKMeans {
 
@@ -56,5 +63,6 @@ bool ValidateSettings(const Ydb::Table::VectorIndexSettings& settings, TString&
 bool ValidateSettings(const Ydb::Table::KMeansTreeSettings& settings, TString& error);
 bool FillSetting(Ydb::Table::KMeansTreeSettings& settings, const TString& name, const TString& value, TString& error);
 void FilterOverlapRows(TVector<TSerializedCellVec>& rows, size_t distancePos, ui32 overlapClusters, double overlapRatio);
+void FilterOverlapRows(TVector<std::pair<NTableIndex::NKMeans::TClusterId, double>>& rowClusters, ui32 overlapClusters, double overlapRatio);
 
 }

ydb/core/base/table_index.h

@@ -79,15 +79,19 @@ inline constexpr const char* PostingTable = "indexImplPostingTable";
 inline constexpr const char* BuildSuffix0 = "0build";
 inline constexpr const char* BuildSuffix1 = "1build";
 inline constexpr auto IsForeignType = Ydb::Type::BOOL;
+inline constexpr auto IsForeignTypeName = "Bool";
 inline constexpr const char* IsForeignColumn = "__ydb_foreign";
 inline constexpr auto DistanceType = Ydb::Type::DOUBLE;
+inline constexpr auto DistanceTypeName = "Double";
 inline constexpr const char* DistanceColumn = "__ydb_distance";
 
 // Prefix table
 inline constexpr const char* PrefixTable = "indexImplPrefixTable";
 inline constexpr const char* IdColumnSequence = "__ydb_id_sequence";
 
 inline constexpr const int DefaultKMeansRounds = 3;
+inline constexpr const int DefaultOverlapClusters = 1;
+inline constexpr const double DefaultOverlapRatio = 0;
 
 inline constexpr TClusterId PostingParentFlag = (1ull << 63ull);
 

ydb/core/kqp/common/kqp_yql.cpp

@@ -735,11 +735,32 @@ NNodes::TCoNameValueTupleList TKqpStreamLookupSettings::BuildNode(TExprContext&
                 .Done());
     }
 
+    if (VectorTopDistinct) {
+        settings.emplace_back(
+            Build<TCoNameValueTuple>(ctx, pos)
+                .Name().Build(VectorTopDistinctSettingName)
+                .Done());
+    }
+
     return Build<TCoNameValueTupleList>(ctx, pos)
         .Add(settings)
         .Done();
 }
 
+bool TKqpStreamLookupSettings::HasVectorTopDistinct(const NNodes::TCoNameValueTupleList& list) {
+    for (const auto& tuple : list) {
+        auto name = tuple.Name().Value();
+        if (name == VectorTopDistinctSettingName) {
+            return true;
+        }
+    }
+    return false;
+}
+
+bool TKqpStreamLookupSettings::HasVectorTopDistinct(const NNodes::TKqlStreamLookupTable& node) {
+    return TKqpStreamLookupSettings::HasVectorTopDistinct(node.Settings());
+}
+
 TKqpStreamLookupSettings TKqpStreamLookupSettings::Parse(const NNodes::TCoNameValueTupleList& list) {
     TKqpStreamLookupSettings settings;
 
@@ -775,6 +796,8 @@ TKqpStreamLookupSettings TKqpStreamLookupSettings::Parse(const NNodes::TCoNameVa
         } else if (name == VectorTopLimitSettingName) {
             YQL_ENSURE(tuple.Value().IsValid());
             settings.VectorTopLimit = tuple.Value().Cast().Ptr();
+        } else if (name == VectorTopDistinctSettingName) {
+            settings.VectorTopDistinct = true;
         } else {
             YQL_ENSURE(false, "Unknown KqpStreamLookup setting name '" << name << "'");
         }

ydb/core/kqp/common/kqp_yql.h

@@ -60,6 +60,7 @@ struct TKqpStreamLookupSettings {
     static constexpr TStringBuf VectorTopIndexSettingName = "VectorTopIndex";
     static constexpr TStringBuf VectorTopLimitSettingName = "VectorTopLimit";
     static constexpr TStringBuf VectorTopTargetSettingName = "VectorTopTarget";
+    static constexpr TStringBuf VectorTopDistinctSettingName = "VectorTopDistinct";
 
     // stream lookup strategy types
     static constexpr std::string_view LookupStrategyName = "LookupRows"sv;
@@ -76,11 +77,15 @@ struct TKqpStreamLookupSettings {
     TExprNode::TPtr VectorTopTarget;
     TExprNode::TPtr VectorTopLimit;
 
+    bool VectorTopDistinct = false;
+
     NNodes::TCoNameValueTupleList BuildNode(TExprContext& ctx, TPositionHandle pos) const;
     static TKqpStreamLookupSettings Parse(const NNodes::TKqlStreamLookupTable& node);
     static TKqpStreamLookupSettings Parse(const NNodes::TKqlStreamLookupIndex& node);
     static TKqpStreamLookupSettings Parse(const NNodes::TKqpCnStreamLookup& node);
     static TKqpStreamLookupSettings Parse(const NNodes::TCoNameValueTupleList& node);
+    static bool HasVectorTopDistinct(const NNodes::TKqlStreamLookupTable& node);
+    static bool HasVectorTopDistinct(const NNodes::TCoNameValueTupleList& node);
 };
 
 struct TKqpDeleteRowsIndexSettings {

ydb/core/kqp/executer_actor/kqp_tasks_graph.cpp

@@ -500,6 +500,9 @@ void TKqpTasksGraph::BuildStreamLookupChannels(const TStageInfo& stageInfo, ui32
         auto target = ExtractPhyValue(stageInfo, in.GetTargetVector(), TxAlloc->HolderFactory, TxAlloc->TypeEnv, NUdf::TUnboxedValuePod());
         out.SetTargetVector(TString(target.AsStringRef()));
         out.SetLimit((ui32)ExtractPhyValue(stageInfo, in.GetLimit(), TxAlloc->HolderFactory, TxAlloc->TypeEnv, NUdf::TUnboxedValuePod()).Get<ui64>());
+        for (auto& colIdx: in.GetDistinctColumns()) {
+            out.AddDistinctColumns(colIdx);
+        }
     }
 
     TTransform streamLookupTransform;
@@ -520,6 +523,8 @@ void TKqpTasksGraph::BuildVectorResolveChannels(const TStageInfo& stageInfo, ui3
     auto* settings = GetMeta().Allocate<NKikimrTxDataShard::TKqpVectorResolveSettings>();
 
     *settings->MutableIndexSettings() = vectorResolve.GetIndexSettings();
+    settings->SetOverlapClusters(vectorResolve.GetOverlapClusters());
+    settings->SetOverlapRatio(vectorResolve.GetOverlapRatio());
 
     YQL_ENSURE(stageInfo.Meta.IndexMetas.size() == 1);
     const auto& levelTableInfo = stageInfo.Meta.IndexMetas.back().TableConstInfo;

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

@@ -74,6 +74,11 @@ TExprBase KqpApplyExtractMembersToReadTable(TExprBase node, TExprContext& ctx, c
         return node;
     }
 
+    auto slt = node.Maybe<TKqlStreamLookupTable>();
+    if (slt && TKqpStreamLookupSettings::HasVectorTopDistinct(slt.Cast())) {
+        return node;
+    }
+
     TCoAtomList columnsNode = TExprBase(node.Ptr()->Child(TKqlReadColumnsNodeIdx)).Cast<TCoAtomList>();
     auto usedColumns = GetUsedColumns(node, columnsNode, parentsMap, allowMultiUsage, ctx);
     if (!usedColumns) {

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

@@ -588,6 +588,7 @@ void VectorReadMain(
     const TIntrusivePtr<TKikimrTableMetadata> & mainTableMeta,
     const TCoAtomList& mainColumns,
     const TKqpStreamLookupSettings& pushdownSettings,
+    bool withOverlap,
     TExprNodePtr& read)
 {
     const bool isCovered = CheckIndexCovering(mainColumns, postingTableMeta);
@@ -606,21 +607,48 @@ void VectorReadMain(
             .Columns(postingColumns)
             .Settings(isVectorCovered ? settingsNode : settings.BuildNode(ctx, pos))
         .Done().Ptr();
+    }
 
-        read = Build<TKqlStreamLookupTable>(ctx, pos)
-            .Table(mainTable)
-            .LookupKeys(read)
-            .Columns(mainColumns)
-            .Settings(settingsNode)
-        .Done().Ptr();
-    } else {
-        read = Build<TKqlStreamLookupTable>(ctx, pos)
-            .Table(postingTable)
-            .LookupKeys(read)
-            .Columns(mainColumns)
-            .Settings(settingsNode)
-        .Done().Ptr();
+    const auto& targetTable = isCovered ? postingTable : mainTable;
+
+    if (withOverlap) {
+        // mainColumns must contain primary key columns for DistinctColumns pushdown
+        THashSet<TStringBuf> cols;
+        for (const auto& col: mainColumns) {
+            cols.insert(col.Value());
+        }
+        TVector<TCoAtom> columnsWithKey;
+        for (const auto& col: mainTableMeta->KeyColumnNames) {
+            if (!cols.contains(col)) {
+                columnsWithKey.push_back(Build<TCoAtom>(ctx, pos)
+                    .Value(col)
+                    .Done());
+            }
+        }
+        if (columnsWithKey.size()) {
+            for (const auto& col: mainColumns) {
+                columnsWithKey.push_back(col);
+            }
+            read = Build<TKqlStreamLookupTable>(ctx, pos)
+                .Table(targetTable)
+                .LookupKeys(read)
+                .Columns<TCoAtomList>().Add(columnsWithKey).Build()
+                .Settings(settingsNode)
+                .Done().Ptr();
+            read = Build<TCoExtractMembers>(ctx, pos)
+                .Input(read)
+                .Members(mainColumns)
+                .Done().Ptr();
+            return;
+        }
     }
+
+    read = Build<TKqlStreamLookupTable>(ctx, pos)
+        .Table(targetTable)
+        .LookupKeys(read)
+        .Columns(mainColumns)
+        .Settings(settingsNode)
+        .Done().Ptr();
 }
 
 void VectorTopMain(TExprContext& ctx, const TCoTopBase& top, TExprNodePtr& read) {
@@ -633,6 +661,10 @@ void VectorTopMain(TExprContext& ctx, const TCoTopBase& top, TExprNodePtr& read)
     .Done().Ptr();
 }
 
+// FIXME Most of this rewriting should probably be handled in kqp/opt/physical
+// Logical optimizer should only rewrite it to something like TKqlReadTableVectorIndex
+// This would remove the need for skipping KqpApplyExtractMembersToReadTable based on settings.VectorTopDistinct
+
 TExprBase DoRewriteTopSortOverKMeansTree(
     const TReadMatch& match, const TMaybeNode<TCoFlatMap>& flatMap, const TExprBase& lambdaArgs, const TExprBase& lambdaBody, const TCoTopBase& top,
     TExprContext& ctx, const TKqpOptimizeContext& kqpCtx,
@@ -697,6 +729,9 @@ TExprBase DoRewriteTopSortOverKMeansTree(
 
     const auto levelTop = kqpCtx.Config->KMeansTreeSearchTopSize.Get().GetOrElse(1);
 
+    const auto& kmeansDesc = std::get<NKikimrKqp::TVectorIndexKmeansTreeDescription>(indexDesc.SpecializedIndexDescription);
+    const bool withOverlap = kmeansDesc.settings().overlap_clusters() > 1;
+
     TKqpStreamLookupSettings settings;
     settings.Strategy = EStreamLookupStrategyType::LookupRows;
     settings.VectorTopColumn = NTableIndex::NKMeans::CentroidColumn;
@@ -716,7 +751,8 @@ TExprBase DoRewriteTopSortOverKMeansTree(
 
     settings.VectorTopColumn = indexDesc.KeyColumns.back();
     settings.VectorTopLimit = top.Count().Ptr();
-    VectorReadMain(ctx, pos, postingTable, postingTableDesc->Metadata, mainTable, tableDesc.Metadata, mainColumns, settings, read);
+    settings.VectorTopDistinct = true;
+    VectorReadMain(ctx, pos, postingTable, postingTableDesc->Metadata, mainTable, tableDesc.Metadata, mainColumns, settings, withOverlap, read);
 
     if (flatMap) {
         read = Build<TCoFlatMap>(ctx, flatMap.Cast().Pos())
@@ -827,6 +863,9 @@ TExprBase DoRewriteTopSortOverPrefixedKMeansTree(
 
     const auto levelTop = kqpCtx.Config->KMeansTreeSearchTopSize.Get().GetOrElse(1);
 
+    const auto& kmeansDesc = std::get<NKikimrKqp::TVectorIndexKmeansTreeDescription>(indexDesc.SpecializedIndexDescription);
+    const bool withOverlap = kmeansDesc.settings().overlap_clusters() > 1;
+
     TKqpStreamLookupSettings settings;
     settings.Strategy = EStreamLookupStrategyType::LookupRows;
     settings.VectorTopColumn = NTableIndex::NKMeans::CentroidColumn;
@@ -849,7 +888,8 @@ TExprBase DoRewriteTopSortOverPrefixedKMeansTree(
 
     settings.VectorTopColumn = indexDesc.KeyColumns.back();
     settings.VectorTopLimit = top.Count().Ptr();
-    VectorReadMain(ctx, pos, postingTable, postingTableDesc->Metadata, mainTable, tableDesc.Metadata, mainColumns, settings, read);
+    settings.VectorTopDistinct = true;
+    VectorReadMain(ctx, pos, postingTable, postingTableDesc->Metadata, mainTable, tableDesc.Metadata, mainColumns, settings, withOverlap, read);
 
     if (mainLambda) {
         read = Build<TCoMap>(ctx, flatMap.Pos())

ydb/core/kqp/query_compiler/kqp_query_compiler.cpp

@@ -1670,15 +1670,21 @@ class TKqpQueryCompiler : public IKqpQueryCompiler {
         *vectorTopK.MutableSettings() = kmeansDesc.GetSettings().Getsettings();
 
         // Column index
+        THashMap<TStringBuf, ui32> readColumnIndexes;
         ui32 columnIdx = 0;
         for (const auto& column: streamLookupProto.GetColumns()) {
-            if (column == settings.VectorTopColumn) {
-                break;
+            readColumnIndexes[column] = columnIdx++;
+        }
+        YQL_ENSURE(readColumnIndexes.contains(settings.VectorTopColumn));
+        vectorTopK.SetColumn(readColumnIndexes.at(settings.VectorTopColumn));
+
+        // Unique columns - required when we read from the index posting table and overlap is enabled
+        if (settings.VectorTopDistinct) {
+            for (const auto& keyColumn: mainTable->Metadata->KeyColumnNames) {
+                YQL_ENSURE(readColumnIndexes.contains(keyColumn));
+                vectorTopK.AddDistinctColumns(readColumnIndexes.at(keyColumn));
             }
-            columnIdx++;
         }
-        YQL_ENSURE(columnIdx < streamLookup.Columns().Size());
-        vectorTopK.SetColumn(columnIdx);
 
         // Limit - may be a parameter which will be linked later
         TExprBase expr(settings.VectorTopLimit);
@@ -1992,6 +1998,8 @@ class TKqpQueryCompiler : public IKqpQueryCompiler {
             // Index settings
             auto& kmeansDesc = std::get<NKikimrKqp::TVectorIndexKmeansTreeDescription>(indexDesc->SpecializedIndexDescription);
             *vectorResolveProto.MutableIndexSettings() = kmeansDesc.GetSettings().Getsettings();
+            vectorResolveProto.SetOverlapClusters(kmeansDesc.GetSettings().overlap_clusters());
+            vectorResolveProto.SetOverlapRatio(kmeansDesc.GetSettings().overlap_ratio());
 
             // Main table
             FillTablesMap(vectorResolve.Table(), tablesMap);