Skip to content

Commit ffe7590

Browse files
flown4qqqqflown4qqqq
committed
init
1 parent eaa7bcb commit ffe7590

File tree

3 files changed

+215
-34
lines changed

3 files changed

+215
-34
lines changed

‎ydb/core/kqp/ut/indexes/kqp_indexes_prefixed_vector_ut.cpp‎

Lines changed: 74 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,31 @@ Y_UNIT_TEST_SUITE(KqpPrefixedVectorIndexes) {
184184
DoPositiveQueriesPrefixedVectorIndexOrderBy(session, "CosineSimilarity", "DESC", flags, init, count);
185185
}
186186

187+
bool IsTableEmpty(TSession& session, const TString& tablePath) {
188+
TString query = Sprintf(R"(
189+
SELECT COUNT(*) FROM `%s`;
190+
)"
191+
, tablePath.c_str());
192+
193+
auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).ExtractValueSync();
194+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
195+
196+
auto resultSet = result.GetResultSet(0);
197+
TResultSetParser parser(resultSet);
198+
UNIT_ASSERT(parser.TryNextRow());
199+
auto count = parser.ColumnParser(0).GetUint64();
200+
return count == 0;
201+
}
202+
203+
void DoTruncateTable(TSession& session) {
204+
const TString truncateTable(Q_(R"(
205+
TRUNCATE TABLE `/Root/TestTable`;
206+
)"));
207+
208+
auto result = session.ExecuteSchemeQuery(truncateTable).ExtractValueSync();
209+
UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
210+
}
211+
187212
TSession DoOnlyCreateTableForPrefixedVectorIndex(TTableClient& db, int flags = 0) {
188213
auto session = db.CreateSession().GetValueSync().GetSession();
189214

@@ -979,7 +1004,55 @@ Y_UNIT_TEST_SUITE(KqpPrefixedVectorIndexes) {
9791004
}
9801005
}
9811006

982-
}
1007+
Y_UNIT_TEST_QUAD(PrefixedVectorIndexTruncateTable, Covered, Overlap) {
1008+
NKikimrConfig::TFeatureFlags featureFlags;
1009+
featureFlags.SetEnableTruncateTable(true);
1010+
auto serverSettings = TKikimrSettings().SetFeatureFlags(featureFlags);
1011+
TKikimrRunner kikimr(serverSettings);
1012+
kikimr.GetTestServer().GetRuntime()->SetLogPriority(NKikimrServices::BUILD_INDEX, NActors::NLog::PRI_TRACE);
1013+
kikimr.GetTestServer().GetRuntime()->SetLogPriority(NKikimrServices::FLAT_TX_SCHEMESHARD, NActors::NLog::PRI_TRACE);
1014+
1015+
const int flags = F_NULLABLE | (Covered ? F_COVERING : 0) | (Overlap ? F_OVERLAP : 0);
1016+
auto db = kikimr.GetTableClient();
1017+
auto session = DoCreateTableForPrefixedVectorIndex(db, flags);
1018+
1019+
{
1020+
const char* cover = "";
1021+
if (flags & F_COVERING) {
1022+
cover = " COVER (user, emb, data)";
1023+
}
1024+
1025+
const TString createIndex(Q_(Sprintf(R"(
1026+
ALTER TABLE `/Root/TestTable`
1027+
ADD INDEX index
1028+
GLOBAL USING vector_kmeans_tree
1029+
ON (user, emb)%s
1030+
WITH (distance=cosine, vector_type="uint8", vector_dimension=2, levels=1, clusters=2%s);
1031+
)", cover, "")));
1032+
1033+
auto result = session.ExecuteSchemeQuery(createIndex)
1034+
.ExtractValueSync();
1035+
1036+
UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
1037+
}
1038+
1039+
DoPositiveQueriesPrefixedVectorIndexOrderByCosine(session, flags);
1040+
1041+
UNIT_ASSERT(!IsTableEmpty(session, "/Root/TestTable"));
1042+
UNIT_ASSERT(!IsTableEmpty(session, "/Root/TestTable/index/indexImplPostingTable"));
1043+
UNIT_ASSERT(!IsTableEmpty(session, "/Root/TestTable/index/indexImplLevelTable"));
1044+
UNIT_ASSERT(!IsTableEmpty(session, "/Root/TestTable/index/indexImplPrefixTable"));
1045+
1046+
DoTruncateTable(session);
1047+
1048+
UNIT_ASSERT(IsTableEmpty(session, "/Root/TestTable"));
1049+
UNIT_ASSERT(IsTableEmpty(session, "/Root/TestTable/index/indexImplPostingTable"));
1050+
UNIT_ASSERT(IsTableEmpty(session, "/Root/TestTable/index/indexImplLevelTable"));
1051+
UNIT_ASSERT(IsTableEmpty(session, "/Root/TestTable/index/indexImplPrefixTable"));
9831052

1053+
InsertDataForPrefixedVectorIndex(session);
1054+
DoPositiveQueriesPrefixedVectorIndexOrderByCosine(session, flags);
1055+
}
1056+
}
9841057
}
9851058
}

‎ydb/core/kqp/ut/indexes/kqp_indexes_vector_ut.cpp‎

Lines changed: 88 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -222,27 +222,54 @@ Y_UNIT_TEST_SUITE(KqpVectorIndexes) {
222222
return session;
223223
}
224224

225-
TSession DoCreateTableForVectorIndex(TTableClient& db, int flags = 0) {
225+
void DoTruncateTable(TSession& session) {
226+
const TString truncateTable(Q_(R"(
227+
TRUNCATE TABLE `/Root/TestTable`;
228+
)"));
229+
230+
auto result = session.ExecuteSchemeQuery(truncateTable).ExtractValueSync();
231+
UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
232+
}
233+
234+
bool IsTableEmpty(TSession& session, const TString& tablePath) {
235+
TString query = Sprintf(R"(
236+
SELECT COUNT(*) FROM `%s`;
237+
)"
238+
, tablePath.c_str());
239+
240+
auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).ExtractValueSync();
241+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
242+
243+
auto resultSet = result.GetResultSet(0);
244+
TResultSetParser parser(resultSet);
245+
UNIT_ASSERT(parser.TryNextRow());
246+
auto count = parser.ColumnParser(0).GetUint64();
247+
return count == 0;
248+
}
249+
250+
void DoOnlyUpsertValuesIntoTable(TSession& session, int flags = 0) {
226251
const char* dataCol = flags & F_UNDERSCORE_DATA ? "___data" : "data";
227-
auto session = DoOnlyCreateTableForVectorIndex(db, flags);
228-
{
229-
const TString query1 = TStringBuilder()
230-
<< "UPSERT INTO `/Root/TestTable` (pk, emb, " << dataCol << ") VALUES "
231-
<< "(0, \"\x03\x30\x02\", \"0\"),"
232-
"(1, \"\x13\x31\x02\", \"1\"),"
233-
"(2, \"\x23\x32\x02\", \"2\"),"
234-
"(3, \"\x53\x33\x02\", \"3\"),"
235-
"(4, \"\x43\x34\x02\", \"4\"),"
236-
"(5, \"\x50\x60\x02\", \"5\"),"
237-
"(6, \"\x61\x11\x02\", \"6\"),"
238-
"(7, \"\x12\x62\x02\", \"7\"),"
239-
"(8, \"\x75\x76\x02\", \"8\"),"
240-
"(9, \"\x76\x76\x02\", \"9\");";
252+
const TString query1 = TStringBuilder()
253+
<< "UPSERT INTO `/Root/TestTable` (pk, emb, " << dataCol << ") VALUES "
254+
<< "(0, \"\x03\x30\x02\", \"0\"),"
255+
"(1, \"\x13\x31\x02\", \"1\"),"
256+
"(2, \"\x23\x32\x02\", \"2\"),"
257+
"(3, \"\x53\x33\x02\", \"3\"),"
258+
"(4, \"\x43\x34\x02\", \"4\"),"
259+
"(5, \"\x50\x60\x02\", \"5\"),"
260+
"(6, \"\x61\x11\x02\", \"6\"),"
261+
"(7, \"\x12\x62\x02\", \"7\"),"
262+
"(8, \"\x75\x76\x02\", \"8\"),"
263+
"(9, \"\x76\x76\x02\", \"9\");";
264+
265+
auto result = session.ExecuteDataQuery(Q_(query1), TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
266+
.ExtractValueSync();
267+
UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
268+
}
241269

242-
auto result = session.ExecuteDataQuery(Q_(query1), TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
243-
.ExtractValueSync();
244-
UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
245-
}
270+
TSession DoCreateTableForVectorIndex(TTableClient& db, int flags = 0) {
271+
auto session = DoOnlyCreateTableForVectorIndex(db, flags);
272+
DoOnlyUpsertValuesIntoTable(session, flags);
246273
return session;
247274
}
248275

@@ -1266,6 +1293,48 @@ Y_UNIT_TEST_SUITE(KqpVectorIndexes) {
12661293
}
12671294
}
12681295

1296+
Y_UNIT_TEST_QUAD(VectorIndexTruncateTable, Covered, Overlap) {
1297+
NKikimrConfig::TFeatureFlags featureFlags;
1298+
featureFlags.SetEnableTruncateTable(true);
1299+
auto serverSettings = TKikimrSettings().SetFeatureFlags(featureFlags);
1300+
TKikimrRunner kikimr(serverSettings);
1301+
kikimr.GetTestServer().GetRuntime()->SetLogPriority(NKikimrServices::BUILD_INDEX, NActors::NLog::PRI_TRACE);
1302+
kikimr.GetTestServer().GetRuntime()->SetLogPriority(NKikimrServices::FLAT_TX_SCHEMESHARD, NActors::NLog::PRI_TRACE);
1303+
1304+
const int flags = F_NULLABLE | (Covered ? F_COVERING : 0) | (Overlap ? F_OVERLAP : 0);
1305+
auto db = kikimr.GetTableClient();
1306+
auto session = DoCreateTableForVectorIndex(db, flags);
1307+
1308+
{
1309+
const TString createIndex(Q_(Sprintf(R"(
1310+
ALTER TABLE `/Root/TestTable`
1311+
ADD INDEX index
1312+
GLOBAL USING vector_kmeans_tree
1313+
ON (emb)%s
1314+
WITH (distance=cosine, vector_type="uint8", vector_dimension=2, levels=1, clusters=2%s);
1315+
)",
1316+
(flags & F_COVERING ? " COVER (data, emb)" : ""),
1317+
(flags & F_OVERLAP ? ", overlap_clusters=2" : ""))));
1318+
1319+
auto result = session.ExecuteSchemeQuery(createIndex)
1320+
.ExtractValueSync();
1321+
1322+
UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
1323+
}
1324+
1325+
DoPositiveQueriesVectorIndexOrderByCosine(session);
1326+
1327+
UNIT_ASSERT(!IsTableEmpty(session, "/Root/TestTable"));
1328+
UNIT_ASSERT(!IsTableEmpty(session, "/Root/TestTable/index/indexImplPostingTable"));
1329+
UNIT_ASSERT(!IsTableEmpty(session, "/Root/TestTable/index/indexImplLevelTable"));
1330+
DoTruncateTable(session);
1331+
UNIT_ASSERT(IsTableEmpty(session, "/Root/TestTable"));
1332+
UNIT_ASSERT(IsTableEmpty(session, "/Root/TestTable/index/indexImplPostingTable"));
1333+
UNIT_ASSERT(!IsTableEmpty(session, "/Root/TestTable/index/indexImplLevelTable"));
1334+
1335+
DoOnlyUpsertValuesIntoTable(session);
1336+
DoPositiveQueriesVectorIndexOrderByCosine(session);
1337+
}
12691338
}
12701339

12711340
}

‎ydb/core/tx/schemeshard/schemeshard__operation_truncate_table.cpp‎

Lines changed: 53 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -331,9 +331,12 @@ class TTruncateTable: public TSubOperation {
331331
};
332332

333333
enum ESchemeObjectType {
334-
ETypeMainTable,
335-
ETypeIndex,
336-
ETypeIndexImplTable
334+
MainTable,
335+
GlobalVectorIndex, // there are special rules for processing global vector index
336+
PrefixVectorIndex, // there are special rules for processing prefix vector index
337+
CommonIndex, // using for secondary, unique and fulltext indexes
338+
IndexImplPrefixTable, // there are special rules for processing index impl prefix table
339+
CommonIndexImplTable // using for other index impl tables
337340
};
338341

339342
bool DfsOnTableChildrenTree(TOperationId opId, const TTxTransaction& tx, TOperationContext& context, const TPathId& currentPathId, TVector<ISubOperation::TPtr>& result, ESchemeObjectType objectType) {
@@ -369,7 +372,7 @@ bool DfsOnTableChildrenTree(TOperationId opId, const TTxTransaction& tx, TOperat
369372
// because in that case TRUNCATE could bring the database into an inconsistent state.
370373
// It is better to return an error to the user than to allow that to happen.
371374
switch (objectType) {
372-
case ESchemeObjectType::ETypeMainTable: {
375+
case ESchemeObjectType::MainTable: {
373376
for (const auto& [childName, childPathId] : currentPath.Base()->GetChildren()) {
374377
Y_ABORT_UNLESS(context.SS->PathsById.contains(childPathId));
375378
TPath srcChildPath = currentPath.Child(childName);
@@ -392,15 +395,26 @@ bool DfsOnTableChildrenTree(TOperationId opId, const TTxTransaction& tx, TOperat
392395
return false;
393396
}
394397
case NKikimrSchemeOp::EIndexTypeGlobalVectorKmeansTree: {
395-
// This condition should be removed in the near future
396-
result = {CreateReject(opId, NKikimrScheme::StatusPreconditionFailed, "Cannot truncate table with vector indexes")};
397-
return false;
398+
size_t numberOfChildren = srcChildPath.Base()->GetChildren().size();
399+
Y_ABORT_UNLESS(numberOfChildren == 2 || numberOfChildren == 3);
400+
401+
if (numberOfChildren == 2) { // GlobalVectorIndex contains two implTables
402+
if (!DfsOnTableChildrenTree(opId, tx, context, childPathId, result, ESchemeObjectType::GlobalVectorIndex)) {
403+
return false;
404+
}
405+
} else if (numberOfChildren == 3) { // PrefixVectorIndex contains three implTables
406+
if (!DfsOnTableChildrenTree(opId, tx, context, childPathId, result, ESchemeObjectType::PrefixVectorIndex)) {
407+
return false;
408+
}
409+
}
410+
411+
break;
398412
}
399413
case NKikimrSchemeOp::EIndexTypeGlobalFulltextPlain:
400414
case NKikimrSchemeOp::EIndexTypeGlobalFulltextRelevance:
401415
case NKikimrSchemeOp::EIndexTypeGlobal:
402416
case NKikimrSchemeOp::EIndexTypeGlobalUnique: {
403-
if (!DfsOnTableChildrenTree(opId, tx, context, childPathId, result, ESchemeObjectType::ETypeIndex)) {
417+
if (!DfsOnTableChildrenTree(opId, tx, context, childPathId, result, ESchemeObjectType::CommonIndex)) {
404418
return false;
405419
}
406420

@@ -430,7 +444,9 @@ bool DfsOnTableChildrenTree(TOperationId opId, const TTxTransaction& tx, TOperat
430444
break;
431445
}
432446

433-
case ESchemeObjectType::ETypeIndex: {
447+
case ESchemeObjectType::GlobalVectorIndex:
448+
case ESchemeObjectType::PrefixVectorIndex:
449+
case ESchemeObjectType::CommonIndex: {
434450
for (const auto& [childName, childPathId] : currentPath.Base()->GetChildren()) {
435451
Y_ABORT_UNLESS(context.SS->PathsById.contains(childPathId));
436452
TPath srcChildPath = currentPath.Child(childName);
@@ -439,10 +455,22 @@ bool DfsOnTableChildrenTree(TOperationId opId, const TTxTransaction& tx, TOperat
439455
continue;
440456
}
441457

458+
constexpr TStringBuf tableExcludedFromTruncate = "indexImplLevelTable";
459+
if (objectType == ESchemeObjectType::GlobalVectorIndex && srcChildPath.PathString().EndsWith(tableExcludedFromTruncate)) {
460+
continue;
461+
}
462+
442463
switch (srcChildPath.Base()->PathType) {
443-
case NKikimrSchemeOp::EPathType::EPathTypeTable: { // indexImplTables
444-
if (!DfsOnTableChildrenTree(opId, tx, context, childPathId, result, ESchemeObjectType::ETypeIndexImplTable)) {
445-
return false;
464+
case NKikimrSchemeOp::EPathType::EPathTypeTable: {
465+
constexpr TStringBuf prefixTableName = "indexImplPrefixTable";
466+
if (objectType == ESchemeObjectType::PrefixVectorIndex && srcChildPath.PathString().EndsWith(prefixTableName)) {
467+
if (!DfsOnTableChildrenTree(opId, tx, context, childPathId, result, ESchemeObjectType::IndexImplPrefixTable)) {
468+
return false;
469+
}
470+
} else {
471+
if (!DfsOnTableChildrenTree(opId, tx, context, childPathId, result, ESchemeObjectType::CommonIndexImplTable)) {
472+
return false;
473+
}
446474
}
447475

448476
break;
@@ -458,7 +486,18 @@ bool DfsOnTableChildrenTree(TOperationId opId, const TTxTransaction& tx, TOperat
458486
break;
459487
}
460488

461-
case ESchemeObjectType::ETypeIndexImplTable: {
489+
case ESchemeObjectType::IndexImplPrefixTable: {
490+
if (currentPath.Base()->GetChildren().size() != 1 ||
491+
currentPath.Child(currentPath.Base()->GetChildren().begin()->first).Base()->PathType != NKikimrSchemeOp::EPathType::EPathTypeSequence
492+
) {
493+
result = {CreateReject(opId, NKikimrScheme::StatusPreconditionFailed, "Index impl prefix tables can contain only sequence")};
494+
return false;
495+
}
496+
497+
break;
498+
}
499+
500+
case ESchemeObjectType::CommonIndexImplTable: {
462501
if (!currentPath.Base()->GetChildren().empty()) {
463502
result = {CreateReject(opId, NKikimrScheme::StatusPreconditionFailed, "Index impl tables cannot contain children")};
464503
return false;
@@ -512,7 +551,7 @@ TVector<ISubOperation::TPtr> CreateConsistentTruncateTable(TOperationId opId, co
512551
return result;
513552
}
514553

515-
DfsOnTableChildrenTree(opId, tx, context, mainTablePath.Base()->PathId, result, ESchemeObjectType::ETypeMainTable);
554+
DfsOnTableChildrenTree(opId, tx, context, mainTablePath.Base()->PathId, result, ESchemeObjectType::MainTable);
516555

517556
return result;
518557
}

0 commit comments

Comments
 (0)