Moved commit "Create vector index in SchemeShard" from ydb repo

Author: azevaykin

include/ydb-cpp-sdk/client/table/table.h

@@ -23,6 +23,7 @@ class ChangefeedDescription;
 class DescribeTableResult;
 class ExplicitPartitions;
 class GlobalIndexSettings;
+class VectorIndexSettings;
 class PartitioningSettings;
 class DateTypeColumnModeSettings;
 class TtlSettings;
@@ -190,6 +191,45 @@ struct TGlobalIndexSettings {
     void SerializeTo(Ydb::Table::GlobalIndexSettings& proto) const;
 };
 
+struct TVectorIndexSettings {
+public:
+    enum class EDistance {
+        Cosine,
+        Manhattan,
+        Euclidean,
+
+        Unknown = std::numeric_limits<int>::max()
+    };
+
+    enum class ESimilarity {
+        Cosine,
+        InnerProduct,
+
+        Unknown = std::numeric_limits<int>::max()
+    };
+
+    enum class EVectorType {
+        Float,
+        Uint8,
+        Int8,
+        Bit,
+
+        Unknown = std::numeric_limits<int>::max()
+    };
+    using TMetric = std::variant<std::monostate, EDistance, ESimilarity>;
+
+    TMetric Metric;
+    EVectorType VectorType;
+    uint32_t VectorDimension;
+
+    template <typename TProto>
+    static TVectorIndexSettings FromProto(const TProto& proto);
+
+    void SerializeTo(Ydb::Table::VectorIndexSettings& settings) const;
+
+    void Out(IOutputStream &o) const;
+};
+
 //! Represents index description
 class TIndexDescription {
     friend class NYdb::TProtoAccessor;
@@ -200,20 +240,22 @@ class TIndexDescription {
         EIndexType type,
         const std::vector<std::string>& indexColumns,
         const std::vector<std::string>& dataColumns = {},
-        const TGlobalIndexSettings& settings = {}
+        const std::vector<TGlobalIndexSettings>& globalIndexSettings = {},
+        const std::optional<TVectorIndexSettings>& vectorIndexSettings = {}
     );
 
     TIndexDescription(
         const std::string& name,
         const std::vector<std::string>& indexColumns,
         const std::vector<std::string>& dataColumns = {},
-        const TGlobalIndexSettings& settings = {}
+        const std::vector<TGlobalIndexSettings>& globalIndexSettings = {}
     );
 
     const std::string& GetIndexName() const;
     EIndexType GetIndexType() const;
     const std::vector<std::string>& GetIndexColumns() const;
     const std::vector<std::string>& GetDataColumns() const;
+    const std::optional<TVectorIndexSettings>& GetVectorIndexSettings() const;
     uint64_t GetSizeBytes() const;
 
     void SerializeTo(Ydb::Table::TableIndex& proto) const;
@@ -232,7 +274,8 @@ class TIndexDescription {
     EIndexType IndexType_;
     std::vector<std::string> IndexColumns_;
     std::vector<std::string> DataColumns_;
-    TGlobalIndexSettings GlobalIndexSettings_;
+    std::vector<TGlobalIndexSettings> GlobalIndexSettings_;
+    std::optional<TVectorIndexSettings> VectorIndexSettings_;
     uint64_t SizeBytes = 0;
 };
 
@@ -606,6 +649,9 @@ class TTableDescription {
     // unique
     void AddUniqueSecondaryIndex(const std::string& indexName, const std::vector<std::string>& indexColumns);
     void AddUniqueSecondaryIndex(const std::string& indexName, const std::vector<std::string>& indexColumns, const std::vector<std::string>& dataColumns);
+    // vector KMeansTree
+    void AddVectorKMeansTreeSecondaryIndex(const std::string& indexName, const std::vector<std::string>& indexColumns, const TVectorIndexSettings& vectorIndexSettings);
+    void AddVectorKMeansTreeSecondaryIndex(const std::string& indexName, const std::vector<std::string>& indexColumns, const std::vector<std::string>& dataColumns, const TVectorIndexSettings& vectorIndexSettings);
 
     // default
     void AddSecondaryIndex(const std::string& indexName, const std::vector<std::string>& indexColumns);
@@ -825,6 +871,10 @@ class TTableBuilder {
     TTableBuilder& AddUniqueSecondaryIndex(const std::string& indexName, const std::vector<std::string>& indexColumns);
     TTableBuilder& AddUniqueSecondaryIndex(const std::string& indexName, const std::vector<std::string>& indexColumns, const std::vector<std::string>& dataColumns);
 
+    // vector KMeansTree
+    TTableBuilder& AddVectorKMeansTreeSecondaryIndex(const std::string& indexName, const std::vector<std::string>& indexColumns, const TVectorIndexSettings& vectorIndexSettings);
+    TTableBuilder& AddVectorKMeansTreeSecondaryIndex(const std::string& indexName, const std::vector<std::string>& indexColumns, const std::vector<std::string>& dataColumns, const TVectorIndexSettings& vectorIndexSettings);
+
     // default
     TTableBuilder& AddSecondaryIndex(const std::string& indexName, const std::vector<std::string>& indexColumns, const std::vector<std::string>& dataColumns);
     TTableBuilder& AddSecondaryIndex(const std::string& indexName, const std::vector<std::string>& indexColumns);

include/ydb-cpp-sdk/client/table/table_enum.h

@@ -28,6 +28,7 @@ enum class EIndexType {
     GlobalSync,
     GlobalAsync,
     GlobalUnique,
+    GlobalVectorKMeansTree,
 
     Unknown = std::numeric_limits<int>::max()
 };

src/client/table/out.cpp

@@ -23,3 +23,69 @@ Y_DECLARE_OUT_SPEC(, NYdb::NTable::TCreateSessionResult, o, x) {
 Y_DECLARE_OUT_SPEC(, NYdb::NTable::TDescribeTableResult, o, x) {
     return x.Out(o);
 }
+
+Y_DECLARE_OUT_SPEC(, NYdb::NTable::TVectorIndexSettings::EDistance, stream, value) {
+    auto convertDistance = [] (auto value) -> auto {
+        switch (value) {
+        case NYdb::NTable::TVectorIndexSettings::EDistance::Cosine:
+            return "COSINE";
+        case NYdb::NTable::TVectorIndexSettings::EDistance::Manhattan:
+            return "MANHATTAN";
+        case NYdb::NTable::TVectorIndexSettings::EDistance::Euclidean:
+            return "EUCLIDEAN";
+        case NYdb::NTable::TVectorIndexSettings::EDistance::Unknown:
+            return "UNKNOWN";
+        }
+    };
+
+    stream << convertDistance(value);
+}
+
+Y_DECLARE_OUT_SPEC(, NYdb::NTable::TVectorIndexSettings::ESimilarity, stream, value) {
+    auto convertSimilarity = [] (auto value) -> auto {
+        switch (value) {
+        case NYdb::NTable::TVectorIndexSettings::ESimilarity::Cosine:
+            return "COSINE";
+        case NYdb::NTable::TVectorIndexSettings::ESimilarity::InnerProduct:
+            return "INNER_PRODUCT";
+        case NYdb::NTable::TVectorIndexSettings::ESimilarity::Unknown:
+            return "UNKNOWN";
+        }
+    };
+
+    stream << convertSimilarity(value);
+}
+
+Y_DECLARE_OUT_SPEC(, NYdb::NTable::TVectorIndexSettings::EVectorType, stream, value) {
+    auto convertVectorType = [] (auto value) -> auto {
+        switch (value) {
+        case NYdb::NTable::TVectorIndexSettings::EVectorType::Float:
+            return "FLOAT";
+        case NYdb::NTable::TVectorIndexSettings::EVectorType::Uint8:
+            return "UINT8";
+        case NYdb::NTable::TVectorIndexSettings::EVectorType::Int8:
+            return "INT8";
+        case NYdb::NTable::TVectorIndexSettings::EVectorType::Bit:
+            return "BIT";
+        case NYdb::NTable::TVectorIndexSettings::EVectorType::Unknown:
+            return "UNKNOWN";
+        }
+    };
+
+    stream << convertVectorType(value);
+}
+
+Y_DECLARE_OUT_SPEC(, NYdb::NTable::TVectorIndexSettings, stream, value) {
+    stream << "{";
+
+    if (const auto* distance = std::get_if<NYdb::NTable::TVectorIndexSettings::EDistance>(&value.Metric)) {
+        stream << " distance: " << *distance << "";
+    } else if (const auto* similarity = std::get_if<NYdb::NTable::TVectorIndexSettings::ESimilarity>(&value.Metric)) {
+        stream << " similarity: " << *similarity << "";
+    }
+
+    stream << ", vector_type: " << value.VectorType << "";
+    stream << ", vector_dimension: " << value.VectorDimension << "";
+
+    stream << " }";
+}