sql: record index usage stats in index joins

Fixes #76173

Previously, the construction of index joins skipped the
recording of the  primary key for index usage stats.
This commit records the use of the primary key if the
statement is not of the EXPLAIN variety.

Release note (bug fix): Index usage stats are now properly
captured for index joins.

pkg/ccl/serverccl/statusccl/tenant_status_test.go

@@ -439,6 +439,7 @@ VALUES (1, 10, 100), (2, 20, 200), (3, 30, 300)
 					Exec(t, "SELECT * FROM test")
 
 				// Record scan on secondary index.
+				// Note that this is an index join and will also read from the primary index.
 				cluster.tenantConn(randomServer).
 					Exec(t, "SELECT * FROM test@test_a_idx")
 				testTableIDStr := cluster.tenantConn(randomServer).
@@ -463,7 +464,7 @@ WHERE
   table_id = ` + testTableIDStr
 				// Assert index usage data was inserted.
 				expected := [][]string{
-					{testTableIDStr, "1", "1", "true"},
+					{testTableIDStr, "1", "2", "true"}, // Primary index
 					{testTableIDStr, "2", "1", "true"},
 				}
 				cluster.tenantConn(randomServer).CheckQueryResults(t, query, expected)
@@ -771,6 +772,7 @@ VALUES (1, 10, 100), (2, 20, 200), (3, 30, 300)
 	testingCluster.tenantConn(0).Exec(t, "SELECT * FROM idx_test.test")
 
 	// Record scan on secondary index.
+	// Note that this is an index join and will also read from the primary index.
 	testingCluster.tenantConn(1).Exec(t, "SELECT * FROM idx_test.test@test_a_idx")
 	testTableIDStr := testingCluster.tenantConn(2).QueryStr(t, "SELECT 'idx_test.test'::regclass::oid")[0][0]
 	testTableID, err := strconv.Atoi(testTableIDStr)
@@ -789,7 +791,7 @@ WHERE
 `
 	actual := testingCluster.tenantConn(2).QueryStr(t, query, testTableID)
 	expected := [][]string{
-		{testTableIDStr, "1", "1", "true"},
+		{testTableIDStr, "1", "2", "true"},
 		{testTableIDStr, "2", "1", "true"},
 	}
 

pkg/server/index_usage_stats_test.go

@@ -76,6 +76,11 @@ func TestStatusAPIIndexUsage(t *testing.T) {
 		LastRead:       timeutil.Now(),
 	}
 
+	expectedStatsIndexPrimary := roachpb.IndexUsageStatistics{
+		TotalReadCount: 1,
+		LastRead:       timeutil.Now(),
+	}
+
 	firstPgURL, firstServerConnCleanup := sqlutils.PGUrl(
 		t, firstServer.ServingSQLAddr(), "CreateConnections" /* prefix */, url.User(security.RootUser))
 	defer firstServerConnCleanup()
@@ -111,6 +116,11 @@ func TestStatusAPIIndexUsage(t *testing.T) {
 	require.NoError(t, err)
 	require.False(t, rows.Next())
 
+	indexKeyPrimary := roachpb.IndexUsageKey{
+		TableID: roachpb.TableID(tableID),
+		IndexID: 1, // t@t_pkey
+	}
+
 	indexKeyA := roachpb.IndexUsageKey{
 		TableID: roachpb.TableID(tableID),
 		IndexID: 2, // t@t_a_idx
@@ -150,7 +160,7 @@ func TestStatusAPIIndexUsage(t *testing.T) {
 	_, err = secondServerSQLConn.Exec("SELECT k FROM t WHERE a = 10 AND b = 200")
 	require.NoError(t, err)
 
-	// Record a full scan over t_b_idx.
+	// Record an index join and full scan of t_b_idx.
 	_, err = secondServerSQLConn.Exec("SELECT * FROM t@t_b_idx")
 	require.NoError(t, err)
 
@@ -164,20 +174,29 @@ func TestStatusAPIIndexUsage(t *testing.T) {
 	thirdLocalStatsReader := thirdServer.SQLServer().(*sql.Server).GetLocalIndexStatistics()
 
 	// First node should have nothing.
-	stats := firstLocalStatsReader.Get(indexKeyA.TableID, indexKeyA.IndexID)
+	stats := firstLocalStatsReader.Get(indexKeyPrimary.TableID, indexKeyPrimary.IndexID)
+	require.Equal(t, roachpb.IndexUsageStatistics{}, stats, "expecting empty stats on node 1, but found %v", stats)
+
+	stats = firstLocalStatsReader.Get(indexKeyA.TableID, indexKeyA.IndexID)
 	require.Equal(t, roachpb.IndexUsageStatistics{}, stats, "expecting empty stats on node 1, but found %v", stats)
 
 	stats = firstLocalStatsReader.Get(indexKeyB.TableID, indexKeyB.IndexID)
 	require.Equal(t, roachpb.IndexUsageStatistics{}, stats, "expecting empty stats on node 1, but found %v", stats)
 
 	// Third node should have nothing.
+	stats = firstLocalStatsReader.Get(indexKeyPrimary.TableID, indexKeyPrimary.IndexID)
+	require.Equal(t, roachpb.IndexUsageStatistics{}, stats, "expecting empty stats on node 3, but found %v", stats)
+
 	stats = thirdLocalStatsReader.Get(indexKeyA.TableID, indexKeyA.IndexID)
 	require.Equal(t, roachpb.IndexUsageStatistics{}, stats, "expecting empty stats on node 3, but found %v", stats)
 
 	stats = thirdLocalStatsReader.Get(indexKeyB.TableID, indexKeyB.IndexID)
 	require.Equal(t, roachpb.IndexUsageStatistics{}, stats, "expecting empty stats on node 1, but found %v", stats)
 
 	// Second server should have nonempty local storage.
+	stats = secondLocalStatsReader.Get(indexKeyPrimary.TableID, indexKeyPrimary.IndexID)
+	compareStatsHelper(t, expectedStatsIndexPrimary, stats, time.Minute)
+
 	stats = secondLocalStatsReader.Get(indexKeyA.TableID, indexKeyA.IndexID)
 	compareStatsHelper(t, expectedStatsIndexA, stats, time.Minute)
 
@@ -197,14 +216,16 @@ func TestStatusAPIIndexUsage(t *testing.T) {
 		}
 		statsEntries++
 		switch stats.Key.IndexID {
+		case indexKeyPrimary.IndexID: // t@t_pkey
+			compareStatsHelper(t, expectedStatsIndexPrimary, stats.Stats, time.Minute)
 		case indexKeyA.IndexID: // t@t_a_idx
 			compareStatsHelper(t, expectedStatsIndexA, stats.Stats, time.Minute)
 		case indexKeyB.IndexID: // t@t_b_idx
 			compareStatsHelper(t, expectedStatsIndexB, stats.Stats, time.Minute)
 		}
 	}
 
-	require.True(t, statsEntries == 2, "expect to find two stats entries in RPC response, but found %d", statsEntries)
+	require.Equal(t, 3, statsEntries, "expect to find 3 stats entries in RPC response, but found %d", statsEntries)
 
 	// Test disabling subsystem.
 	_, err = secondServerSQLConn.Exec("SET CLUSTER SETTING sql.metrics.index_usage_stats.enabled = false")
@@ -232,7 +253,7 @@ func TestStatusAPIIndexUsage(t *testing.T) {
 			compareStatsHelper(t, expectedStatsIndexB, stats.Stats, time.Minute)
 		}
 	}
-	require.True(t, statsEntries == 2, "expect to find two stats entries in RPC response, but found %d", statsEntries)
+	require.Equal(t, 3, statsEntries, "expect to find 3 stats entries in RPC response, but found %d", statsEntries)
 }
 
 func TestGetTableID(t *testing.T) {

pkg/sql/opt_exec_factory.go

@@ -610,9 +610,18 @@ func (ef *execFactory) ConstructIndexJoin(
 		return nil, err
 	}
 
-	tableScan.index = tabDesc.GetPrimaryIndex()
+	idx := tabDesc.GetPrimaryIndex()
+	tableScan.index = idx
 	tableScan.disableBatchLimit()
 
+	if !ef.isExplain {
+		idxUsageKey := roachpb.IndexUsageKey{
+			TableID: roachpb.TableID(tabDesc.GetID()),
+			IndexID: roachpb.IndexID(idx.GetID()),
+		}
+		ef.planner.extendedEvalCtx.indexUsageStats.RecordRead(idxUsageKey)
+	}
+
 	n := &indexJoinNode{
 		input:         input.(planNode),
 		table:         tableScan,