@@ -1226,7 +1226,7 @@ func TestTopLevelQueryStats(t *testing.T) {
12261226 var testQuery atomic.Value
12271227 // The callback will send number of rows read and rows written (for each
12281228 // ProducerMetadata.Metrics object) on these channels, respectively.
1229- rowsReadCh , rowsWrittenCh := make (chan int64 ), make (chan int64 )
1229+ rowsReadCh , rowsWrittenCh , indexRowsWrittenCh := make ( chan int64 ), make (chan int64 ), make (chan int64 )
12301230 srv , sqlDB , _ := serverutils .StartServer (t , base.TestServerArgs {
12311231 Knobs : base.TestingKnobs {
12321232 SQLExecutor : & ExecutorTestingKnobs {
@@ -1238,6 +1238,7 @@ func TestTopLevelQueryStats(t *testing.T) {
12381238 if meta != nil && meta .Metrics != nil {
12391239 rowsReadCh <- meta .Metrics .RowsRead
12401240 rowsWrittenCh <- meta .Metrics .RowsWritten
1241+ indexRowsWrittenCh <- meta .Metrics .IndexRowsWritten
12411242 }
12421243 return row , batch , meta
12431244 }
@@ -1254,23 +1255,25 @@ CREATE TABLE t (k INT PRIMARY KEY, i INT, v INT, INDEX(i));
12541255INSERT INTO t SELECT i, 1, 1 FROM generate_series(1, 10) AS g(i);
12551256CREATE FUNCTION no_reads() RETURNS INT AS 'SELECT 1' LANGUAGE SQL;
12561257CREATE FUNCTION reads() RETURNS INT AS 'SELECT count(*) FROM t' LANGUAGE SQL;
1257- CREATE FUNCTION write(x INT) RETURNS INT AS 'INSERT INTO t VALUES (x); SELECT x' LANGUAGE SQL;
1258+ CREATE FUNCTION write(x INT) RETURNS INT AS 'INSERT INTO t VALUES (x, x ); SELECT x' LANGUAGE SQL;
12581259` ); err != nil {
12591260 t .Fatal (err )
12601261 }
12611262
12621263 for _ , tc := range []struct {
1263- name string
1264- query string
1265- setup , cleanup string // optional
1266- expRowsRead int64
1267- expRowsWritten int64
1264+ name string
1265+ query string
1266+ setup , cleanup string // optional
1267+ expRowsRead int64
1268+ expRowsWritten int64
1269+ expIndexRowsWritten int64
12681270 }{
12691271 {
1270- name : "simple read" ,
1271- query : "SELECT k FROM t" ,
1272- expRowsRead : 10 ,
1273- expRowsWritten : 0 ,
1272+ name : "simple read" ,
1273+ query : "SELECT k FROM t" ,
1274+ expRowsRead : 10 ,
1275+ expRowsWritten : 0 ,
1276+ expIndexRowsWritten : 0 ,
12741277 },
12751278 {
12761279 name : "routine and index join (used to be powered by streamer)" ,
@@ -1279,47 +1282,54 @@ CREATE FUNCTION write(x INT) RETURNS INT AS 'INSERT INTO t VALUES (x); SELECT x'
12791282 cleanup : "RESET distsql" ,
12801283 // 10 rows for secondary index, 10 for index join into primary, and
12811284 // then for each row do ten-row-scan in the routine.
1282- expRowsRead : 120 ,
1283- expRowsWritten : 0 ,
1285+ expRowsRead : 120 ,
1286+ expRowsWritten : 0 ,
1287+ expIndexRowsWritten : 0 ,
12841288 },
12851289 {
1286- name : "simple write" ,
1287- query : "INSERT INTO t SELECT generate_series(11, 42)" ,
1288- expRowsRead : 0 ,
1289- expRowsWritten : 32 ,
1290+ name : "simple write" ,
1291+ query : "INSERT INTO t SELECT generate_series(11, 42)" ,
1292+ expRowsRead : 0 ,
1293+ expRowsWritten : 32 ,
1294+ expIndexRowsWritten : 64 ,
12901295 },
12911296 {
12921297 name : "read with apply join" ,
12931298 query : `SELECT (
12941299 WITH foo AS MATERIALIZED (SELECT k FROM t AS x WHERE x.k = y.k)
12951300 SELECT * FROM foo
12961301 ) FROM t AS y` ,
1297- expRowsRead : 84 , // scanning the table twice
1298- expRowsWritten : 0 ,
1302+ expRowsRead : 84 , // scanning the table twice
1303+ expRowsWritten : 0 ,
1304+ expIndexRowsWritten : 0 ,
12991305 },
13001306 {
1301- name : "routine, no reads" ,
1302- query : "SELECT no_reads()" ,
1303- expRowsRead : 0 ,
1304- expRowsWritten : 0 ,
1307+ name : "routine, no reads" ,
1308+ query : "SELECT no_reads()" ,
1309+ expRowsRead : 0 ,
1310+ expRowsWritten : 0 ,
1311+ expIndexRowsWritten : 0 ,
13051312 },
13061313 {
1307- name : "routine, reads" ,
1308- query : "SELECT reads()" ,
1309- expRowsRead : 42 ,
1310- expRowsWritten : 0 ,
1314+ name : "routine, reads" ,
1315+ query : "SELECT reads()" ,
1316+ expRowsRead : 42 ,
1317+ expRowsWritten : 0 ,
1318+ expIndexRowsWritten : 0 ,
13111319 },
13121320 {
1313- name : "routine, write" ,
1314- query : "SELECT write(43)" ,
1315- expRowsRead : 0 ,
1316- expRowsWritten : 1 ,
1321+ name : "routine, write" ,
1322+ query : "SELECT write(43)" ,
1323+ expRowsRead : 0 ,
1324+ expRowsWritten : 1 ,
1325+ expIndexRowsWritten : 2 ,
13171326 },
13181327 {
1319- name : "routine, multiple reads and writes" ,
1320- query : "SELECT reads(), write(44), reads(), write(45), write(46), reads()" ,
1321- expRowsRead : 133 , // first read is 43 rows, second is 44, third is 46
1322- expRowsWritten : 3 ,
1328+ name : "routine, multiple reads and writes" ,
1329+ query : "SELECT reads(), write(44), reads(), write(45), write(46), reads()" ,
1330+ expRowsRead : 133 , // first read is 43 rows, second is 44, third is 46
1331+ expRowsWritten : 3 ,
1332+ expIndexRowsWritten : 6 ,
13231333 },
13241334 } {
13251335 t .Run (tc .name , func (t * testing.T ) {
@@ -1344,21 +1354,24 @@ CREATE FUNCTION write(x INT) RETURNS INT AS 'INSERT INTO t VALUES (x); SELECT x'
13441354 }()
13451355 // In the main goroutine, loop until the query is completed while
13461356 // accumulating the top-level query stats.
1347- var rowsRead , rowsWritten int64
1357+ var rowsRead , rowsWritten , indexRowsWritten int64
13481358 LOOP:
13491359 for {
13501360 select {
13511361 case read := <- rowsReadCh :
13521362 rowsRead += read
13531363 case written := <- rowsWrittenCh :
13541364 rowsWritten += written
1365+ case written := <- indexRowsWrittenCh :
1366+ indexRowsWritten += written
13551367 case err := <- errCh :
13561368 require .NoError (t , err )
13571369 break LOOP
13581370 }
13591371 }
13601372 require .Equal (t , tc .expRowsRead , rowsRead )
13611373 require .Equal (t , tc .expRowsWritten , rowsWritten )
1374+ require .Equal (t , tc .expIndexRowsWritten , indexRowsWritten )
13621375 })
13631376 }
13641377}
0 commit comments