Merge #156672

156672: sql/colfetcher: don't scan extra KVs beyond limit r=yuzefovich a=DrewKimball

#### sql/colfetcher: don't scan extra KVs beyond limit

When there are multiple column families in an index, the fetcher
implementations determine a SQL row is finished by comparing against the
next KV. As an optimization, `cfetcher` keeps track of the max column
family ID in the table. When it encounters this family, it knows the
current SQL row is finished.

Since v22.1, `cfetcher` will fetch exactly the number of KVs needed in
order to satisfy its hard limit in the first batch. The thinking was that
we can avoid having to grab an extra KV for comparison because either:
* One or more column families are elided due to NULL values, in which case
  we'll fetch more KVs than we need and have at least one for determining
  that the last row is done.
* Or there are no elided entries, in which case we will see the max family
  ID for the last row and skip fetching an additional KV for comparison.

This could result in secondary indexes fetching an additional batch with
~10k KVs, because secondary indexes may not contain the column family with
the maximum ID, necessitating a fetch for an additional KV for comparison.

This commit prevents the unnecessary fetch by plumbing the maximum family
ID for the *index* being read, not just the table.

Fixes #156613

Release note (bug fix): Fixed a bug that could cause a scan over a secondary
index to read significantly more KVs than necessary in order to satisfy a
limit when the scanned index has more than one column family.

Co-authored-by: Drew Kimball <[email protected]>

Author: craig[bot]

pkg/sql/catalog/descriptor.go

@@ -596,6 +596,10 @@ type TableDescriptor interface {
 	// one k/v pair per family in the table, just like a primary index.
 	IndexKeysPerRow(idx Index) int
 
+	// MaxFamilyIDForIndex returns the maximum column family ID used to encode a
+	// row for the given index.
+	MaxFamilyIDForIndex(idx Index) (descpb.FamilyID, error)
+
 	// IndexFetchSpecKeyAndSuffixColumns returns information about the key and
 	// suffix columns, suitable for populating a fetchpb.IndexFetchSpec.
 	IndexFetchSpecKeyAndSuffixColumns(idx Index) []fetchpb.IndexFetchSpec_KeyColumn

pkg/sql/catalog/fetchpb/index_fetch.proto

@@ -111,6 +111,12 @@ message IndexFetchSpec {
   // index ID. It is used to speed up the decoding process.
   optional uint32 key_prefix_length = 11 [(gogoproto.nullable) = false];
 
+  // Before v26.2, MaxFamilyID is the highest family ID in the table. In 26.2
+  // and later, MaxFamilyID is the highest family ID that is included in the
+  // index being scanned. It is used to quickly determine SQL row boundaries for
+  // tables with multiple families, which otherwise requires comparing against a
+  // key from the next SQL row, and can interfere with optimizations for row
+  // limits.
   optional uint32 max_family_id = 12 [(gogoproto.nullable) = false,
                                       (gogoproto.customname) = "MaxFamilyID",
                                       (gogoproto.casttype) = "github.com/cockroachdb/cockroach/pkg/sql/sem/catid.FamilyID"];

pkg/sql/catalog/tabledesc/structured.go

@@ -122,6 +122,34 @@ func (desc *wrapper) IndexKeysPerRow(idx catalog.Index) int {
 	return numUsedFamilies
 }
 
+// MaxFamilyIDForIndex implements the TableDescriptor interface.
+func (desc *wrapper) MaxFamilyIDForIndex(idx catalog.Index) (descpb.FamilyID, error) {
+	if desc.PrimaryIndex.ID == idx.GetID() || idx.GetEncodingType() == catenumpb.PrimaryIndexEncoding {
+		// The list of families is sorted by ID.
+		return desc.Families[len(desc.Families)-1].ID, nil
+	}
+	if idx.NumSecondaryStoredColumns() == 0 || len(desc.Families) == 1 {
+		// Secondary indexes always include the 0th family. If there are no stored
+		// columns or no additional families, then the max family ID is that of the
+		// 0th family.
+		return desc.Families[0].ID, nil
+	}
+	// Calculate the max column family used by the secondary index. We
+	// only need to look at the stored columns because column families are only
+	// applicable to the value part of the KV. Iterate in reverse order, since
+	// the list of families is sorted by ID, and we want the max family ID.
+	storedColumnIDs := idx.CollectSecondaryStoredColumnIDs()
+	for i := len(desc.Families) - 1; i >= 0; i-- {
+		family := desc.Families[i]
+		for _, columnID := range family.ColumnIDs {
+			if storedColumnIDs.Contains(columnID) {
+				return family.ID, nil
+			}
+		}
+	}
+	return 0, errors.AssertionFailedf("could not find max column family for index %d", idx.GetID())
+}
+
 // BuildIndexName returns an index name that is not equal to any
 // of tableDesc's indexes, roughly following Postgres's conventions for naming
 // anonymous indexes. For example:

pkg/sql/colfetcher/BUILD.bazel

@@ -85,13 +85,15 @@ go_test(
     srcs = [
         "bytes_read_test.go",
         "main_test.go",
+        "scan_limit_test.go",
         "vectorized_batch_size_test.go",
     ],
     deps = [
         "//pkg/base",
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",
         "//pkg/server",
+        "//pkg/sql/sem/eval",
         "//pkg/testutils",
         "//pkg/testutils/serverutils",
         "//pkg/testutils/skip",

pkg/sql/colfetcher/cfetcher.go

@@ -1440,8 +1440,12 @@ func (cf *cFetcher) finalizeBatch() {
 // getCurrentColumnFamilyID returns the column family id of the key in
 // cf.machine.nextKV.Key.
 func (cf *cFetcher) getCurrentColumnFamilyID() (descpb.FamilyID, error) {
-	// If the table only has 1 column family, and its ID is 0, we know that the
+	// If the index only has 1 column family, and its ID is 0, we know that the
 	// key has to be the 0th column family.
+	// TODO(drewk): once we know all nodes are at least 26.2, we can perform this
+	// optimization whenever cf.table.spec.MaxKeysPerRow == 1, since at that point
+	// MaxFamilyID is always for the just index being scanned, not the whole
+	// table.
 	if cf.table.spec.MaxFamilyID == 0 {
 		return 0, nil
 	}

pkg/sql/colfetcher/scan_limit_test.go

@@ -0,0 +1,71 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package colfetcher_test
+
+import (
+	"context"
+	"regexp"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
+	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/stretchr/testify/require"
+)
+
+// TestScanLimitKVPairsRead is a regression test for fetching an extra KV batch
+// beyond the limit when the index has more than one column family and doesn't
+// include the max column family. It verifies that the correct number of KV
+// pairs are read as reported in EXPLAIN ANALYZE output.
+func TestScanLimitKVPairsRead(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	s, conn, _ := serverutils.StartServer(t, base.TestServerArgs{
+		Knobs: base.TestingKnobs{SQLEvalContext: &eval.TestingKnobs{ForceProductionValues: true}},
+	})
+	defer s.Stopper().Stop(ctx)
+	sqlDB := sqlutils.MakeSQLRunner(conn)
+
+	checkRows := func(rows [][]string) {
+		// Look for "KV pairs read" in the output.
+		kvPairsReadRE := regexp.MustCompile(`KV pairs read: ([\d,]+)`)
+		var kvPairsCount string
+		for _, row := range rows {
+			if len(row) != 1 {
+				t.Fatalf("expected one column in EXPLAIN ANALYZE output")
+			}
+			if matches := kvPairsReadRE.FindStringSubmatch(row[0]); len(matches) > 0 {
+				kvPairsCount = matches[1]
+				break
+			}
+		}
+		require.Equal(t, "15", kvPairsCount,
+			"expected exactly 15 KV pairs to be read for LIMIT 5 with 3 column families",
+		)
+	}
+
+	sqlDB.Exec(t, `
+  CREATE TABLE t_multi_cf (
+    k INT, a INT, b INT, c INT,
+    INDEX (k) STORING (a, b),
+    FAMILY (k), FAMILY (a), FAMILY (b), FAMILY (c)
+  )`)
+
+	// Case with no NULL values.
+	sqlDB.Exec(t, `INSERT INTO t_multi_cf (SELECT t, t%19, t*7, t//3 FROM generate_series(1, 100) g(t))`)
+	rows := sqlDB.QueryStr(t, `EXPLAIN ANALYZE SELECT a, b FROM t_multi_cf LIMIT 5`)
+	checkRows(rows)
+
+	// Case with some NULL values in the stored columns.
+	sqlDB.Exec(t, `UPDATE t_multi_cf SET a = NULL WHERE k % 10 = 0`)
+	rows = sqlDB.QueryStr(t, `EXPLAIN ANALYZE SELECT a, b FROM t_multi_cf LIMIT 5`)
+	checkRows(rows)
+}

pkg/sql/rowenc/index_fetch.go

@@ -63,11 +63,10 @@ func InitIndexFetchSpec(
 
 	s.FamilyDefaultColumns = table.FamilyDefaultColumns()
 
-	families := table.GetFamilies()
-	for i := range families {
-		if id := families[i].ID; id > s.MaxFamilyID {
-			s.MaxFamilyID = id
-		}
+	var err error
+	s.MaxFamilyID, err = table.MaxFamilyIDForIndex(index)
+	if err != nil {
+		return err
 	}
 
 	s.KeyAndSuffixColumns = table.IndexFetchSpecKeyAndSuffixColumns(index)

pkg/sql/rowenc/testdata/index-fetch

@@ -407,6 +407,7 @@ CREATE TABLE fam (
 
   INDEX b(b),
   INDEX b2(b) STORING (d),
+  INDEX b3(b) STORING (c),
   INDEX c(c),
   INDEX c2(c) STORING (d),
 
@@ -491,7 +492,7 @@ columns:
   "num_key_suffix_columns": 1,
   "max_keys_per_row": 1,
   "key_prefix_length": 2,
-  "max_family_id": 2,
+  "max_family_id": 0,
   "family_default_columns": [
     {
       "family_id": 0,
@@ -601,6 +602,71 @@ columns:
   ]
 }
 
+# Index b3 spans two families, but doesn't include the max column family ID.
+index-fetch
+table: fam
+index: b3
+columns:
+  - a
+----
+{
+  "version": 1,
+  "table_id": 107,
+  "table_name": "fam",
+  "index_id": 4,
+  "index_name": "b3",
+  "is_secondary_index": true,
+  "is_unique_index": false,
+  "geo_config": {},
+  "encoding_type": 0,
+  "num_key_suffix_columns": 1,
+  "max_keys_per_row": 2,
+  "key_prefix_length": 2,
+  "max_family_id": 1,
+  "family_default_columns": [
+    {
+      "family_id": 0,
+      "default_column_id": 2
+    },
+    {
+      "family_id": 1,
+      "default_column_id": 3
+    }
+  ],
+  "key_and_suffix_columns": [
+    {
+      "column": {
+        "column_id": 2,
+        "name": "b",
+        "type": "family: StringFamily\nwidth: 0\nprecision: 0\nlocale: \"\"\nvisible_type: 0\noid: 25\ntime_precision_is_set: false\n",
+        "is_non_nullable": false
+      },
+      "direction": 0,
+      "is_composite": false,
+      "is_inverted": false
+    },
+    {
+      "column": {
+        "column_id": 1,
+        "name": "a",
+        "type": "family: IntFamily\nwidth: 64\nprecision: 0\nlocale: \"\"\nvisible_type: 3\noid: 20\ntime_precision_is_set: false\n",
+        "is_non_nullable": true
+      },
+      "direction": 0,
+      "is_composite": false,
+      "is_inverted": false
+    }
+  ],
+  "fetched_columns": [
+    {
+      "column_id": 1,
+      "name": "a",
+      "type": "family: IntFamily\nwidth: 64\nprecision: 0\nlocale: \"\"\nvisible_type: 3\noid: 20\ntime_precision_is_set: false\n",
+      "is_non_nullable": true
+    }
+  ]
+}
+
 # Index c has one key per row.
 index-fetch
 table: fam
@@ -612,7 +678,7 @@ columns:
   "version": 1,
   "table_id": 107,
   "table_name": "fam",
-  "index_id": 4,
+  "index_id": 5,
   "index_name": "c",
   "is_secondary_index": true,
   "is_unique_index": false,
@@ -621,7 +687,7 @@ columns:
   "num_key_suffix_columns": 1,
   "max_keys_per_row": 1,
   "key_prefix_length": 2,
-  "max_family_id": 2,
+  "max_family_id": 0,
   "family_default_columns": [
     {
       "family_id": 0,
@@ -677,7 +743,7 @@ columns:
   "version": 1,
   "table_id": 107,
   "table_name": "fam",
-  "index_id": 5,
+  "index_id": 6,
   "index_name": "c2",
   "is_secondary_index": true,
   "is_unique_index": false,