kvstreamer: optimize singleRangeBatch.Less

pkg/bench/bench_test.go

@@ -1200,6 +1200,75 @@ func BenchmarkIndexJoinColumnFamilies(b *testing.B) {
 	})
 }
 
+// BenchmarkLookupJoinEqColsAreKeyNoOrdering measures a lookup-join with 1000
+// rows when equality columns are key and ordering doesn't have to be
+// maintained.
+func BenchmarkLookupJoinEqColsAreKeyNoOrdering(b *testing.B) {
+	defer log.Scope(b).Close(b)
+	ForEachDB(b, func(b *testing.B, db *sqlutils.SQLRunner) {
+		db.Exec(b, `CREATE TABLE t1 (a INT)`)
+		db.Exec(b, `INSERT INTO t1 SELECT generate_series(1, 1000)`)
+		db.Exec(b, `CREATE TABLE t2 (a INT PRIMARY KEY, b INT)`)
+		db.Exec(b, `INSERT INTO t2 SELECT generate_series(1, 1000), (random()*1000)::INT`)
+		b.ResetTimer()
+
+		for i := 0; i < b.N; i++ {
+			db.Exec(b, `SELECT * FROM t1 INNER LOOKUP JOIN t2 ON t1.a = t2.a`)
+		}
+	})
+}
+
+// BenchmarkLookupJoinEqColsAreKeyOrdering measures a lookup-join with 1000 rows
+// when equality columns are key and ordering needs to be maintained.
+func BenchmarkLookupJoinEqColsAreKeyOrdering(b *testing.B) {
+	defer log.Scope(b).Close(b)
+	ForEachDB(b, func(b *testing.B, db *sqlutils.SQLRunner) {
+		db.Exec(b, `CREATE TABLE t1 (a INT PRIMARY KEY)`)
+		db.Exec(b, `INSERT INTO t1 SELECT generate_series(1, 1000)`)
+		db.Exec(b, `CREATE TABLE t2 (a INT PRIMARY KEY, b INT)`)
+		db.Exec(b, `INSERT INTO t2 SELECT generate_series(1, 1000), (random()*1000)::INT`)
+		b.ResetTimer()
+
+		for i := 0; i < b.N; i++ {
+			db.Exec(b, `SELECT * FROM t1 INNER LOOKUP JOIN t2 ON t1.a = t2.a ORDER BY t1.a`)
+		}
+	})
+}
+
+// BenchmarkLookupJoinNoOrdering measures a lookup-join with 1000 rows and
+// ordering doesn't have to be maintained.
+func BenchmarkLookupJoinNoOrdering(b *testing.B) {
+	defer log.Scope(b).Close(b)
+	ForEachDB(b, func(b *testing.B, db *sqlutils.SQLRunner) {
+		db.Exec(b, `CREATE TABLE t1 (a INT)`)
+		db.Exec(b, `INSERT INTO t1 SELECT generate_series(1, 1000)`)
+		db.Exec(b, `CREATE TABLE t2 (a INT, INDEX (a))`)
+		db.Exec(b, `INSERT INTO t2 SELECT generate_series(1, 1000)`)
+		b.ResetTimer()
+
+		for i := 0; i < b.N; i++ {
+			db.Exec(b, `SELECT * FROM t1 INNER LOOKUP JOIN t2 ON t1.a = t2.a`)
+		}
+	})
+}
+
+// BenchmarkLookupJoinOrdering measures a lookup-join with 1000 rows when
+// ordering needs to be maintained.
+func BenchmarkLookupJoinOrdering(b *testing.B) {
+	defer log.Scope(b).Close(b)
+	ForEachDB(b, func(b *testing.B, db *sqlutils.SQLRunner) {
+		db.Exec(b, `CREATE TABLE t1 (a INT PRIMARY KEY)`)
+		db.Exec(b, `INSERT INTO t1 SELECT generate_series(1, 1000)`)
+		db.Exec(b, `CREATE TABLE t2 (a INT, INDEX (a))`)
+		db.Exec(b, `INSERT INTO t2 SELECT generate_series(1, 1000)`)
+		b.ResetTimer()
+
+		for i := 0; i < b.N; i++ {
+			db.Exec(b, `SELECT * FROM t1 INNER LOOKUP JOIN t2 ON t1.a = t2.a ORDER BY t1.a`)
+		}
+	})
+}
+
 func BenchmarkSortJoinAggregation(b *testing.B) {
 	defer log.Scope(b).Close(b)
 	ForEachDB(b, func(b *testing.B, db *sqlutils.SQLRunner) {

pkg/kv/kvclient/kvstreamer/requests_provider.go

@@ -29,6 +29,12 @@ import (
 // singleRangeBatch.
 type singleRangeBatch struct {
 	reqs []roachpb.RequestUnion
+	// reqsKeys stores the start key of the corresponding request in reqs. It is
+	// only set prior to sorting reqs within this object (currently, this is
+	// done only in the OutOfOrder mode for the original requests - resume
+	// requests don't set this), and the reference is niled out right after
+	// sorting is complete. Thus, this slice doesn't have to be accounted for.
+	reqsKeys []roachpb.Key
 	// positions is a 1-to-1 mapping with reqs to indicate which ordinal among
 	// the originally enqueued requests a particular reqs[i] corresponds to. In
 	// other words, if reqs[i] is (or a part of) enqueuedReqs[j], then
@@ -90,6 +96,9 @@ func (r *singleRangeBatch) Len() int {
 
 func (r *singleRangeBatch) Swap(i, j int) {
 	r.reqs[i], r.reqs[j] = r.reqs[j], r.reqs[i]
+	if r.reqsKeys != nil {
+		r.reqsKeys[i], r.reqsKeys[j] = r.reqsKeys[j], r.reqsKeys[i]
+	}
 	r.positions[i], r.positions[j] = r.positions[j], r.positions[i]
 	if r.subRequestIdx != nil {
 		r.subRequestIdx[i], r.subRequestIdx[j] = r.subRequestIdx[j], r.subRequestIdx[i]
@@ -98,9 +107,7 @@ func (r *singleRangeBatch) Swap(i, j int) {
 
 // Less returns true if r.reqs[i]'s key comes before r.reqs[j]'s key.
 func (r *singleRangeBatch) Less(i, j int) bool {
-	// TODO(yuzefovich): figure out whether it's worth extracting the keys when
-	// constructing singleRangeBatch object.
-	return r.reqs[i].GetInner().Header().Key.Compare(r.reqs[j].GetInner().Header().Key) < 0
+	return r.reqsKeys[i].Compare(r.reqsKeys[j]) < 0
 }
 
 // priority returns the priority value of this batch.

pkg/kv/kvclient/kvstreamer/streamer.go

@@ -450,6 +450,7 @@ func (s *Streamer) Enqueue(ctx context.Context, reqs []roachpb.RequestUnion) (re
 			s.mu.Unlock()
 		}
 	}()
+	var reqsKeysScratch []roachpb.Key
 	for ; ri.Valid(); ri.Seek(ctx, seekKey, scanDir) {
 		// Truncate the request span to the current range.
 		singleRangeSpan, err := rs.Intersect(ri.Token().Desc())
@@ -516,7 +517,13 @@ func (s *Streamer) Enqueue(ctx context.Context, reqs []roachpb.RequestUnion) (re
 			// can become head-of-the-line in the future. We probably will need
 			// to introduce a way to "restore" the original order within
 			// singleRangeBatch if it is sorted and issued with headOfLine=true.
+			r.reqsKeys = reqsKeysScratch[:0]
+			for i := range r.reqs {
+				r.reqsKeys = append(r.reqsKeys, r.reqs[i].GetInner().Header().Key)
+			}
 			sort.Sort(&r)
+			reqsKeysScratch = r.reqsKeys
+			r.reqsKeys = nil
 		}
 
 		requestsToServe = append(requestsToServe, r)