Skip to content

sql: record vector index usage of index_usage_stats #158416

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
+42 −55
Open

sql: record vector index usage of index_usage_stats #158416

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions pkg/sql/distsql_spec_exec_factory.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
7 changes: 5 additions & 2 deletions pkg/sql/opt/exec/execbuilder/relational.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -1026,7 +1027,7 @@ func (b *Builder) buildPlaceholderScan(
}
root, err := b.factory.ConstructScan(
tab,
tab.Index(scan.Index),
idx,
params,
reqOrdering,
)
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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 {
Expand Down
78 changes: 30 additions & 48 deletions pkg/sql/opt_exec_factory.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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{
Expand Down Expand Up @@ -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{
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down