Skip to content

Commit 8c4c654

Browse files
committed
Reject vector index settings with unknown protobuf fields
1 parent 7b591b5 commit 8c4c654

File tree

3 files changed

+113
-1
lines changed

3 files changed

+113
-1
lines changed

ydb/core/base/kmeans_clusters.cpp

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -453,6 +453,11 @@ std::unique_ptr<IClusters> CreateClusters(const Ydb::Table::VectorIndexSettings&
453453
bool ValidateSettings(const Ydb::Table::KMeansTreeSettings& settings, TString& error) {
454454
error = "";
455455

456+
if (auto unknownCount = settings.GetReflection()->GetUnknownFields(settings).field_count(); unknownCount > 0) {
457+
error = TStringBuilder() << "vector index settings contain " << unknownCount << " unsupported parameter(s)";
458+
return false;
459+
}
460+
456461
if (!settings.has_settings()) {
457462
error = TStringBuilder() << "vector index settings should be set";
458463
return false;
@@ -497,7 +502,10 @@ bool ValidateSettings(const Ydb::Table::KMeansTreeSettings& settings, TString& e
497502
}
498503

499504
bool ValidateSettings(const Ydb::Table::VectorIndexSettings& settings, TString& error) {
500-
error = "";
505+
if (auto unknownCount = settings.GetReflection()->GetUnknownFields(settings).field_count(); unknownCount > 0) {
506+
error = TStringBuilder() << "vector index settings contain " << unknownCount << " unsupported parameter(s)";
507+
return false;
508+
}
501509

502510
if (!settings.has_metric() || settings.metric() == Ydb::Table::VectorIndexSettings::METRIC_UNSPECIFIED) {
503511
error = TStringBuilder() << "either distance or similarity should be set";

ydb/core/base/ut/kmeans_ut.cpp

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
#include "kmeans_clusters.h"
2+
3+
#include <ydb/public/api/protos/ydb_table.pb.h>
4+
5+
#include <library/cpp/testing/unittest/registar.h>
6+
#include <util/generic/xrange.h>
7+
8+
namespace NKikimr::NKMeans {
9+
10+
Y_UNIT_TEST_SUITE(NKMeans) {
11+
12+
Y_UNIT_TEST(ValidateSettings) {
13+
Ydb::Table::KMeansTreeSettings settings;
14+
TString error;
15+
16+
UNIT_ASSERT(!ValidateSettings(settings, error));
17+
UNIT_ASSERT_VALUES_EQUAL(error, "vector index settings should be set");
18+
19+
auto& vectorIndex = *settings.mutable_settings();
20+
21+
UNIT_ASSERT(!ValidateSettings(settings, error));
22+
UNIT_ASSERT_VALUES_EQUAL(error, "either distance or similarity should be set");
23+
24+
vectorIndex.set_metric((::Ydb::Table::VectorIndexSettings_Metric)UINT32_MAX);
25+
UNIT_ASSERT(!ValidateSettings(settings, error));
26+
UNIT_ASSERT_STRING_CONTAINS(error, "Invalid metric");
27+
28+
vectorIndex.set_metric(Ydb::Table::VectorIndexSettings::DISTANCE_COSINE);
29+
UNIT_ASSERT(!ValidateSettings(settings, error));
30+
UNIT_ASSERT_VALUES_EQUAL(error, "vector_type should be set");
31+
32+
vectorIndex.set_vector_type((::Ydb::Table::VectorIndexSettings_VectorType)UINT32_MAX);
33+
UNIT_ASSERT(!ValidateSettings(settings, error));
34+
UNIT_ASSERT_STRING_CONTAINS(error, "Invalid vector_type");
35+
36+
vectorIndex.set_vector_type(Ydb::Table::VectorIndexSettings::VECTOR_TYPE_BIT);
37+
UNIT_ASSERT(!ValidateSettings(settings, error));
38+
UNIT_ASSERT_STRING_CONTAINS(error, "vector_dimension should be set");
39+
40+
vectorIndex.set_vector_dimension(UINT32_MAX);
41+
UNIT_ASSERT(!ValidateSettings(settings, error));
42+
UNIT_ASSERT_STRING_CONTAINS(error, "Invalid vector_dimension");
43+
44+
vectorIndex.set_vector_dimension(16384);
45+
vectorIndex.GetReflection()->MutableUnknownFields(&vectorIndex)->AddVarint(10000, 100000);
46+
UNIT_ASSERT(!ValidateSettings(settings, error));
47+
UNIT_ASSERT_STRING_CONTAINS(error, "vector index settings contain 1 unsupported parameter(s)");
48+
49+
vectorIndex.GetReflection()->MutableUnknownFields(&vectorIndex)->DeleteByNumber(10000);
50+
settings.GetReflection()->MutableUnknownFields(&settings)->AddVarint(10000, 100000);
51+
UNIT_ASSERT(!ValidateSettings(settings, error));
52+
UNIT_ASSERT_STRING_CONTAINS(error, "vector index settings contain 1 unsupported parameter(s)");
53+
54+
settings.GetReflection()->MutableUnknownFields(&settings)->DeleteByNumber(10000);
55+
UNIT_ASSERT(!ValidateSettings(settings, error));
56+
UNIT_ASSERT_STRING_CONTAINS(error, "levels should be set");
57+
58+
settings.set_levels(UINT32_MAX);
59+
UNIT_ASSERT(!ValidateSettings(settings, error));
60+
UNIT_ASSERT_STRING_CONTAINS(error, "Invalid levels");
61+
62+
settings.set_levels(16);
63+
UNIT_ASSERT(!ValidateSettings(settings, error));
64+
UNIT_ASSERT_STRING_CONTAINS(error, "clusters should be set");
65+
66+
settings.set_clusters(UINT32_MAX);
67+
UNIT_ASSERT(!ValidateSettings(settings, error));
68+
UNIT_ASSERT_STRING_CONTAINS(error, "Invalid clusters");
69+
70+
settings.set_clusters(2048);
71+
UNIT_ASSERT(!ValidateSettings(settings, error));
72+
UNIT_ASSERT_STRING_CONTAINS(error, "Invalid clusters^levels");
73+
74+
settings.set_clusters(4);
75+
settings.set_levels(16);
76+
UNIT_ASSERT(!ValidateSettings(settings, error));
77+
UNIT_ASSERT_STRING_CONTAINS(error, "Invalid clusters^levels");
78+
79+
settings.set_clusters(1024);
80+
settings.set_levels(3);
81+
UNIT_ASSERT(!ValidateSettings(settings, error));
82+
UNIT_ASSERT_STRING_CONTAINS(error, "Invalid vector_dimension*clusters");
83+
84+
settings.set_clusters(50);
85+
settings.set_levels(3);
86+
settings.set_overlap_clusters(51);
87+
UNIT_ASSERT(!ValidateSettings(settings, error));
88+
UNIT_ASSERT_VALUES_EQUAL(error, "overlap_clusters should be less than or equal to clusters");
89+
90+
settings.set_overlap_clusters(0);
91+
settings.set_overlap_ratio(-1.2);
92+
UNIT_ASSERT(!ValidateSettings(settings, error));
93+
UNIT_ASSERT_VALUES_EQUAL(error, "overlap_ratio should be >= 0");
94+
95+
settings.set_overlap_ratio(0);
96+
UNIT_ASSERT(ValidateSettings(settings, error));
97+
}
98+
}
99+
100+
}

ydb/core/base/ut/ya.make

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,15 @@ SIZE(MEDIUM)
66
PEERDIR(
77
library/cpp/getopt
88
library/cpp/svnversion
9+
ydb/core/scheme
10+
yql/essentials/public/udf/service/stub
11+
yql/essentials/sql/pg_dummy
912
)
1013

1114
SRCS(
1215
path_ut.cpp
1316
blobstorage_grouptype_ut.cpp
17+
kmeans_ut.cpp
1418
localdb_ut.cpp
1519
logoblob_ut.cpp
1620
memory_stats_ut.cpp

0 commit comments

Comments
 (0)