Merge cockroachdb#33350

33350: sql: point spans should never include interleaves r=jordanlewis a=jordanlewis

Closes cockroachdb#33347.
Closes cockroachdb#33326.
Closes cockroachdb#33330.

Previously, many sql processors that needed to generate spans for key
lookups unconditionally used the PrefixEnd method on the key to look up
to get the end of the keyspan. While correct, this generated span is
wide enough to contain all child interleaves of a table on that key, if
they exist. This is very bad for performance, causing "point scans" to
accidentally need to scan through an entire interleaves if it exists.

Now, code is adjusted to use spans that are bounded on the right by the
interleave marker appended to the start key, if the index to look up in
is unique and fully specified. Note that the condition chosen is broader
than strictly necessary - we only insert interleaves today into primary
indexes, and we could tighten the condition to only emit an interleave
marker if the index does in fact have interleaves - but keeping the
condition simple makes things easier to test and provides a modicum of
assurance that this hairy code won't have to change again if we decide
to permit interleaving tables or indexes into other unique, non primary
indexes later.

New helper methods were introduced that permit retrieval of the key span
for a set of key values, and not just the key itself, in the hopes that
future code won't commit the same sins.

```
name                                      old time/op    new time/op    delta
SQL/Cockroach/InterleavedFK/count=1-8        790µs ± 9%     515µs ± 8%  -34.77%  (p=0.000 n=10+10)
SQL/Cockroach/InterleavedFK/count=10-8      6.53ms ± 7%    1.32ms ± 5%  -79.79%  (p=0.000 n=10+10)
SQL/Cockroach/InterleavedFK/count=100-8     65.1ms ±10%     9.7ms ±14%  -85.16%  (p=0.000 n=10+10)
SQL/Cockroach/InterleavedFK/count=1000-8     565ms ± 8%     102ms ±11%  -81.99%  (p=0.000 n=10+10)

name                                      old alloc/op   new alloc/op   delta
SQL/Cockroach/InterleavedFK/count=1-8        158kB ± 2%      54kB ± 6%  -65.89%  (p=0.000 n=10+10)
SQL/Cockroach/InterleavedFK/count=10-8      2.72MB ± 1%    0.14MB ± 1%  -94.88%  (p=0.000 n=10+9)
SQL/Cockroach/InterleavedFK/count=100-8     26.8MB ± 1%     1.0MB ± 7%  -96.17%  (p=0.000 n=10+8)
SQL/Cockroach/InterleavedFK/count=1000-8     222MB ± 1%      11MB ± 6%  -94.95%  (p=0.000 n=10+9)

name                                      old allocs/op  new allocs/op  delta
SQL/Cockroach/InterleavedFK/count=1-8          515 ± 2%       508 ± 3%     ~     (p=0.195 n=10+10)
SQL/Cockroach/InterleavedFK/count=10-8       1.64k ±20%     1.33k ± 1%  -18.59%  (p=0.000 n=10+9)
SQL/Cockroach/InterleavedFK/count=100-8      14.2k ±12%     10.0k ±15%  -29.25%  (p=0.000 n=10+9)
SQL/Cockroach/InterleavedFK/count=1000-8      129k ±16%      102k ± 7%  -21.07%  (p=0.000 n=10+9)
```

Release note (performance improvement): index joins, lookup joins,
foreign key checks, cascade scans, zig zag joins, and upserts no longer
needlessly scan over child interleaved tables when searching for keys.

Co-authored-by: Jordan Lewis <[email protected]>

pkg/bench/bench_test.go

@@ -356,6 +356,7 @@ func BenchmarkSQL(b *testing.B) {
 			runBenchmarkInsertFK,
 			runBenchmarkInsertSecondaryIndex,
 			runBenchmarkInterleavedSelect,
+			runBenchmarkInterleavedFK,
 			runBenchmarkTrackChoices,
 			runBenchmarkUpdate,
 			runBenchmarkUpsert,
@@ -713,6 +714,43 @@ func runBenchmarkInterleavedSelect(b *testing.B, db *gosql.DB, count int) {
 	b.StopTimer()
 }
 
+func runBenchmarkInterleavedFK(b *testing.B, db *gosql.DB, count int) {
+	defer func() {
+		if _, err := db.Exec(`DROP TABLE IF EXISTS bench.parent CASCADE; DROP TABLE IF EXISTS bench.child CASCADE`); err != nil {
+			b.Fatal(err)
+		}
+	}()
+
+	if _, err := db.Exec(`
+CREATE TABLE bench.parent (a INT PRIMARY KEY);
+INSERT INTO bench.parent VALUES(0);
+CREATE TABLE bench.child
+  (a INT REFERENCES bench.parent(a),
+   b INT, PRIMARY KEY(a, b))
+INTERLEAVE IN PARENT bench.parent (a)
+`); err != nil {
+		b.Fatal(err)
+	}
+
+	b.ResetTimer()
+	var buf bytes.Buffer
+	val := 0
+	for i := 0; i < b.N; i++ {
+		buf.Reset()
+		buf.WriteString(`INSERT INTO bench.child VALUES `)
+		for j := 0; j < count; j++ {
+			if j > 0 {
+				buf.WriteString(", ")
+			}
+			fmt.Fprintf(&buf, "(0, %d)", val)
+			val++
+		}
+		if _, err := db.Exec(buf.String()); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
 // runBenchmarkOrderBy benchmarks scanning a table and sorting the results.
 func runBenchmarkOrderBy(b *testing.B, db *gosql.DB, count int, limit int, distinct bool) {
 	defer func() {

pkg/sql/distsqlrun/indexjoiner.go

@@ -204,13 +204,9 @@ func (ij *indexJoiner) generateSpan(row sqlbase.EncDatumRow) (roachpb.Span, erro
 	// numKeyCols.
 	keyRow := row[:numKeyCols]
 	types := ij.input.OutputTypes()[:numKeyCols]
-	key, err := sqlbase.MakeKeyFromEncDatums(
+	return sqlbase.MakeSpanFromEncDatums(
 		ij.keyPrefix, keyRow, types, ij.desc.PrimaryIndex.ColumnDirections, &ij.desc,
 		&ij.desc.PrimaryIndex, &ij.alloc)
-	if err != nil {
-		return roachpb.Span{}, err
-	}
-	return roachpb.Span{Key: key, EndKey: key.PrefixEnd()}, nil
 }
 
 // outputStatsToTrace outputs the collected indexJoiner stats to the trace. Will

pkg/sql/distsqlrun/joinreader.go

@@ -313,23 +313,23 @@ func (jr *joinReader) neededRightCols() util.FastIntSet {
 	return neededRightCols
 }
 
-// Generate a key to create a span for a given row.
+// Generate a span for a given row.
 // If lookup columns are specified will use those to collect the relevant
 // columns. Otherwise the first rows are assumed to correspond with the index.
-func (jr *joinReader) generateKey(row sqlbase.EncDatumRow) (roachpb.Key, error) {
+func (jr *joinReader) generateSpan(row sqlbase.EncDatumRow) (roachpb.Span, error) {
 	numKeyCols := len(jr.indexTypes)
 	numLookupCols := len(jr.lookupCols)
 
 	if numLookupCols > numKeyCols {
-		return nil, errors.Errorf(
+		return roachpb.Span{}, errors.Errorf(
 			"%d lookup columns specified, expecting at most %d", numLookupCols, numKeyCols)
 	}
 
 	jr.indexKeyRow = jr.indexKeyRow[:0]
 	for _, id := range jr.lookupCols {
 		jr.indexKeyRow = append(jr.indexKeyRow, row[id])
 	}
-	return sqlbase.MakeKeyFromEncDatums(
+	return sqlbase.MakeSpanFromEncDatums(
 		jr.indexKeyPrefix, jr.indexKeyRow, jr.indexTypes[:numLookupCols], jr.indexDirs, &jr.desc,
 		jr.index, &jr.alloc)
 }
@@ -414,16 +414,16 @@ func (jr *joinReader) readInput() (joinReaderState, *ProducerMetadata) {
 		if jr.hasNullLookupColumn(inputRow) {
 			continue
 		}
-		key, err := jr.generateKey(inputRow)
+		span, err := jr.generateSpan(inputRow)
 		if err != nil {
 			jr.MoveToDraining(err)
 			return jrStateUnknown, jr.DrainHelper()
 		}
-		inputRowIndices := jr.keyToInputRowIndices[string(key)]
+		inputRowIndices := jr.keyToInputRowIndices[string(span.Key)]
 		if inputRowIndices == nil {
-			spans = append(spans, roachpb.Span{Key: key, EndKey: key.PrefixEnd()})
+			spans = append(spans, span)
 		}
-		jr.keyToInputRowIndices[string(key)] = append(inputRowIndices, i)
+		jr.keyToInputRowIndices[string(span.Key)] = append(inputRowIndices, i)
 	}
 	if len(spans) == 0 {
 		// All of the input rows were filtered out. Skip the index lookup.
@@ -452,7 +452,7 @@ func (jr *joinReader) performLookup() (joinReaderState, *ProducerMetadata) {
 	for len(jr.lookupRows) < jr.batchSize {
 		// Construct a "partial key" of nCols, so we can match the key format that
 		// was stored in our keyToInputRowIndices map. This matches the format that
-		// is output in jr.generateKey.
+		// is output in jr.generateSpan.
 		key, err := jr.fetcher.PartialKey(nCols)
 		if err != nil {
 			jr.MoveToDraining(err)
@@ -603,14 +603,14 @@ func (jr *joinReader) primaryLookup(
 		for i, columnID := range jr.desc.PrimaryIndex.ColumnIDs {
 			values[i] = row[jr.colIdxMap[columnID]]
 		}
-		key, err := sqlbase.MakeKeyFromEncDatums(
+		span, err := sqlbase.MakeSpanFromEncDatums(
 			jr.primaryKeyPrefix, values, jr.primaryColumnTypes, jr.desc.PrimaryIndex.ColumnDirections,
 			&jr.desc, &jr.desc.PrimaryIndex, &jr.alloc)
 		if err != nil {
 			return nil, err
 		}
-		keyToInputRowIdx[string(key)] = rowIdx
-		spans[rowIdx] = roachpb.Span{Key: key, EndKey: key.PrefixEnd()}
+		keyToInputRowIdx[string(span.Key)] = rowIdx
+		spans[rowIdx] = span
 	}
 
 	// Perform the primary index scan.

pkg/sql/distsqlrun/zigzagjoiner.go

@@ -446,12 +446,13 @@ func (z *zigzagJoiner) setupInfo(spec *distsqlpb.ZigzagJoinerSpec, side int, col
 	}
 
 	info.prefix = sqlbase.MakeIndexKeyPrefix(info.table, info.index.ID)
-	info.key, err = z.produceKeyFromBaseRow()
+	span, err := z.produceSpanFromBaseRow()
 
 	if err != nil {
 		return err
 	}
-	info.endKey = info.key.PrefixEnd()
+	info.key = span.Key
+	info.endKey = span.EndKey
 	return nil
 }
 
@@ -535,7 +536,7 @@ func (z *zigzagJoiner) extractEqDatums(row sqlbase.EncDatumRow, side int) sqlbas
 // info. Used by produceKeyFromBaseRow.
 func (z *zigzagJoiner) produceInvertedIndexKey(
 	info *zigzagJoinerInfo, datums sqlbase.EncDatumRow,
-) (roachpb.Key, error) {
+) (roachpb.Span, error) {
 	// For inverted indexes, the JSON field (first column in the index) is
 	// encoded a little differently. We need to explicitly call
 	// EncodeInvertedIndexKeys to generate the prefix. The rest of the
@@ -548,7 +549,7 @@ func (z *zigzagJoiner) produceInvertedIndexKey(
 	for i, encDatum := range datums {
 		err := encDatum.EnsureDecoded(&info.indexTypes[i], info.alloc)
 		if err != nil {
-			return nil, err
+			return roachpb.Span{}, err
 		}
 
 		decodedDatums[i] = encDatum.Datum
@@ -569,26 +570,27 @@ func (z *zigzagJoiner) produceInvertedIndexKey(
 		info.prefix,
 	)
 	if err != nil {
-		return nil, err
+		return roachpb.Span{}, err
 	}
 	if len(keys) != 1 {
-		return nil, errors.Errorf("%d fixed values passed in for inverted index", len(keys))
+		return roachpb.Span{}, errors.Errorf("%d fixed values passed in for inverted index", len(keys))
 	}
 
 	// Append remaining (non-JSON) datums to the key.
-	key, _, err := sqlbase.EncodeColumns(
+	keyBytes, _, err := sqlbase.EncodeColumns(
 		info.index.ExtraColumnIDs[:len(datums)-1],
 		info.indexDirs[1:],
 		colMap,
 		decodedDatums,
 		keys[0],
 	)
-	return key, err
+	key := roachpb.Key(keyBytes)
+	return roachpb.Span{Key: key, EndKey: key.PrefixEnd()}, err
 }
 
 // Generates a Key, corresponding to the current `z.baseRow` in
 // the index on the current side.
-func (z *zigzagJoiner) produceKeyFromBaseRow() (roachpb.Key, error) {
+func (z *zigzagJoiner) produceSpanFromBaseRow() (roachpb.Span, error) {
 	info := z.infos[z.side]
 	neededDatums := info.fixedValues
 	if z.baseRow != nil {
@@ -602,7 +604,7 @@ func (z *zigzagJoiner) produceKeyFromBaseRow() (roachpb.Key, error) {
 		return z.produceInvertedIndexKey(info, neededDatums)
 	}
 
-	key, err := sqlbase.MakeKeyFromEncDatums(
+	return sqlbase.MakeSpanFromEncDatums(
 		info.prefix,
 		neededDatums,
 		info.indexTypes[:len(neededDatums)],
@@ -611,7 +613,6 @@ func (z *zigzagJoiner) produceKeyFromBaseRow() (roachpb.Key, error) {
 		info.index,
 		info.alloc,
 	)
-	return key, err
 }
 
 // Returns the column types of the equality columns.
@@ -741,13 +742,13 @@ func (z *zigzagJoiner) nextRow(
 
 		curInfo := z.infos[z.side]
 
-		var err error
 		// Generate a key from the last row seen from the last side. We're about to
 		// use it to jump to the next possible match on the current side.
-		curInfo.key, err = z.produceKeyFromBaseRow()
+		span, err := z.produceSpanFromBaseRow()
 		if err != nil {
 			return nil, z.producerMeta(err)
 		}
+		curInfo.key = span.Key
 
 		err = curInfo.fetcher.StartScan(
 			ctx,

pkg/sql/index_join.go

@@ -277,16 +277,12 @@ func (n *indexJoinNode) Next(params runParams) (bool, error) {
 				break
 			}
 			vals := n.index.Values()
-			primaryIndexKey, _, err := sqlbase.EncodeIndexKey(
+			primaryIndexSpan, _, err := sqlbase.EncodeIndexSpan(
 				n.table.desc.TableDesc(), n.table.index, n.run.colIDtoRowIndex, vals, n.run.primaryKeyPrefix)
 			if err != nil {
 				return false, err
 			}
-			key := roachpb.Key(primaryIndexKey)
-			n.table.spans = append(n.table.spans, roachpb.Span{
-				Key:    key,
-				EndKey: key.PrefixEnd(),
-			})
+			n.table.spans = append(n.table.spans, primaryIndexSpan)
 		}
 
 		if log.V(3) {

pkg/sql/logictest/testdata/logic_test/distsql_lookup_join

@@ -150,6 +150,35 @@ SELECT count(*) FROM data as d1 NATURAL JOIN data as d2
 ----
 10000
 
+# A lookup join that doesn't fully constrain the lookup side uses PrefixEnd, not
+# the interleave marker, to form the lookup span.
+query IIIIII
+SET TRACING = on,kv; SELECT * FROM distsql_lookup_test_1 JOIN distsql_lookup_test_2 ON f = b WHERE a > 1 AND e > 1; SET TRACING=off
+----
+2  1  1  NULL  2  1
+
+query T rowsort
+SELECT message FROM [SHOW KV TRACE FOR SESSION] WHERE
+message LIKE 'Scan%'
+----
+Scan /Table/54/{1/2-2}
+Scan /Table/55/1/{1-2}
+
+# A lookup join that fully constrains the lookup side uses the interleave marker
+# to form the lookup span.
+query IIIIII
+SET TRACING = on,kv; SELECT * FROM distsql_lookup_test_1 JOIN distsql_lookup_test_2 ON f = b AND e = a; SET TRACING=off
+----
+1  1  2  2     1  1
+2  1  1  NULL  2  1
+
+query T rowsort
+SELECT message FROM [SHOW KV TRACE FOR SESSION] WHERE
+message LIKE 'Scan%'
+----
+Scan /Table/54/{1-2}
+Scan /Table/55/1/1/1{-/#}, /Table/55/1/1/2{-/#}
+
 statement ok
 CREATE TABLE foo (a int, b int); INSERT INTO foo VALUES (0, 1), (0, 2), (1, 1)
 
@@ -271,6 +300,19 @@ SELECT t1.a, t2.d FROM multiples t1 JOIN multiples@bc t2 ON t1.a = t2.b
 8   16
 10  20
 
+# The final primary index lookup uses the interleave marker to constrain its
+# span, because it was fully specified and might contain interleaves.
+statement ok
+SET TRACING = on,kv; SELECT t1.a, t2.d FROM multiples t1 JOIN multiples@bc t2 ON t1.a = t2.b AND t1.a=2; SET TRACING=off
+
+query T rowsort
+SELECT message FROM [SHOW KV TRACE FOR SESSION] WHERE
+message LIKE 'Scan%'
+----
+Scan /Table/62/1/{2-3}
+Scan /Table/62/2/{2-3}
+Scan /Table/62/1/1/2{-/#}
+
 ############################
 #  LEFT OUTER LOOKUP JOIN  #
 ############################

pkg/sql/logictest/testdata/logic_test/fk

@@ -159,8 +159,8 @@ query T rowsort
 SELECT message FROM [SHOW KV TRACE FOR SESSION]
 WHERE message LIKE 'FKScan%'
 ----
-FKScan /Table/53/1/{2-3}
-FKScan /Table/54/1/"fake"{-/PrefixEnd}
+FKScan /Table/54/1/"fake"{-/#}
+FKScan /Table/53/1/2{-/#}
 
 statement error pgcode 23503 foreign key violation: values \['780'\] in columns \[sku\] referenced in table "orders"
 DELETE FROM products
@@ -628,6 +628,22 @@ DELETE FROM pairs WHERE id = 2
 statement error pgcode 23503 foreign key violation: values \[100 'one'\] in columns \[src dest\] referenced in table "refpairs"
 DELETE FROM pairs WHERE id = 1
 
+statement error foreign key violation: values \[100 'one'\] in columns \[src dest\] referenced in table "refpairs"
+SET tracing = on,kv; DELETE FROM pairs WHERE id = 1
+
+statement ok
+SET tracing=off
+
+# Test that fk scans on indexes that are longer than the foreign key use
+# PrefixEnd instead of interleave end.
+query T rowsort
+SELECT message FROM [SHOW KV TRACE FOR SESSION]
+WHERE message LIKE 'FKScan%'
+----
+FKScan /Table/85/3/100/"one"{-/#}
+FKScan /Table/86/3/100/"one"{-/#}
+FKScan /Table/87/2/100/"one"{-/PrefixEnd}
+
 # since PKs are handled differently than other indexes, check pk<->pk ref with no other indexes in play.
 statement ok
 CREATE TABLE foo (id INT PRIMARY KEY)
@@ -698,6 +714,22 @@ domain_modules  domain_modules_module_id_fk  FOREIGN KEY  FOREIGN KEY (module_id
 domain_modules  domain_modules_uq            UNIQUE       UNIQUE (domain_id ASC, module_id ASC)            true
 domain_modules  primary                      PRIMARY KEY  PRIMARY KEY (id ASC)                             true
 
+statement ok
+INSERT INTO modules VALUES(3)
+
+statement error foreign key violation: value \[2\] not found in domains@primary
+SET tracing = on,kv; INSERT INTO domain_modules VALUES (1, 2, 3)
+
+statement ok
+SET tracing=off
+
+query T rowsort
+SELECT message FROM [SHOW KV TRACE FOR SESSION]
+WHERE message LIKE 'FKScan%'
+----
+FKScan /Table/93/1/3{-/#}
+FKScan /Table/94/1/2{-/#}
+
 statement ok
 CREATE TABLE tx (
   id INT NOT NULL PRIMARY KEY