Merge pull request #160404 from yuzefovich/blathers/backport-release-25.4-158416

release-25.4: sql: record vector index usage of index_usage_stats

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

@@ -956,7 +956,7 @@ func (b *Builder) buildScan(scan *memo.ScanExpr) (_ execPlan, outputCols colOrdM
 	}
 	root, err := b.factory.ConstructScan(
 		tab,
-		tab.Index(scan.Index),
+		idx,
 		params,
 		reqOrdering,
 	)
@@ -1014,6 +1014,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))
@@ -1051,7 +1052,7 @@ func (b *Builder) buildPlaceholderScan(
 	}
 	root, err := b.factory.ConstructScan(
 		tab,
-		tab.Index(scan.Index),
+		idx,
 		params,
 		reqOrdering,
 	)
@@ -4144,6 +4145,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) {
@@ -4190,6 +4192,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
 }
@@ -693,19 +699,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,
@@ -756,19 +755,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{
@@ -893,19 +885,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{
@@ -966,22 +951,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)
@@ -2068,23 +2046,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,
@@ -2115,11 +2094,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,