sql: optimize point lookups on column families

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
 	}