Unify the def of max vector dims

pingcap · Aug 13, 2024 · 6458415 · 6458415
1 parent 61d05ff
commit 6458415
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/dbms/src/Storages/DeltaMerge/Index/VectorIndex.cpp b/dbms/src/Storages/DeltaMerge/Index/VectorIndex.cpp
@@ -18,6 +18,7 @@
 #include <Storages/DeltaMerge/File/dtpb/dmfile.pb.h>
 #include <Storages/DeltaMerge/Index/VectorIndex.h>
 #include <Storages/DeltaMerge/Index/VectorIndexHNSW/Index.h>
+#include <TiDB/Schema/VectorIndex.h>
 #include <tipb/executor.pb.h>
 
 namespace DB::ErrorCodes
@@ -41,15 +42,15 @@ bool VectorIndexBuilder::isSupportedType(const IDataType & type)
 VectorIndexBuilderPtr VectorIndexBuilder::create(const TiDB::VectorIndexDefinitionPtr & definition)
 {
     RUNTIME_CHECK(definition->dimension > 0);
-    RUNTIME_CHECK(definition->dimension <= std::numeric_limits<UInt32>::max());
+    RUNTIME_CHECK(definition->dimension <= TiDB::MAX_VECTOR_DIMENSION);
 
     switch (definition->kind)
     {
     case tipb::VectorIndexKind::HNSW:
         return std::make_shared<VectorIndexHNSWBuilder>(definition);
     default:
         throw Exception(
-            ErrorCodes::INCORRECT_QUERY,
+            ErrorCodes::BAD_ARGUMENTS,
             "Unsupported vector index {}",
             tipb::VectorIndexKind_Name(definition->kind));
     }
@@ -58,7 +59,7 @@ VectorIndexBuilderPtr VectorIndexBuilder::create(const TiDB::VectorIndexDefiniti
 VectorIndexViewerPtr VectorIndexViewer::view(const dtpb::VectorIndexFileProps & file_props, std::string_view path)
 {
     RUNTIME_CHECK(file_props.dimensions() > 0);
-    RUNTIME_CHECK(file_props.dimensions() <= std::numeric_limits<UInt32>::max());
+    RUNTIME_CHECK(file_props.dimensions() <= TiDB::MAX_VECTOR_DIMENSION);
 
     tipb::VectorIndexKind kind;
     RUNTIME_CHECK(tipb::VectorIndexKind_Parse(file_props.index_kind(), &kind));
@@ -68,7 +69,7 @@ VectorIndexViewerPtr VectorIndexViewer::view(const dtpb::VectorIndexFileProps &
     case tipb::VectorIndexKind::HNSW:
         return VectorIndexHNSWViewer::view(file_props, path);
     default:
-        throw Exception(ErrorCodes::INCORRECT_QUERY, "Unsupported vector index {}", file_props.index_kind());
+        throw Exception(ErrorCodes::BAD_ARGUMENTS, "Unsupported vector index {}", file_props.index_kind());
     }
 }
 

diff --git a/dbms/src/TiDB/Schema/TiDB.cpp b/dbms/src/TiDB/Schema/TiDB.cpp
@@ -34,6 +34,7 @@
 #include <algorithm>
 #include <cmath>
 #include <magic_enum.hpp>
+#include "TiDB/Schema/VectorIndex.h"
 
 namespace DB
 {
@@ -481,7 +482,7 @@ try
 
         auto dimension = vector_index_json->getValue<UInt64>("dimension");
         RUNTIME_CHECK(dimension > 0);
-        RUNTIME_CHECK(dimension <= 16383); // Just a protection
+        RUNTIME_CHECK(dimension <= TiDB::MAX_VECTOR_DIMENSION); // Just a protection
 
         tipb::VectorDistanceMetric distance_metric = tipb::VectorDistanceMetric::INVALID_DISTANCE_METRIC;
         auto distance_metric_field = vector_index_json->getValue<String>("distance_metric");

diff --git a/dbms/src/TiDB/Schema/VectorIndex.h b/dbms/src/TiDB/Schema/VectorIndex.h
@@ -39,6 +39,9 @@ struct VectorIndexDefinition
 // ever try to modify it anyway.
 using VectorIndexDefinitionPtr = std::shared_ptr<const VectorIndexDefinition>;
 
+// Defined in TiDB pkg/types/vector.go
+static constexpr Int64 MAX_VECTOR_DIMENSION = 16383;
+
 } // namespace TiDB
 
 template <>