From 30a6d49140df7fb3ea05741118823ef5c58b5e84 Mon Sep 17 00:00:00 2001 From: Yahor Yuzefovich Date: Wed, 26 Nov 2025 15:10:07 -0800 Subject: [PATCH] sql: record vector index usage of index_usage_stats This commit records the vector index reads for `crdb_internal.index_usage_stats`. It also removes some duplicated code in favor of using the shared helper. It also fixes the omission of vector index usage for purposes of IndexesUsed which we include into statement statistics. Also the omission of the index for PlaceholderScan. Release note: None --- pkg/sql/distsql_spec_exec_factory.go | 12 ++-- pkg/sql/opt/exec/execbuilder/relational.go | 7 +- pkg/sql/opt_exec_factory.go | 78 +++++++++------------- 3 files changed, 42 insertions(+), 55 deletions(-) diff --git a/pkg/sql/distsql_spec_exec_factory.go b/pkg/sql/distsql_spec_exec_factory.go index b4f631d7b829..193cfe2d071b 100644 --- a/pkg/sql/distsql_spec_exec_factory.go +++ b/pkg/sql/distsql_spec_exec_factory.go @@ -317,6 +317,7 @@ func (e *distSQLSpecExecFactory) ConstructScan( Reverse: params.Reverse, TableDescriptorModificationTime: tabDesc.GetModificationTime(), } + // TODO(yuzefovich): record the index usage when applicable. if err := rowenc.InitIndexFetchSpec(&trSpec.FetchSpec, e.planner.ExecCfg().Codec, tabDesc, idx, columnIDs); err != nil { return nil, err } @@ -871,7 +872,7 @@ func (e *distSQLSpecExecFactory) ConstructIndexJoin( fetch.lockingWaitPolicy = descpb.ToScanLockingWaitPolicy(locking.WaitPolicy) fetch.lockingDurability = descpb.ToScanLockingDurability(locking.Durability) - // TODO(drewk): in an EXPLAIN context, record the index usage. + // TODO(drewk): record the index usage when applicable. planInfo := &indexJoinPlanningInfo{ fetch: fetch, keyCols: keyCols, @@ -954,7 +955,7 @@ func (e *distSQLSpecExecFactory) ConstructLookupJoin( fetch.lockingWaitPolicy = descpb.ToScanLockingWaitPolicy(locking.WaitPolicy) fetch.lockingDurability = descpb.ToScanLockingDurability(locking.Durability) - // TODO(drewk): if in an EXPLAIN context, record the index usage. + // TODO(drewk): record the index usage when applicable. planInfo := &lookupJoinPlanningInfo{ fetch: fetch, joinType: joinType, @@ -1031,7 +1032,7 @@ func (e *distSQLSpecExecFactory) ConstructInvertedJoin( fetch.lockingWaitPolicy = descpb.ToScanLockingWaitPolicy(locking.WaitPolicy) fetch.lockingDurability = descpb.ToScanLockingDurability(locking.Durability) - // TODO(drewk): if we're in an EXPLAIN context, record the index usage. + // TODO(drewk): record the index usage when applicable. planInfo := &invertedJoinPlanningInfo{ fetch: fetch, joinType: joinType, @@ -1093,8 +1094,7 @@ func (e *distSQLSpecExecFactory) constructZigzagJoinSide( return zigzagPlanningSide{}, err } - // TODO (cucaroach): update indexUsageStats. - + // TODO(yuzefovich): record the index usage when applicable. return zigzagPlanningSide{ desc: desc, index: index.(*optIndex).idx, @@ -1524,6 +1524,7 @@ func (e *distSQLSpecExecFactory) ConstructVectorSearch( if err != nil { return nil, err } + // TODO(yuzefovich): record the index usage when applicable. planInfo := &vectorSearchPlanningInfo{ table: tabDesc, index: indexDesc, @@ -1562,6 +1563,7 @@ func (e *distSQLSpecExecFactory) ConstructVectorMutationSearch( if isIndexPut { cols = append(cols, colinfo.ResultColumn{Name: "quantized-vector", Typ: types.Bytes}) } + // TODO(yuzefovich): record the index usage when applicable. planInfo := &vectorMutationSearchPlanningInfo{ table: table.(*optTable).desc, index: index.(*optIndex).idx, diff --git a/pkg/sql/opt/exec/execbuilder/relational.go b/pkg/sql/opt/exec/execbuilder/relational.go index 4862d9e6bf3b..4ad413e98370 100644 --- a/pkg/sql/opt/exec/execbuilder/relational.go +++ b/pkg/sql/opt/exec/execbuilder/relational.go @@ -944,7 +944,7 @@ func (b *Builder) buildScan(scan *memo.ScanExpr) (_ execPlan, outputCols colOrdM } root, err := b.factory.ConstructScan( tab, - tab.Index(scan.Index), + idx, params, reqOrdering, ) @@ -975,6 +975,7 @@ func (b *Builder) buildPlaceholderScan( md := b.mem.Metadata() tab := md.Table(scan.Table) idx := tab.Index(scan.Index) + b.IndexesUsed.add(tab.ID(), idx.ID()) // Build the index constraint. spanColumns := make([]opt.OrderingColumn, len(scan.Span)) @@ -1026,7 +1027,7 @@ func (b *Builder) buildPlaceholderScan( } root, err := b.factory.ConstructScan( tab, - tab.Index(scan.Index), + idx, params, reqOrdering, ) @@ -4117,6 +4118,7 @@ func (b *Builder) buildVectorSearch( return execPlan{}, colOrdMap{}, errors.AssertionFailedf( "vector search is only supported on vector indexes") } + b.IndexesUsed.add(table.ID(), index.ID()) primaryKeyCols := md.TableMeta(search.Table).IndexKeyColumns(cat.PrimaryIndex) for col, ok := search.Cols.Next(0); ok; col, ok = search.Cols.Next(col + 1) { if !primaryKeyCols.Contains(col) { @@ -4163,6 +4165,7 @@ func (b *Builder) buildVectorMutationSearch( return execPlan{}, colOrdMap{}, errors.AssertionFailedf( "vector mutation search is only supported on vector indexes") } + b.IndexesUsed.add(table.ID(), index.ID()) input, inputCols, err := b.buildRelational(search.Input) if err != nil { diff --git a/pkg/sql/opt_exec_factory.go b/pkg/sql/opt_exec_factory.go index 21360f9506b6..661b07187122 100644 --- a/pkg/sql/opt_exec_factory.go +++ b/pkg/sql/opt_exec_factory.go @@ -111,6 +111,18 @@ func (ef *execFactory) ConstructLiteralValues( } } +// recordIndexRead - if applicable - records the fact that the given index has +// been used for reading. +func (ef *execFactory) recordIndexRead(tabDesc catalog.TableDescriptor, idx catalog.Index) { + if !ef.isExplain && !ef.planner.SessionData().Internal { + idxUsageKey := roachpb.IndexUsageKey{ + TableID: roachpb.TableID(tabDesc.GetID()), + IndexID: roachpb.IndexID(idx.GetID()), + } + ef.planner.extendedEvalCtx.indexUsageStats.RecordRead(idxUsageKey) + } +} + // ConstructScan is part of the exec.Factory interface. func (ef *execFactory) ConstructScan( table cat.Table, index cat.Index, params exec.ScanParams, reqOrdering exec.OutputOrdering, @@ -158,13 +170,7 @@ func (ef *execFactory) ConstructScan( scan.lockingWaitPolicy = descpb.ToScanLockingWaitPolicy(params.Locking.WaitPolicy) scan.lockingDurability = descpb.ToScanLockingDurability(params.Locking.Durability) scan.localityOptimized = params.LocalityOptimized - if !ef.isExplain && !ef.planner.SessionData().Internal { - idxUsageKey := roachpb.IndexUsageKey{ - TableID: roachpb.TableID(tabDesc.GetID()), - IndexID: roachpb.IndexID(idx.GetID()), - } - ef.planner.extendedEvalCtx.indexUsageStats.RecordRead(idxUsageKey) - } + ef.recordIndexRead(tabDesc, idx) return scan, nil } @@ -681,19 +687,12 @@ func (ef *execFactory) ConstructIndexJoin( } idx := tabDesc.GetPrimaryIndex() + ef.recordIndexRead(tabDesc, idx) fetch.index = idx fetch.lockingStrength = descpb.ToScanLockingStrength(locking.Strength) fetch.lockingWaitPolicy = descpb.ToScanLockingWaitPolicy(locking.WaitPolicy) fetch.lockingDurability = descpb.ToScanLockingDurability(locking.Durability) - if !ef.isExplain && !ef.planner.SessionData().Internal { - idxUsageKey := roachpb.IndexUsageKey{ - TableID: roachpb.TableID(tabDesc.GetID()), - IndexID: roachpb.IndexID(idx.GetID()), - } - ef.planner.extendedEvalCtx.indexUsageStats.RecordRead(idxUsageKey) - } - n := &indexJoinNode{ singleInputPlanNode: singleInputPlanNode{input.(planNode)}, columns: fetch.columns, @@ -744,19 +743,12 @@ func (ef *execFactory) ConstructLookupJoin( return nil, err } + ef.recordIndexRead(tabDesc, idx) fetch.index = idx fetch.lockingStrength = descpb.ToScanLockingStrength(locking.Strength) fetch.lockingWaitPolicy = descpb.ToScanLockingWaitPolicy(locking.WaitPolicy) fetch.lockingDurability = descpb.ToScanLockingDurability(locking.Durability) - if !ef.isExplain && !ef.planner.SessionData().Internal { - idxUsageKey := roachpb.IndexUsageKey{ - TableID: roachpb.TableID(tabDesc.GetID()), - IndexID: roachpb.IndexID(idx.GetID()), - } - ef.planner.extendedEvalCtx.indexUsageStats.RecordRead(idxUsageKey) - } - n := &lookupJoinNode{ singleInputPlanNode: singleInputPlanNode{input.(planNode)}, lookupJoinPlanningInfo: lookupJoinPlanningInfo{ @@ -881,19 +873,12 @@ func (ef *execFactory) ConstructInvertedJoin( if err := fetch.initDescDefaults(tabDesc, colCfg); err != nil { return nil, err } + ef.recordIndexRead(tabDesc, idx) fetch.index = idx fetch.lockingStrength = descpb.ToScanLockingStrength(locking.Strength) fetch.lockingWaitPolicy = descpb.ToScanLockingWaitPolicy(locking.WaitPolicy) fetch.lockingDurability = descpb.ToScanLockingDurability(locking.Durability) - if !ef.isExplain && !ef.planner.SessionData().Internal { - idxUsageKey := roachpb.IndexUsageKey{ - TableID: roachpb.TableID(tabDesc.GetID()), - IndexID: roachpb.IndexID(idx.GetID()), - } - ef.planner.extendedEvalCtx.indexUsageStats.RecordRead(idxUsageKey) - } - n := &invertedJoinNode{ singleInputPlanNode: singleInputPlanNode{input.(planNode)}, invertedJoinPlanningInfo: invertedJoinPlanningInfo{ @@ -954,22 +939,15 @@ func (ef *execFactory) constructFetchForZigzag( return fetchPlanningInfo{}, nil, err } - tableDesc := table.(*optTable).desc - idxDesc := index.(*optIndex).idx + tabDesc := table.(*optTable).desc + idx := index.(*optIndex).idx var fetch fetchPlanningInfo - if err := fetch.initDescDefaults(tableDesc, colCfg); err != nil { + if err := fetch.initDescDefaults(tabDesc, colCfg); err != nil { return fetchPlanningInfo{}, nil, err } - if !ef.isExplain && !ef.planner.SessionData().Internal { - idxUsageKey := roachpb.IndexUsageKey{ - TableID: roachpb.TableID(tableDesc.GetID()), - IndexID: roachpb.IndexID(idxDesc.GetID()), - } - ef.planner.extendedEvalCtx.indexUsageStats.RecordRead(idxUsageKey) - } - - fetch.index = idxDesc + ef.recordIndexRead(tabDesc, idx) + fetch.index = idx fetch.lockingStrength = descpb.ToScanLockingStrength(locking.Strength) fetch.lockingWaitPolicy = descpb.ToScanLockingWaitPolicy(locking.WaitPolicy) fetch.lockingDurability = descpb.ToScanLockingDurability(locking.Durability) @@ -1962,23 +1940,24 @@ func (ef *execFactory) ConstructVectorSearch( targetNeighborCount uint64, ) (exec.Node, error) { tabDesc := table.(*optTable).desc - indexDesc := index.(*optIndex).idx + idx := index.(*optIndex).idx cols := makeColList(table, outCols) resultCols := colinfo.ResultColumnsFromColumns(tabDesc.GetID(), cols) // Encode the prefix constraint as a list of roachpb.Keys. var sb span.Builder sb.InitAllowingExternalRowData( - ef.planner.EvalContext(), ef.planner.ExecCfg().Codec, tabDesc, indexDesc, + ef.planner.EvalContext(), ef.planner.ExecCfg().Codec, tabDesc, idx, ) prefixKeys, err := sb.KeysFromVectorPrefixConstraint(ef.ctx, prefixConstraint) if err != nil { return nil, err } + ef.recordIndexRead(tabDesc, idx) return &vectorSearchNode{ vectorSearchPlanningInfo: vectorSearchPlanningInfo{ table: tabDesc, - index: indexDesc, + index: idx, prefixKeys: prefixKeys, queryVector: queryVector, targetNeighborCount: targetNeighborCount, @@ -2009,11 +1988,14 @@ func (ef *execFactory) ConstructVectorMutationSearch( cols = append(cols, colinfo.ResultColumn{Name: "quantized-vector", Typ: types.Bytes}) } + tabDesc := table.(*optTable).desc + idx := index.(*optIndex).idx + ef.recordIndexRead(tabDesc, idx) return &vectorMutationSearchNode{ singleInputPlanNode: singleInputPlanNode{input: inputPlan}, vectorMutationSearchPlanningInfo: vectorMutationSearchPlanningInfo{ - table: table.(*optTable).desc, - index: index.(*optIndex).idx, + table: tabDesc, + index: idx, prefixKeyCols: prefixKeyCols, queryVectorCol: queryVectorCol, suffixKeyCols: suffixKeyCols,