Implement index-only searches with covering vector indexes (#17770) (#18137)

Author: azevaykin

ydb/core/base/table_index.cpp

@@ -154,7 +154,11 @@ bool IsCompatibleIndex(NKikimrSchemeOp::EIndexType indexType, const TTableColumn
     }
     tmp.clear();
     tmp.insert(table.Keys.begin(), table.Keys.end());
-    tmp.insert(index.KeyColumns.begin(), index.KeyColumns.end() - (isSecondaryIndex ? 0 : 1));
+    if (isSecondaryIndex) {
+        tmp.insert(index.KeyColumns.begin(), index.KeyColumns.end());
+    } else {
+        // Vector indexes allow to add all columns both to index & data
+    }
     if (const auto* broken = IsContains(index.DataColumns, tmp, true)) {
         explain = TStringBuilder()
                   << "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

@@ -17,8 +17,7 @@ using namespace NYql::NNodes;
 
 namespace {
 
-TCoAtomList BuildKeyColumnsList(const TKikimrTableDescription& /* table */, TPositionHandle pos, TExprContext& ctx,
-                                const auto& columnsToSelect) {
+TCoAtomList BuildKeyColumnsList(TPositionHandle pos, TExprContext& ctx, const auto& columnsToSelect) {
     TVector<TExprBase> columnsList;
     columnsList.reserve(columnsToSelect.size());
     for (auto column : columnsToSelect) {
@@ -35,7 +34,7 @@ TCoAtomList BuildKeyColumnsList(const TKikimrTableDescription& /* table */, TPos
 }
 
 TCoAtomList BuildKeyColumnsList(const TKikimrTableDescription& table, TPositionHandle pos, TExprContext& ctx) {
-    return BuildKeyColumnsList(table, pos, ctx, table.Metadata->KeyColumnNames);
+    return BuildKeyColumnsList(pos, ctx, table.Metadata->KeyColumnNames);
 }
 
 TCoAtomList MergeColumns(const NNodes::TCoAtomList& col1, const TVector<TString>& col2, TExprContext& ctx) {
@@ -306,21 +305,20 @@ struct TReadMatch {
     }
 };
 
-template<typename TRead>
-bool CheckIndexCovering(const TRead& read, const TIntrusivePtr<TKikimrTableMetadata>& indexMeta) {
-    for (const auto& col : read.Columns()) {
+bool CheckIndexCovering(const TCoAtomList& readColumns, const TIntrusivePtr<TKikimrTableMetadata>& indexMeta) {
+    for (const auto& col : readColumns) {
         if (!indexMeta->Columns.contains(col.StringValue())) {
-            return true;
+            return false;
         }
     }
-    return false;
+    return true;
 }
 
 TExprBase DoRewriteIndexRead(const TReadMatch& read, TExprContext& ctx,
     const TKikimrTableDescription& tableDesc, TIntrusivePtr<TKikimrTableMetadata> indexMeta, bool useStreamLookup,
     const TVector<TString>& extraColumns, const std::function<TExprBase(const TExprBase&)>& middleFilter = {})
 {
-    const bool needDataRead = CheckIndexCovering(read, indexMeta);
+    const bool isCovered = CheckIndexCovering(read.Columns(), indexMeta);
 
     if (read.FullScan()) {
         const auto indexName = read.Index().StringValue();
@@ -329,7 +327,7 @@ TExprBase DoRewriteIndexRead(const TReadMatch& read, TExprContext& ctx,
         ctx.AddWarning(issue);
     }
 
-    if (!needDataRead) {
+    if (isCovered) {
         // We can read all data from index table.
         auto ret = read.BuildRead(ctx, BuildTableMeta(*indexMeta, read.Pos(), ctx), read.Columns());
 
@@ -522,22 +520,37 @@ void VectorReadLevel(
 
 void VectorReadMain(
     TExprContext& ctx, TPositionHandle pos,
-    const TKqpTable& postingTable, const TCoAtomList& postingColumns,
-    const TKqpTable& mainTable, const TCoAtomList& mainColumns,
+    const TKqpTable& postingTable,
+    const TIntrusivePtr<TKikimrTableMetadata> & postingTableMeta,
+    const TKqpTable& mainTable,
+    const TIntrusivePtr<TKikimrTableMetadata> & mainTableMeta,
+    const TCoAtomList& mainColumns,
     TExprNodePtr& read)
 {
-    // TODO(mbkkt) handle covered index columns
-    read = Build<TKqlLookupTable>(ctx, pos)
-        .Table(postingTable)
-        .LookupKeys(read)
-        .Columns(postingColumns)
-    .Done().Ptr();
+    const bool isCovered = CheckIndexCovering(mainColumns, postingTableMeta);
 
-    read = Build<TKqlLookupTable>(ctx, pos)
-        .Table(mainTable)
-        .LookupKeys(read)
-        .Columns(mainColumns)
-    .Done().Ptr();
+    if (!isCovered) {
+        const auto postingColumns = BuildKeyColumnsList(pos, ctx, mainTableMeta->KeyColumnNames);
+
+        read = Build<TKqlLookupTable>(ctx, pos)
+            .Table(postingTable)
+            .LookupKeys(read)
+            .Columns(postingColumns)
+            .Done().Ptr();
+
+        read = Build<TKqlLookupTable>(ctx, pos)
+            .Table(mainTable)
+            .LookupKeys(read)
+            .Columns(mainColumns)
+            .Done().Ptr();
+    } else {
+        read = Build<TKqlStreamLookupTable>(ctx, pos)
+            .Table(postingTable)
+            .LookupKeys(read)
+            .Columns(mainColumns)
+            .Settings(settings.BuildNode(ctx, pos))
+        .Done().Ptr();
+    }
 }
 
 void VectorTopMain(TExprContext& ctx, const TCoTopBase& top, TExprNodePtr& read) {
@@ -569,9 +582,8 @@ TExprBase DoRewriteTopSortOverKMeansTree(
     const auto postingTable = BuildTableMeta(*postingTableDesc->Metadata, pos, ctx);
     const auto mainTable = BuildTableMeta(*tableDesc.Metadata, pos, ctx);
 
-    const auto levelColumns = BuildKeyColumnsList(*levelTableDesc, pos, ctx,
+    const auto levelColumns = BuildKeyColumnsList(pos, ctx,
             std::initializer_list<std::string_view>{NTableIndex::NTableVectorKmeansTreeIndex::IdColumn, NTableIndex::NTableVectorKmeansTreeIndex::CentroidColumn});
-    const auto postingColumns = BuildKeyColumnsList(*postingTableDesc, pos, ctx, tableDesc.Metadata->KeyColumnNames);
     const auto& mainColumns = match.Columns();
 
     TNodeOnNodeOwnedMap replaces;
@@ -602,7 +614,7 @@ TExprBase DoRewriteTopSortOverKMeansTree(
 
     VectorReadLevel(indexDesc, ctx, pos, kqpCtx, levelLambda, top, levelTable, levelColumns, read);
 
-    VectorReadMain(ctx, pos, postingTable, postingColumns, mainTable, mainColumns, read);
+    VectorReadMain(ctx, pos, postingTable, postingTableDesc->Metadata, mainTable, tableDesc.Metadata, mainColumns, read);
 
     if (flatMap) {
         read = Build<TCoFlatMap>(ctx, flatMap.Cast().Pos())
@@ -639,13 +651,12 @@ TExprBase DoRewriteTopSortOverPrefixedKMeansTree(
     const auto prefixTable = BuildTableMeta(*prefixTableDesc->Metadata, pos, ctx);
     const auto mainTable = BuildTableMeta(*tableDesc.Metadata, pos, ctx);
 
-    const auto levelColumns = BuildKeyColumnsList(*levelTableDesc, pos, ctx,
+    const auto levelColumns = BuildKeyColumnsList(pos, ctx,
             std::initializer_list<std::string_view>{NTableIndex::NTableVectorKmeansTreeIndex::IdColumn, NTableIndex::NTableVectorKmeansTreeIndex::CentroidColumn});
-    const auto postingColumns = BuildKeyColumnsList(*postingTableDesc, pos, ctx, tableDesc.Metadata->KeyColumnNames);
     const auto prefixColumns = [&] {
         auto columns = indexDesc.KeyColumns;
         columns.back().assign(NTableIndex::NTableVectorKmeansTreeIndex::IdColumn);
-        return BuildKeyColumnsList(*prefixTableDesc, pos, ctx, columns);
+        return BuildKeyColumnsList(pos, ctx, columns);
     }();
     const auto& mainColumns = match.Columns();
 
@@ -686,7 +697,7 @@ TExprBase DoRewriteTopSortOverPrefixedKMeansTree(
 
     VectorReadLevel(indexDesc, ctx, pos, kqpCtx, levelLambda, top, levelTable, levelColumns, read);
 
-    VectorReadMain(ctx, pos, postingTable, postingColumns, mainTable, mainColumns, read);
+    VectorReadMain(ctx, pos, postingTable, postingTableDesc->Metadata, mainTable, tableDesc.Metadata, mainColumns, read);
 
     if (mainLambda) {
         read = Build<TCoMap>(ctx, flatMap.Pos())
@@ -733,9 +744,9 @@ TExprBase KqpRewriteLookupIndex(const TExprBase& node, TExprContext& ctx, const
         YQL_ENSURE(indexDesc->Type != TIndexDescription::EType::GlobalSyncVectorKMeansTree,
             "lookup doesn't support vector index: " << indexName);
 
-        const bool needDataRead = CheckIndexCovering(lookupIndex, implTable);
+        const bool isCovered = CheckIndexCovering(lookupIndex.Columns(), implTable);
 
-        if (!needDataRead) {
+        if (isCovered) {
             if (kqpCtx.Config->EnableKqpDataQueryStreamLookup) {
                 TKqpStreamLookupSettings settings;
                 settings.Strategy = EStreamLookupStrategyType::LookupRows;
@@ -805,8 +816,8 @@ TExprBase KqpRewriteStreamLookupIndex(const TExprBase& node, TExprContext& ctx,
     YQL_ENSURE(indexDesc->Type != TIndexDescription::EType::GlobalSyncVectorKMeansTree,
         "stream lookup doesn't support vector index: " << indexName);
 
-    const bool needDataRead = CheckIndexCovering(streamLookupIndex, implTable);
-    if (!needDataRead) {
+    const bool isCovered = CheckIndexCovering(streamLookupIndex.Columns(), implTable);
+    if (isCovered) {
         return Build<TKqlStreamLookupTable>(ctx, node.Pos())
             .Table(BuildTableMeta(*implTable, node.Pos(), ctx))
             .LookupKeys(streamLookupIndex.LookupKeys())