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

Author: yuzefovich

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,

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,
 	)
@@ -1002,6 +1002,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))
@@ -1039,7 +1040,7 @@ func (b *Builder) buildPlaceholderScan(
 	}
 	root, err := b.factory.ConstructScan(
 		tab,
-		tab.Index(scan.Index),
+		idx,
 		params,
 		reqOrdering,
 	)
@@ -4115,6 +4116,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) {
@@ -4161,6 +4163,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 {

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,