sql: optimize point lookups on column families

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 cockroachdb#18168

Release note: None

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/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{})
 	if err != nil {
 		return nil, err
 	}

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,8 @@ 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{})
 	if err != nil {
 		return err
 	}

pkg/sql/opt/exec/execbuilder/testdata/select

@@ -579,7 +579,6 @@ query T
 SELECT message FROM [SHOW KV TRACE FOR SESSION]
  WHERE message LIKE 'fetched:%' OR message LIKE 'output row%'
 ----
-fetched: /dt/primary/'2015-08-25 04:45:45.53453+00:00' -> NULL
 fetched: /dt/primary/'2015-08-25 04:45:45.53453+00:00'/b -> '2015-08-25'
 fetched: /dt/primary/'2015-08-25 04:45:45.53453+00:00'/c -> '2h45m2s234ms'
 output row: ['2015-08-25 04:45:45.53453+00:00' '2015-08-25' '2h45m2s234ms']

pkg/sql/opt_exec_factory.go

@@ -101,7 +101,7 @@ 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)
 	if err != nil {
 		return nil, err
 	}

pkg/sql/opt_index_selection.go

@@ -21,9 +21,11 @@ import (
 
 	"github.com/pkg/errors"
 
+	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/constraint"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec/execbuilder"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/idxconstraint"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
@@ -229,7 +231,7 @@ func (p *planner) selectIndex(
 	s.run.isSecondaryIndex = (c.index != &s.desc.PrimaryIndex)
 
 	var err error
-	s.spans, err = spansFromConstraint(s.desc, c.index, c.ic.Constraint())
+	s.spans, err = spansFromConstraint(s.desc, c.index, c.ic.Constraint(), s.valNeededForCol)
 	if err != nil {
 		return nil, errors.Wrapf(
 			err, "constraint = %s, table ID = %d, index ID = %d",
@@ -495,15 +497,18 @@ func (v *indexInfo) makeIndexConstraints(
 func unconstrainedSpans(
 	tableDesc *sqlbase.TableDescriptor, index *sqlbase.IndexDescriptor,
 ) (roachpb.Spans, error) {
-	return spansFromConstraint(tableDesc, index, nil)
+	return spansFromConstraint(tableDesc, index, nil, exec.ColumnOrdinalSet{})
 }
 
 // spansFromConstraint converts the spans in a Constraint to roachpb.Spans.
 //
 // interstices are pieces of the key that need to be inserted after each column
 // (for interleavings).
 func spansFromConstraint(
-	tableDesc *sqlbase.TableDescriptor, index *sqlbase.IndexDescriptor, c *constraint.Constraint,
+	tableDesc *sqlbase.TableDescriptor,
+	index *sqlbase.IndexDescriptor,
+	c *constraint.Constraint,
+	needed exec.ColumnOrdinalSet,
 ) (roachpb.Spans, error) {
 	interstices := make([][]byte, len(index.ColumnDirections)+len(index.ExtraColumnIDs)+1)
 	interstices[0] = sqlbase.MakeIndexKeyPrefix(tableDesc, index.ID)
@@ -529,20 +534,21 @@ func spansFromConstraint(
 
 	if c == nil || c.IsUnconstrained() {
 		// Encode a full span.
-		sp, err := spanFromConstraintSpan(tableDesc, index, &constraint.UnconstrainedSpan, interstices)
+		sp, err := spansFromConstraintSpan(
+			tableDesc, index, &constraint.UnconstrainedSpan, interstices, needed)
 		if err != nil {
 			return nil, err
 		}
-		return roachpb.Spans{sp}, nil
+		return sp, nil
 	}
 
-	spans := make(roachpb.Spans, c.Spans.Count())
-	for i := range spans {
-		s, err := spanFromConstraintSpan(tableDesc, index, c.Spans.Get(i), interstices)
+	spans := make(roachpb.Spans, 0, c.Spans.Count())
+	for i := 0; i < c.Spans.Count(); i++ {
+		s, err := spansFromConstraintSpan(tableDesc, index, c.Spans.Get(i), interstices, needed)
 		if err != nil {
 			return nil, err
 		}
-		spans[i] = s
+		spans = append(spans, s...)
 	}
 	return spans, nil
 }
@@ -590,19 +596,22 @@ func encodeConstraintKey(
 	return key, nil
 }
 
-// spanFromConstraintSpan converts a constraint.Span to a roachpb.Span.
-func spanFromConstraintSpan(
+// spansFromConstraintSpan converts a constraint.Span to one or more
+// roachpb.Spans. It returns multiple spans in the case that multiple,
+// non-adjacent column families should be scanned.
+func spansFromConstraintSpan(
 	tableDesc *sqlbase.TableDescriptor,
 	index *sqlbase.IndexDescriptor,
 	cs *constraint.Span,
 	interstices [][]byte,
-) (roachpb.Span, error) {
+	needed exec.ColumnOrdinalSet,
+) (roachpb.Spans, error) {
 	var s roachpb.Span
 	var err error
 	// Encode each logical part of the start key.
 	s.Key, err = encodeConstraintKey(index, cs.StartKey(), interstices)
 	if err != nil {
-		return roachpb.Span{}, err
+		return nil, err
 	}
 	if cs.StartBoundary() == constraint.IncludeBoundary {
 		s.Key = append(s.Key, interstices[cs.StartKey().Length()]...)
@@ -613,18 +622,69 @@ func spanFromConstraintSpan(
 	// Encode each logical part of the end key.
 	s.EndKey, err = encodeConstraintKey(index, cs.EndKey(), interstices)
 	if err != nil {
-		return roachpb.Span{}, err
+		return nil, err
 	}
 	s.EndKey = append(s.EndKey, interstices[cs.EndKey().Length()]...)
 
+	// Optimization: for single row lookups on a table with multiple column
+	// families, only scan the relevant column families.
+	if needed.Len() > 0 &&
+		index.ID == tableDesc.PrimaryIndex.ID &&
+		len(tableDesc.Families) > 1 &&
+		cs.StartKey().Length() == len(tableDesc.PrimaryIndex.ColumnIDs) &&
+		s.Key.Equal(s.EndKey) {
+		neededFamilyIDs := neededColumnFamilyIDs(tableDesc, needed)
+		if len(neededFamilyIDs) < len(tableDesc.Families) {
+			spans := make(roachpb.Spans, 0, len(neededFamilyIDs))
+			for i, familyID := range neededFamilyIDs {
+				var span roachpb.Span
+				span.Key = make(roachpb.Key, len(s.Key))
+				copy(span.Key, s.Key)
+				span.Key = keys.MakeFamilyKey(span.Key, uint32(familyID))
+				span.EndKey = span.Key.PrefixEnd()
+				if i > 0 && familyID == neededFamilyIDs[i-1]+1 {
+					// This column family is adjacent to the previous one. We can merge
+					// the two spans into one.
+					spans[len(spans)-1].EndKey = span.EndKey
+				} else {
+					spans = append(spans, span)
+				}
+			}
+			return spans, nil
+		}
+	}
+
 	// We tighten the end key to prevent reading interleaved children after the
 	// last parent key. If cs.End.Inclusive is true, we also advance the key as
 	// necessary.
 	endInclusive := cs.EndBoundary() == constraint.IncludeBoundary
 	s.EndKey, err = sqlbase.AdjustEndKeyForInterleave(tableDesc, index, s.EndKey, endInclusive)
 	if err != nil {
-		return roachpb.Span{}, err
+		return nil, err
 	}
+	return roachpb.Spans{s}, nil
+}
+
+func neededColumnFamilyIDs(
+	tableDesc *sqlbase.TableDescriptor, neededCols exec.ColumnOrdinalSet,
+) []sqlbase.FamilyID {
+	colIdxMap := tableDesc.ColumnIdxMap()
+
+	var needed []sqlbase.FamilyID
+	for _, family := range tableDesc.Families {
+		for _, columnID := range family.ColumnIDs {
+			columnOrdinal := colIdxMap[columnID]
+			if neededCols.Contains(columnOrdinal) {
+				needed = append(needed, family.ID)
+				break
+			}
+		}
+	}
+
+	// TODO(solon): There is a further optimization possible here: if there is at
+	// least one non-nullable column in the needed column families, we can
+	// potentially omit the primary family, since the primary keys are encoded
+	// in all families.
 
-	return s, nil
+	return needed
 }

pkg/sql/opt_index_selection_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/constraint"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/optbuilder"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/xform"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
@@ -110,7 +111,7 @@ func makeSpans(
 		t.Fatal(err)
 	}
 
-	spans, err = spansFromConstraint(desc, index, c.ic.Constraint())
+	spans, err = spansFromConstraint(desc, index, c.ic.Constraint(), exec.ColumnOrdinalSet{})
 	if err != nil {
 		t.Fatal(err)
 	}