Vector index build preparation in SchemeShard

Author: azevaykin

ydb/core/base/table_index.cpp

@@ -140,5 +140,9 @@ bool IsCompatibleIndex(const NKikimrSchemeOp::EIndexType indexType, const TTable
     return true;
 }
 
+bool IsImplTable(std::string_view tableName) {
+    return std::find(std::begin(ImplTables), std::end(ImplTables), tableName) != std::end(ImplTables);
+}
+
 }
 }

ydb/core/base/table_index.h

@@ -21,8 +21,13 @@ struct TIndexColumns {
     TVector<TString> DataColumns;
 };
 
+inline constexpr const char* ImplTable = "indexImplTable";
+inline constexpr std::string_view ImplTables[] = {ImplTable, NTableVectorKmeansTreeIndex::LevelTable, NTableVectorKmeansTreeIndex::PostingTable};
+
 bool IsCompatibleIndex(const NKikimrSchemeOp::EIndexType indexType, const TTableColumns& table, const TIndexColumns& index, TString& explain);
 TTableColumns CalcTableImplDescription(const NKikimrSchemeOp::EIndexType indexType, const TTableColumns& table, const TIndexColumns& index);
 
+bool IsImplTable(std::string_view tableName);
+
 }
 }

ydb/core/tx/schemeshard/schemeshard__operation_apply_build_index.cpp

@@ -52,17 +52,25 @@ TVector<ISubOperation::TPtr> ApplyBuildIndex(TOperationId nextId, const TTxTrans
 
     if (!indexName.empty())
     {
+        auto alterImplTableTransactionTemplate = [] (TPath index, TPath implIndexTable, TTableInfo::TPtr implIndexTableInfo) {
+            auto indexImplTableAltering = TransactionTemplate(index.PathString(), NKikimrSchemeOp::EOperationType::ESchemeOpFinalizeBuildIndexImplTable);
+            auto alterTable = indexImplTableAltering.MutableAlterTable();
+            alterTable->SetName(implIndexTable.LeafName());
+            alterTable->MutablePartitionConfig()->MutableCompactionPolicy()->CopyFrom(implIndexTableInfo->PartitionConfig().GetCompactionPolicy());
+            alterTable->MutablePartitionConfig()->MutableCompactionPolicy()->SetKeepEraseMarkers(false);
+            alterTable->MutablePartitionConfig()->SetShadowData(false);
+            return indexImplTableAltering;
+        };
+
         TPath index = table.Child(indexName);
-        TPath implIndexTable = index.Child("indexImplTable");
-        TTableInfo::TPtr implIndexTableInfo = context.SS->Tables.at(implIndexTable.Base()->PathId);
-        auto indexImplTableAltering = TransactionTemplate(index.PathString(), NKikimrSchemeOp::EOperationType::ESchemeOpFinalizeBuildIndexImplTable);
-        auto alterTable = indexImplTableAltering.MutableAlterTable();
-        alterTable->SetName(implIndexTable.LeafName());
-        alterTable->MutablePartitionConfig()->MutableCompactionPolicy()->CopyFrom(implIndexTableInfo->PartitionConfig().GetCompactionPolicy());
-        alterTable->MutablePartitionConfig()->MutableCompactionPolicy()->SetKeepEraseMarkers(false);
-        alterTable->MutablePartitionConfig()->SetShadowData(false);
-
-        result.push_back(CreateFinalizeBuildIndexImplTable(NextPartId(nextId, result), indexImplTableAltering));
+        for (const std::string_view implTable : NTableIndex::ImplTables) {
+            TPath implIndexTable = index.Child(implTable.data());
+            if (!implIndexTable.IsResolved()) {
+                continue;
+            }
+            TTableInfo::TPtr implIndexTableInfo = context.SS->Tables.at(implIndexTable.Base()->PathId);
+            result.push_back(CreateFinalizeBuildIndexImplTable(NextPartId(nextId, result), alterImplTableTransactionTemplate(index, implIndexTable, implIndexTableInfo)));
+        }
     }
 
     return result;
@@ -109,7 +117,7 @@ TVector<ISubOperation::TPtr> CancelBuildIndex(TOperationId nextId, const TTxTran
         Y_ABORT_UNLESS(index.Base()->GetChildren().size() == 1);
         for (auto& indexChildItems: index.Base()->GetChildren()) {
             const TString& implTableName = indexChildItems.first;
-            Y_ABORT_UNLESS(implTableName == "indexImplTable", "unexpected name %s", implTableName.c_str());
+            Y_ABORT_UNLESS(NTableIndex::IsImplTable(implTableName), "unexpected name %s", implTableName.c_str());
 
             TPath implTable = index.Child(implTableName);
             {

ydb/core/tx/schemeshard/schemeshard__operation_drop_indexed_table.cpp

@@ -495,10 +495,8 @@ TVector<ISubOperation::TPtr> CreateDropIndexedTable(TOperationId nextId, const T
         }
 
         for (auto& [implName, implPathId] : child.Base()->GetChildren()) {
-            Y_ABORT_UNLESS(implName == "indexImplTable" 
+            Y_ABORT_UNLESS(NTableIndex::IsImplTable(implName)
                         || implName == "streamImpl"
-                        || implName == NTableIndex::NTableVectorKmeansTreeIndex::LevelTable
-                        || implName == NTableIndex::NTableVectorKmeansTreeIndex::PostingTable
                 , "unexpected name %s", implName.c_str());
 
             TPath implPath = child.Child(implName);

ydb/core/tx/schemeshard/schemeshard_build_index__create.cpp

@@ -226,9 +226,13 @@ class TSchemeShard::TIndexBuilder::TTxCreate: public TSchemeShard::TIndexBuilder
             case Ydb::Table::TableIndex::TypeCase::kGlobalUniqueIndex:
                 explain = "unsupported index type to build";
                 return false;
-            case Ydb::Table::TableIndex::TypeCase::kGlobalVectorKmeansTreeIndex:
-                explain = "unsupported vector index type to build";
-                return false;
+            case Ydb::Table::TableIndex::TypeCase::kGlobalVectorKmeansTreeIndex: {
+                buildInfo->IndexType = NKikimrSchemeOp::EIndexType::EIndexTypeGlobalVectorKmeansTree;
+                NKikimrSchemeOp::TVectorIndexKmeansTreeDescription vectorIndexKmeansTreeDescription;
+                *vectorIndexKmeansTreeDescription.MutableSettings() = index.global_vector_kmeans_tree_index().vector_settings();
+                buildInfo->SpecializedIndexDescription = vectorIndexKmeansTreeDescription;
+                break;
+            }
             case Ydb::Table::TableIndex::TypeCase::TYPE_NOT_SET:
                 explain = "invalid or unset index type";
                 return false;

ydb/core/tx/schemeshard/schemeshard_build_index__progress.cpp

@@ -253,6 +253,13 @@ struct TSchemeShard::TIndexBuilder::TTxProgress: public TSchemeShard::TIndexBuil
             } else if (!buildInfo->InitiateTxDone) {
                 Send(Self->SelfId(), MakeHolder<TEvSchemeShard::TEvNotifyTxCompletion>(ui64(buildInfo->InitiateTxId)));
             } else {
+                // TODO add vector index filling
+                if (buildInfo->IndexType == NKikimrSchemeOp::EIndexType::EIndexTypeGlobalVectorKmeansTree) {
+                    ChangeState(BuildId, TIndexBuildInfo::EState::Applying);
+                    Progress(BuildId);
+                    break;
+                }
+
                 ChangeState(BuildId, TIndexBuildInfo::EState::Filling);
                 Progress(BuildId);
             }

ydb/core/tx/schemeshard/schemeshard_info_types.cpp

@@ -2122,6 +2122,10 @@ void TIndexBuildInfo::SerializeToProto(TSchemeShard* ss, NKikimrSchemeOp::TIndex
     for (const auto& implTableDescription : ImplTableDescriptions) {
         *index.AddIndexImplTableDescriptions() = implTableDescription;
     }
+
+    if (IndexType == NKikimrSchemeOp::EIndexType::EIndexTypeGlobalVectorKmeansTree) {
+        *index.MutableVectorIndexKmeansTreeDescription() = std::get<NKikimrSchemeOp::TVectorIndexKmeansTreeDescription>(SpecializedIndexDescription);
+    }
 }
 
 void TIndexBuildInfo::SerializeToProto(TSchemeShard* ss, NKikimrIndexBuilder::TColumnBuildSettings* result) const {

ydb/core/tx/schemeshard/schemeshard_info_types.h

@@ -2927,6 +2927,8 @@ struct TIndexBuildInfo: public TSimpleRefCount<TIndexBuildInfo> {
     NTableIndex::TTableColumns ImplTableColumns;
     TVector<NKikimrSchemeOp::TTableDescription> ImplTableDescriptions;
 
+    std::variant<std::monostate, NKikimrSchemeOp::TVectorIndexKmeansTreeDescription> SpecializedIndexDescription;
+
     EState State = EState::Invalid;
     TString Issue;
 

ydb/core/tx/schemeshard/schemeshard_utils.cpp

@@ -451,7 +451,7 @@ NKikimrSchemeOp::TTableDescription CalcImplTableDesc(
 {
     NKikimrSchemeOp::TTableDescription implTableDesc;
 
-    implTableDesc.SetName("indexImplTable");
+    implTableDesc.SetName(NTableIndex::ImplTable);
 
     SetImplTablePartitionConfig(baseTableInfo->PartitionConfig(), indexTableDesc, implTableDesc);
 
@@ -467,7 +467,7 @@ NKikimrSchemeOp::TTableDescription CalcImplTableDesc(
 {
     NKikimrSchemeOp::TTableDescription implTableDesc;
 
-    implTableDesc.SetName("indexImplTable");
+    implTableDesc.SetName(NTableIndex::ImplTable);
 
     SetImplTablePartitionConfig(baseTableDescr.GetPartitionConfig(), indexTableDesc, implTableDesc);