Skip to content

Commit b7d9010

Browse files
vitalifvitalif
authored
Implement index-only searches with covering vector indexes (#17770) (#18137)
1 parent 4ad14e0 commit b7d9010

11 files changed

+288
-97
lines changed

‎ydb/core/base/table_index.cpp‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -154,7 +154,11 @@ bool IsCompatibleIndex(NKikimrSchemeOp::EIndexType indexType, const TTableColumn
154154
}
155155
tmp.clear();
156156
tmp.insert(table.Keys.begin(), table.Keys.end());
157-
tmp.insert(index.KeyColumns.begin(), index.KeyColumns.end() - (isSecondaryIndex ? 0 : 1));
157+
if (isSecondaryIndex) {
158+
tmp.insert(index.KeyColumns.begin(), index.KeyColumns.end());
159+
} else {
160+
// Vector indexes allow to add all columns both to index & data
161+
}
158162
if (const auto* broken = IsContains(index.DataColumns, tmp, true)) {
159163
explain = TStringBuilder()
160164
<< "the same column can't be used as key and data column for one index, for example " << *broken;

‎ydb/core/kqp/opt/logical/kqp_opt_log_indexes.cpp‎

Lines changed: 47 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,7 @@ using namespace NYql::NNodes;
1717

1818
namespace {
1919

20-
TCoAtomList BuildKeyColumnsList(const TKikimrTableDescription& /* table */, TPositionHandle pos, TExprContext& ctx,
21-
const auto& columnsToSelect) {
20+
TCoAtomList BuildKeyColumnsList(TPositionHandle pos, TExprContext& ctx, const auto& columnsToSelect) {
2221
TVector<TExprBase> columnsList;
2322
columnsList.reserve(columnsToSelect.size());
2423
for (auto column : columnsToSelect) {
@@ -35,7 +34,7 @@ TCoAtomList BuildKeyColumnsList(const TKikimrTableDescription& /* table */, TPos
3534
}
3635

3736
TCoAtomList BuildKeyColumnsList(const TKikimrTableDescription& table, TPositionHandle pos, TExprContext& ctx) {
38-
return BuildKeyColumnsList(table, pos, ctx, table.Metadata->KeyColumnNames);
37+
return BuildKeyColumnsList(pos, ctx, table.Metadata->KeyColumnNames);
3938
}
4039

4140
TCoAtomList MergeColumns(const NNodes::TCoAtomList& col1, const TVector<TString>& col2, TExprContext& ctx) {
@@ -306,21 +305,20 @@ struct TReadMatch {
306305
}
307306
};
308307

309-
template<typename TRead>
310-
bool CheckIndexCovering(const TRead& read, const TIntrusivePtr<TKikimrTableMetadata>& indexMeta) {
311-
for (const auto& col : read.Columns()) {
308+
bool CheckIndexCovering(const TCoAtomList& readColumns, const TIntrusivePtr<TKikimrTableMetadata>& indexMeta) {
309+
for (const auto& col : readColumns) {
312310
if (!indexMeta->Columns.contains(col.StringValue())) {
313-
return true;
311+
return false;
314312
}
315313
}
316-
return false;
314+
return true;
317315
}
318316

319317
TExprBase DoRewriteIndexRead(const TReadMatch& read, TExprContext& ctx,
320318
const TKikimrTableDescription& tableDesc, TIntrusivePtr<TKikimrTableMetadata> indexMeta,
321319
const TVector<TString>& extraColumns, const std::function<TExprBase(const TExprBase&)>& middleFilter = {})
322320
{
323-
const bool needDataRead = CheckIndexCovering(read, indexMeta);
321+
const bool isCovered = CheckIndexCovering(read.Columns(), indexMeta);
324322

325323
if (read.FullScan()) {
326324
const auto indexName = read.Index().StringValue();
@@ -329,7 +327,7 @@ TExprBase DoRewriteIndexRead(const TReadMatch& read, TExprContext& ctx,
329327
ctx.AddWarning(issue);
330328
}
331329

332-
if (!needDataRead) {
330+
if (isCovered) {
333331
// We can read all data from index table.
334332
auto ret = read.BuildRead(ctx, BuildTableMeta(*indexMeta, read.Pos(), ctx), read.Columns());
335333

@@ -517,26 +515,41 @@ void VectorReadLevel(
517515

518516
void VectorReadMain(
519517
TExprContext& ctx, TPositionHandle pos,
520-
const TKqpTable& postingTable, const TCoAtomList& postingColumns,
521-
const TKqpTable& mainTable, const TCoAtomList& mainColumns,
518+
const TKqpTable& postingTable,
519+
const TIntrusivePtr<TKikimrTableMetadata> & postingTableMeta,
520+
const TKqpTable& mainTable,
521+
const TIntrusivePtr<TKikimrTableMetadata> & mainTableMeta,
522+
const TCoAtomList& mainColumns,
522523
TExprNodePtr& read)
523524
{
524-
// TODO(mbkkt) handle covered index columns
525525
TKqpStreamLookupSettings settings;
526526
settings.Strategy = EStreamLookupStrategyType::LookupRows;
527-
read = Build<TKqlStreamLookupTable>(ctx, pos)
528-
.Table(postingTable)
529-
.LookupKeys(read)
530-
.Columns(postingColumns)
531-
.Settings(settings.BuildNode(ctx, pos))
532-
.Done().Ptr();
527+
const bool isCovered = CheckIndexCovering(mainColumns, postingTableMeta);
533528

534-
read = Build<TKqlStreamLookupTable>(ctx, pos)
535-
.Table(mainTable)
536-
.LookupKeys(read)
537-
.Columns(mainColumns)
538-
.Settings(settings.BuildNode(ctx, pos))
539-
.Done().Ptr();
529+
if (!isCovered) {
530+
const auto postingColumns = BuildKeyColumnsList(pos, ctx, mainTableMeta->KeyColumnNames);
531+
532+
read = Build<TKqlStreamLookupTable>(ctx, pos)
533+
.Table(postingTable)
534+
.LookupKeys(read)
535+
.Columns(postingColumns)
536+
.Settings(settings.BuildNode(ctx, pos))
537+
.Done().Ptr();
538+
539+
read = Build<TKqlStreamLookupTable>(ctx, pos)
540+
.Table(mainTable)
541+
.LookupKeys(read)
542+
.Columns(mainColumns)
543+
.Settings(settings.BuildNode(ctx, pos))
544+
.Done().Ptr();
545+
} else {
546+
read = Build<TKqlStreamLookupTable>(ctx, pos)
547+
.Table(postingTable)
548+
.LookupKeys(read)
549+
.Columns(mainColumns)
550+
.Settings(settings.BuildNode(ctx, pos))
551+
.Done().Ptr();
552+
}
540553
}
541554

542555
void VectorTopMain(TExprContext& ctx, const TCoTopBase& top, TExprNodePtr& read) {
@@ -568,9 +581,8 @@ TExprBase DoRewriteTopSortOverKMeansTree(
568581
const auto postingTable = BuildTableMeta(*postingTableDesc->Metadata, pos, ctx);
569582
const auto mainTable = BuildTableMeta(*tableDesc.Metadata, pos, ctx);
570583

571-
const auto levelColumns = BuildKeyColumnsList(*levelTableDesc, pos, ctx,
584+
const auto levelColumns = BuildKeyColumnsList(pos, ctx,
572585
std::initializer_list<std::string_view>{NTableIndex::NTableVectorKmeansTreeIndex::IdColumn, NTableIndex::NTableVectorKmeansTreeIndex::CentroidColumn});
573-
const auto postingColumns = BuildKeyColumnsList(*postingTableDesc, pos, ctx, tableDesc.Metadata->KeyColumnNames);
574586
const auto& mainColumns = match.Columns();
575587

576588
TNodeOnNodeOwnedMap replaces;
@@ -601,7 +613,7 @@ TExprBase DoRewriteTopSortOverKMeansTree(
601613

602614
VectorReadLevel(indexDesc, ctx, pos, kqpCtx, levelLambda, top, levelTable, levelColumns, read);
603615

604-
VectorReadMain(ctx, pos, postingTable, postingColumns, mainTable, mainColumns, read);
616+
VectorReadMain(ctx, pos, postingTable, postingTableDesc->Metadata, mainTable, tableDesc.Metadata, mainColumns, read);
605617

606618
if (flatMap) {
607619
read = Build<TCoFlatMap>(ctx, flatMap.Cast().Pos())
@@ -638,13 +650,12 @@ TExprBase DoRewriteTopSortOverPrefixedKMeansTree(
638650
const auto prefixTable = BuildTableMeta(*prefixTableDesc->Metadata, pos, ctx);
639651
const auto mainTable = BuildTableMeta(*tableDesc.Metadata, pos, ctx);
640652

641-
const auto levelColumns = BuildKeyColumnsList(*levelTableDesc, pos, ctx,
653+
const auto levelColumns = BuildKeyColumnsList(pos, ctx,
642654
std::initializer_list<std::string_view>{NTableIndex::NTableVectorKmeansTreeIndex::IdColumn, NTableIndex::NTableVectorKmeansTreeIndex::CentroidColumn});
643-
const auto postingColumns = BuildKeyColumnsList(*postingTableDesc, pos, ctx, tableDesc.Metadata->KeyColumnNames);
644655
const auto prefixColumns = [&] {
645656
auto columns = indexDesc.KeyColumns;
646657
columns.back().assign(NTableIndex::NTableVectorKmeansTreeIndex::IdColumn);
647-
return BuildKeyColumnsList(*prefixTableDesc, pos, ctx, columns);
658+
return BuildKeyColumnsList(pos, ctx, columns);
648659
}();
649660
const auto& mainColumns = match.Columns();
650661

@@ -688,7 +699,7 @@ TExprBase DoRewriteTopSortOverPrefixedKMeansTree(
688699

689700
VectorReadLevel(indexDesc, ctx, pos, kqpCtx, levelLambda, top, levelTable, levelColumns, read);
690701

691-
VectorReadMain(ctx, pos, postingTable, postingColumns, mainTable, mainColumns, read);
702+
VectorReadMain(ctx, pos, postingTable, postingTableDesc->Metadata, mainTable, tableDesc.Metadata, mainColumns, read);
692703

693704
if (mainLambda) {
694705
read = Build<TCoMap>(ctx, flatMap.Pos())
@@ -735,9 +746,9 @@ TExprBase KqpRewriteLookupIndex(const TExprBase& node, TExprContext& ctx, const
735746
YQL_ENSURE(indexDesc->Type != TIndexDescription::EType::GlobalSyncVectorKMeansTree,
736747
"lookup doesn't support vector index: " << indexName);
737748

738-
const bool needDataRead = CheckIndexCovering(lookupIndex, implTable);
749+
const bool isCovered = CheckIndexCovering(lookupIndex.Columns(), implTable);
739750

740-
if (!needDataRead) {
751+
if (isCovered) {
741752
TKqpStreamLookupSettings settings;
742753
settings.Strategy = EStreamLookupStrategyType::LookupRows;
743754
return Build<TKqlStreamLookupTable>(ctx, node.Pos())
@@ -785,8 +796,8 @@ TExprBase KqpRewriteStreamLookupIndex(const TExprBase& node, TExprContext& ctx,
785796
YQL_ENSURE(indexDesc->Type != TIndexDescription::EType::GlobalSyncVectorKMeansTree,
786797
"stream lookup doesn't support vector index: " << indexName);
787798

788-
const bool needDataRead = CheckIndexCovering(streamLookupIndex, implTable);
789-
if (!needDataRead) {
799+
const bool isCovered = CheckIndexCovering(streamLookupIndex.Columns(), implTable);
800+
if (isCovered) {
790801
return Build<TKqlStreamLookupTable>(ctx, node.Pos())
791802
.Table(BuildTableMeta(*implTable, node.Pos(), ctx))
792803
.LookupKeys(streamLookupIndex.LookupKeys())

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

Lines changed: 77 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,20 @@ using namespace NYdb::NTable;
2626

2727
Y_UNIT_TEST_SUITE(KqpPrefixedVectorIndexes) {
2828

29-
std::vector<i64> DoPositiveQueryVectorIndex(TSession& session, const TString& query) {
29+
std::vector<i64> DoPositiveQueryVectorIndex(TSession& session, const TString& query, bool covered = false) {
3030
{
3131
auto result = session.ExplainDataQuery(query).ExtractValueSync();
3232
UNIT_ASSERT_C(result.IsSuccess(),
3333
"Failed to explain: `" << query << "` with " << result.GetIssues().ToString());
34+
35+
if (covered) {
36+
// Check that the query doesn't use main table
37+
NJson::TJsonValue plan;
38+
NJson::ReadJsonTree(result.GetPlan(), &plan, true);
39+
UNIT_ASSERT(ValidatePlanNodeIds(plan));
40+
auto mainTableAccess = CountPlanNodesByKv(plan, "Table", "TestTable");
41+
UNIT_ASSERT_VALUES_EQUAL(mainTableAccess, 0);
42+
}
3443
}
3544
{
3645
auto result = session.ExecuteDataQuery(query,
@@ -53,7 +62,7 @@ Y_UNIT_TEST_SUITE(KqpPrefixedVectorIndexes) {
5362
}
5463
}
5564

56-
void DoPositiveQueriesVectorIndex(TSession& session, const TString& mainQuery, const TString& indexQuery) {
65+
void DoPositiveQueriesVectorIndex(TSession& session, const TString& mainQuery, const TString& indexQuery, bool covered = false) {
5766
auto toStr = [](const auto& rs) -> TString {
5867
TStringBuilder b;
5968
for (const auto& r : rs) {
@@ -66,7 +75,7 @@ Y_UNIT_TEST_SUITE(KqpPrefixedVectorIndexes) {
6675
UNIT_ASSERT_EQUAL_C(mainResults.size(), 3, toStr(mainResults));
6776
UNIT_ASSERT_C(std::unique(mainResults.begin(), mainResults.end()) == mainResults.end(), toStr(mainResults));
6877

69-
auto indexResults = DoPositiveQueryVectorIndex(session, indexQuery);
78+
auto indexResults = DoPositiveQueryVectorIndex(session, indexQuery, covered);
7079
absl::c_sort(indexResults);
7180
UNIT_ASSERT_EQUAL_C(indexResults.size(), 3, toStr(indexResults));
7281
UNIT_ASSERT_C(std::unique(indexResults.begin(), indexResults.end()) == indexResults.end(), toStr(indexResults));
@@ -79,13 +88,17 @@ Y_UNIT_TEST_SUITE(KqpPrefixedVectorIndexes) {
7988
std::string_view function,
8089
std::string_view direction,
8190
std::string_view left,
82-
std::string_view right) {
91+
std::string_view right,
92+
bool covered = false) {
8393
constexpr std::string_view init =
8494
"$target = \"\x67\x68\x03\";\n"
8595
"$user = \"user_b\";";
8696
std::string metric = std::format("Knn::{}({}, {})", function, left, right);
8797
// no metric in result
8898
{
99+
// TODO(vitaliff): Exclude index-covered WHERE fields from KqpReadTableRanges.
100+
// Currently even if we SELECT only pk, emb, data WHERE user=xxx we also get `user`
101+
// in SELECT columns and thus it's required to add it to covered columns.
89102
const TString plainQuery(Q1_(std::format(R"({}
90103
SELECT * FROM `/Root/TestTable`
91104
WHERE user = $user
@@ -100,7 +113,7 @@ Y_UNIT_TEST_SUITE(KqpPrefixedVectorIndexes) {
100113
ORDER BY {} {}
101114
LIMIT 3;
102115
)", init, metric, direction)));
103-
DoPositiveQueriesVectorIndex(session, plainQuery, indexQuery);
116+
DoPositiveQueriesVectorIndex(session, plainQuery, indexQuery, covered);
104117
}
105118
// metric in result
106119
{
@@ -117,7 +130,7 @@ Y_UNIT_TEST_SUITE(KqpPrefixedVectorIndexes) {
117130
ORDER BY {} {}
118131
LIMIT 3;
119132
)", init, metric, metric, direction)));
120-
DoPositiveQueriesVectorIndex(session, plainQuery, indexQuery);
133+
DoPositiveQueriesVectorIndex(session, plainQuery, indexQuery, covered);
121134
}
122135
// metric as result
123136
// TODO(mbkkt) fix this behavior too
@@ -136,27 +149,28 @@ Y_UNIT_TEST_SUITE(KqpPrefixedVectorIndexes) {
136149
ORDER BY m {}
137150
LIMIT 3;
138151
)", init, metric, direction)));
139-
DoPositiveQueriesVectorIndex(session, plainQuery, indexQuery);
152+
DoPositiveQueriesVectorIndex(session, plainQuery, indexQuery, covered);
140153
}
141154
}
142155

143156
void DoPositiveQueriesPrefixedVectorIndexOrderBy(
144157
TSession& session,
145158
std::string_view function,
146-
std::string_view direction) {
159+
std::string_view direction,
160+
bool covered = false) {
147161
// target is left, member is right
148-
DoPositiveQueriesPrefixedVectorIndexOrderBy(session, function, direction, "$target", "emb");
162+
DoPositiveQueriesPrefixedVectorIndexOrderBy(session, function, direction, "$target", "emb", covered);
149163
// target is right, member is left
150-
DoPositiveQueriesPrefixedVectorIndexOrderBy(session, function, direction, "emb", "$target");
164+
DoPositiveQueriesPrefixedVectorIndexOrderBy(session, function, direction, "emb", "$target", covered);
151165
}
152166

153-
void DoPositiveQueriesPrefixedVectorIndexOrderByCosine(TSession& session) {
167+
void DoPositiveQueriesPrefixedVectorIndexOrderByCosine(TSession& session, bool covered = false) {
154168
// distance, default direction
155-
DoPositiveQueriesPrefixedVectorIndexOrderBy(session, "CosineDistance", "");
169+
DoPositiveQueriesPrefixedVectorIndexOrderBy(session, "CosineDistance", "", covered);
156170
// distance, asc direction
157-
DoPositiveQueriesPrefixedVectorIndexOrderBy(session, "CosineDistance", "ASC");
171+
DoPositiveQueriesPrefixedVectorIndexOrderBy(session, "CosineDistance", "ASC", covered);
158172
// similarity, desc direction
159-
DoPositiveQueriesPrefixedVectorIndexOrderBy(session, "CosineSimilarity", "DESC");
173+
DoPositiveQueriesPrefixedVectorIndexOrderBy(session, "CosineSimilarity", "DESC", covered);
160174
}
161175

162176
TSession DoCreateTableForPrefixedVectorIndex(TTableClient& db, bool nullable) {
@@ -406,6 +420,7 @@ Y_UNIT_TEST_SUITE(KqpPrefixedVectorIndexes) {
406420
.ExtractValueSync();
407421

408422
UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
423+
// FIXME: result does not return failure/issues when index is created but fails to be filled with data
409424
}
410425
{
411426
auto result = session.DescribeTable("/Root/TestTable").ExtractValueSync();
@@ -425,6 +440,54 @@ Y_UNIT_TEST_SUITE(KqpPrefixedVectorIndexes) {
425440
DoPositiveQueriesPrefixedVectorIndexOrderByCosine(session);
426441
}
427442

443+
Y_UNIT_TEST_TWIN(PrefixedVectorIndexOrderByCosineDistanceWithCover, Nullable) {
444+
NKikimrConfig::TFeatureFlags featureFlags;
445+
featureFlags.SetEnableVectorIndex(true);
446+
auto setting = NKikimrKqp::TKqpSetting();
447+
auto serverSettings = TKikimrSettings()
448+
.SetFeatureFlags(featureFlags)
449+
.SetKqpSettings({setting});
450+
451+
TKikimrRunner kikimr(serverSettings);
452+
kikimr.GetTestServer().GetRuntime()->SetLogPriority(NKikimrServices::BUILD_INDEX, NActors::NLog::PRI_TRACE);
453+
kikimr.GetTestServer().GetRuntime()->SetLogPriority(NKikimrServices::FLAT_TX_SCHEMESHARD, NActors::NLog::PRI_TRACE);
454+
455+
auto db = kikimr.GetTableClient();
456+
auto session = DoCreateTableForPrefixedVectorIndex(db, Nullable);
457+
{
458+
const TString createIndex(Q_(R"(
459+
ALTER TABLE `/Root/TestTable`
460+
ADD INDEX index
461+
GLOBAL USING vector_kmeans_tree
462+
ON (user, emb) COVER (user, emb, data)
463+
WITH (distance=cosine, vector_type="uint8", vector_dimension=2, levels=2, clusters=2);
464+
)"));
465+
466+
auto result = session.ExecuteSchemeQuery(createIndex)
467+
.ExtractValueSync();
468+
469+
UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
470+
}
471+
{
472+
auto result = session.DescribeTable("/Root/TestTable").ExtractValueSync();
473+
UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), NYdb::EStatus::SUCCESS);
474+
const auto& indexes = result.GetTableDescription().GetIndexDescriptions();
475+
UNIT_ASSERT_EQUAL(indexes.size(), 1);
476+
UNIT_ASSERT_EQUAL(indexes[0].GetIndexName(), "index");
477+
std::vector<std::string> indexKeyColumns{"user", "emb"};
478+
UNIT_ASSERT_EQUAL(indexes[0].GetIndexColumns(), indexKeyColumns);
479+
std::vector<std::string> indexDataColumns{"user", "emb", "data"};
480+
UNIT_ASSERT_EQUAL(indexes[0].GetDataColumns(), indexDataColumns);
481+
const auto& settings = std::get<TKMeansTreeSettings>(indexes[0].GetIndexSettings());
482+
UNIT_ASSERT_EQUAL(settings.Settings.Metric, NYdb::NTable::TVectorIndexSettings::EMetric::CosineDistance);
483+
UNIT_ASSERT_EQUAL(settings.Settings.VectorType, NYdb::NTable::TVectorIndexSettings::EVectorType::Uint8);
484+
UNIT_ASSERT_EQUAL(settings.Settings.VectorDimension, 2);
485+
UNIT_ASSERT_EQUAL(settings.Levels, 2);
486+
UNIT_ASSERT_EQUAL(settings.Clusters, 2);
487+
}
488+
DoPositiveQueriesPrefixedVectorIndexOrderByCosine(session, true /*covered*/);
489+
}
490+
428491
}
429492

430493
}

0 commit comments

Comments
 (0)