Merge #30744

30744: sql: optimize point lookups on column families r=solongordon a=solongordon

For tables with multiple column families, point lookups will now only
scan column families which contain the needed columns. Previously we
would scan the entire row. This optimization allows for faster lookups
and, perhaps more importantly, reduces contention between operations on
the same row but disjoint column families.

Fixes #18168

Release note: None

Co-authored-by: Solon Gordon <[email protected]>

pkg/sql/delete.go

@@ -447,8 +447,9 @@ func canDeleteFast(ctx context.Context, source planNode, r *deleteRun) (*scanNod
 		return nil, false
 	}
 
-	// Check whether the source plan is "simple": that it contains no
-	// remaining filtering, limiting, sorting, etc.
+	// Check whether the source plan is "simple": that it contains no remaining
+	// filtering, limiting, sorting, etc. Note that this logic must be kept in
+	// sync with the logic for setting scanNode.isDeleteSource (see doExpandPlan.)
 	// TODO(dt): We could probably be smarter when presented with an
 	// index-join, but this goes away anyway once we push-down more of
 	// SQL.

pkg/sql/distsql_plan_stats.go

@@ -40,7 +40,7 @@ func (dsp *DistSQLPlanner) createStatsPlan(
 	if err != nil {
 		return PhysicalPlan{}, err
 	}
-	scan.spans, err = unconstrainedSpans(desc, scan.index)
+	scan.spans, err = unconstrainedSpans(desc, scan.index, scan.isDeleteSource)
 	if err != nil {
 		return PhysicalPlan{}, err
 	}

pkg/sql/distsqlplan/fake_span_resolver.go

@@ -91,7 +91,7 @@ func (fit *fakeSpanResolverIterator) Seek(
 			fit.err = err
 			return
 		}
-		if !splitKey.Equal(lastKey) {
+		if !splitKey.Equal(lastKey) && span.ContainsKey(splitKey) {
 			splitKeys = append(splitKeys, splitKey)
 			lastKey = splitKey
 		}

pkg/sql/expand_plan.go

@@ -121,6 +121,18 @@ func doExpandPlan(
 		n.source, err = doExpandPlan(ctx, p, noParams, n.source)
 
 	case *deleteNode:
+		// If the source of the delete is a scan node (optionally with a render on
+		// top), mark it as such. Note that this parallels the logic in
+		// canDeleteFast.
+		maybeScan := n.source
+		if sel, ok := maybeScan.(*renderNode); ok {
+			maybeScan = sel.source.plan
+		}
+		scan, ok := maybeScan.(*scanNode)
+		if ok {
+			scan.isDeleteSource = true
+		}
+
 		n.source, err = doExpandPlan(ctx, p, noParams, n.source)
 
 	case *rowCountNode:

pkg/sql/join_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/internal/client"
 	"github.com/cockroachdb/cockroach/pkg/security"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
@@ -41,7 +42,8 @@ func newTestScanNode(kvDB *client.DB, tableName string) (*scanNode, error) {
 		return nil, err
 	}
 	scan.initOrdering(0 /* exactPrefix */, p.EvalContext())
-	scan.spans, err = spansFromConstraint(desc, &desc.PrimaryIndex, nil /* constraint */)
+	scan.spans, err = spansFromConstraint(
+		desc, &desc.PrimaryIndex, nil /* constraint */, exec.ColumnOrdinalSet{}, false /* forDelete */)
 	if err != nil {
 		return nil, err
 	}

pkg/sql/logictest/testdata/logic_test/family

@@ -37,6 +37,13 @@ SELECT * FROM abcd
 1 2 3 4
 5 6 7 8
 
+# Test point lookup, which triggers an optimization for only scanning one
+# column family.
+query I
+SELECT c FROM abcd WHERE a = 1
+----
+3
+
 query I
 SELECT count(*) FROM abcd
 ----
@@ -81,6 +88,30 @@ SELECT * FROM abcd WHERE a = 1
 ----
 1 NULL NULL NULL
 
+# Test updating a NULL family
+statement ok
+INSERT INTO abcd (a) VALUES (2)
+
+query IIII
+SELECT * FROM abcd WHERE a = 2
+----
+2 NULL NULL NULL
+
+statement ok
+UPDATE abcd SET d = 5 WHERE a = 2
+
+query IIII
+SELECT * FROM abcd WHERE a = 2
+----
+2 NULL NULL 5
+
+statement ok
+DELETE FROM abcd WHERE a = 2
+
+query IIII
+SELECT * FROM abcd WHERE a = 2
+----
+
 statement ok
 ALTER TABLE abcd ADD e STRING FAMILY f1
 

pkg/sql/lookup_join.go

@@ -17,6 +17,7 @@ package sql
 import (
 	"context"
 
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
 )
@@ -63,7 +64,9 @@ type lookupJoinRun struct {
 func (lj *lookupJoinNode) startExec(params runParams) error {
 	// Make sure the table node has a span (full scan).
 	var err error
-	lj.table.spans, err = spansFromConstraint(lj.table.desc, lj.table.index, nil /* constraint */)
+	lj.table.spans, err = spansFromConstraint(
+		lj.table.desc, lj.table.index, nil /* constraint */, exec.ColumnOrdinalSet{},
+		false /* forDelete */)
 	if err != nil {
 		return err
 	}

pkg/sql/opt/exec/execbuilder/testdata/select_index

@@ -1336,3 +1336,119 @@ render           ·         ·           (w)     ·
       │          spans     /1-/10      ·       ·
       └── scan   ·         ·           (v, w)  ·
 ·                table     t3@primary  ·       ·
+
+# ------------------------------------------------------------------------------
+# These tests are for the point lookup optimization: for single row lookups on
+# a table with multiple column families, we only scan the relevant column
+# families. Note that this applies to SELECTs and UPDATEs but not DELETEs, since
+# we need to ensure that we delete across all column families.
+# ------------------------------------------------------------------------------
+statement ok
+CREATE TABLE t4 (
+  a INT,
+  b INT,
+  c INT,
+  d INT,
+  e INT,
+  PRIMARY KEY (a, b),
+  FAMILY (a, b),
+  FAMILY (c),
+  FAMILY (d),
+  FAMILY (e)
+)
+
+statement ok
+INSERT INTO t4 VALUES (10, 20, 30, 40, 50)
+
+# Point lookup on c does not touch the d or e families.
+query TTT
+EXPLAIN SELECT c FROM t4 WHERE a = 10 and b = 20
+----
+render     ·      ·
+ └── scan  ·      ·
+·          table  t4@primary
+·          spans  /10/20/0-/10/20/1/2
+
+statement ok
+SET tracing = on,kv,results; SELECT c FROM t4 WHERE a = 10 and b = 20; SET tracing = off
+
+query T
+SELECT message FROM [SHOW KV TRACE FOR SESSION]
+ WHERE message LIKE 'fetched:%' OR message LIKE 'output row%'
+----
+fetched: /t4/primary/10/20 -> NULL
+fetched: /t4/primary/10/20/c -> 30
+output row: [30]
+
+# Point lookup on d does not touch the c or e families.
+query TTT
+EXPLAIN SELECT d FROM t4 WHERE a = 10 and b = 20
+----
+render     ·      ·
+ └── scan  ·      ·
+·          table  t4@primary
+·          spans  /10/20/0-/10/20/1 /10/20/2/1-/10/20/2/2
+
+statement ok
+SET tracing = on,kv,results; SELECT d FROM t4 WHERE a = 10 and b = 20; SET tracing = off
+
+query T
+SELECT message FROM [SHOW KV TRACE FOR SESSION]
+ WHERE message LIKE 'fetched:%' OR message LIKE 'output row%'
+----
+fetched: /t4/primary/10/20 -> NULL
+fetched: /t4/primary/10/20/c -> 40
+output row: [40]
+
+# Point lookup on both d and e uses a single span for the two adjacent column
+# families.
+query TTT
+EXPLAIN SELECT d, e FROM t4 WHERE a = 10 and b = 20
+----
+render     ·      ·
+ └── scan  ·      ·
+·          table  t4@primary
+·          spans  /10/20/0-/10/20/1 /10/20/2/1-/10/20/3/2
+
+# Optimization should also be applied for updates.
+query TTT
+EXPLAIN UPDATE t4 SET c = 30 WHERE a = 10 and b = 20
+----
+count                ·      ·
+ └── update          ·      ·
+      │              table  t4
+      │              set    c
+      └── render     ·      ·
+           └── scan  ·      ·
+·                    table  t4@primary
+·                    spans  /10/20/0-/10/20/1/2
+
+# Optimization should not be applied for deletes.
+query TTT
+EXPLAIN DELETE FROM t4 WHERE a = 10 and b = 20
+----
+count                ·      ·
+ └── delete          ·      ·
+      │              from   t4
+      └── render     ·      ·
+           └── scan  ·      ·
+·                    table  t4@primary
+·                    spans  /10/20-/10/20/#
+
+# Optimization should not be applied for non point lookups.
+query TTT
+EXPLAIN SELECT c FROM t4 WHERE a = 10 and b >= 20 and b < 22
+----
+render     ·      ·
+ └── scan  ·      ·
+·          table  t4@primary
+·          spans  /10/20-/10/21/#
+
+# Optimization should not be applied for partial primary key filter.
+query TTT
+EXPLAIN SELECT c FROM t4 WHERE a = 10
+----
+render     ·      ·
+ └── scan  ·      ·
+·          table  t4@primary
+·          spans  /10-/11

pkg/sql/opt_exec_factory.go

@@ -101,7 +101,8 @@ func (ef *execFactory) ConstructScan(
 	scan.hardLimit = hardLimit
 	scan.reverse = reverse
 	var err error
-	scan.spans, err = spansFromConstraint(tabDesc, indexDesc, indexConstraint)
+	scan.spans, err = spansFromConstraint(
+		tabDesc, indexDesc, indexConstraint, cols, scan.isDeleteSource)
 	if err != nil {
 		return nil, err
 	}