From 0d4e0d9a7e90de963961ac0d16aa912c91bfd502 Mon Sep 17 00:00:00 2001
From: Marcus Gartner <marcus@cockroachlabs.com>
Date: Mon, 20 Dec 2021 16:33:18 -0500
Subject: [PATCH] sql: do not fetch virtual columns during backfill

Fixes #73372

Release note (bug fix): A bug has been fixed that caused internal errors
when altering the primary key of a table. The bug was only present if
the table had a partial index with a predicate that referenced a virtual
computed column. This bug was present since virtual computed columns
were added in version 21.1.0.
---
 pkg/sql/backfill/backfill.go                  | 45 ++++++++++++-------
 .../testdata/logic_test/virtual_columns       | 23 ++++++++++
 2 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/pkg/sql/backfill/backfill.go b/pkg/sql/backfill/backfill.go
index eab39f7563bf..aaf990649c58 100644
--- a/pkg/sql/backfill/backfill.go
+++ b/pkg/sql/backfill/backfill.go
@@ -485,8 +485,10 @@ func (ib *IndexBackfiller) InitForLocalUse(
 }
 
 // constructExprs is a helper to construct the index and column expressions
-// required for an index backfill. It also returns the set of columns referenced
-// by any of these exprs.
+// required for an index backfill. It also returns the set of non-virtual
+// columns referenced by any of these exprs that should be fetched from the
+// primary index. Virtual columns are not included because they don't exist in
+// the primary index.
 //
 // The cols argument is the full set of cols in the table (including those being
 // added). The addedCols argument is the set of non-public, non-computed
@@ -554,21 +556,34 @@ func constructExprs(
 		colExprs[id] = computedExprs[i]
 	}
 
-	// Ensure that only existing columns are added to the needed set. Otherwise
-	// the fetcher may complain that the columns don't exist. There's a somewhat
-	// subtle invariant that if any dependencies exist between computed columns
-	// and default values that the computed column be a later column and thus the
-	// default value will have been populated. Computed columns are not permitted
-	// to reference each other.
-	addToReferencedColumns := func(cols catalog.TableColSet) {
-		cols.ForEach(func(col descpb.ColumnID) {
-			if !addedColSet.Contains(col) {
-				referencedColumns.Add(col)
+	// Ensure that only existing, non-virtual columns are added to the needed
+	// set. Otherwise the fetcher may complain that the columns don't exist.
+	// There's a somewhat subtle invariant that if any dependencies exist
+	// between computed columns and default values that the computed column be a
+	// later column and thus the default value will have been populated.
+	// Computed columns are not permitted to reference each other.
+	addToReferencedColumns := func(cols catalog.TableColSet) error {
+		for colID, ok := cols.Next(0); ok; colID, ok = cols.Next(colID + 1) {
+			if addedColSet.Contains(colID) {
+				continue
+			}
+			col, err := desc.FindColumnWithID(colID)
+			if err != nil {
+				return errors.AssertionFailedf("column %d does not exist", colID)
 			}
-		})
+			if col.IsVirtual() {
+				continue
+			}
+			referencedColumns.Add(colID)
+		}
+		return nil
+	}
+	if err := addToReferencedColumns(predicateRefColIDs); err != nil {
+		return nil, nil, catalog.TableColSet{}, err
+	}
+	if err := addToReferencedColumns(computedExprRefColIDs); err != nil {
+		return nil, nil, catalog.TableColSet{}, err
 	}
-	addToReferencedColumns(predicateRefColIDs)
-	addToReferencedColumns(computedExprRefColIDs)
 	return predicates, colExprs, referencedColumns, nil
 }
 
diff --git a/pkg/sql/logictest/testdata/logic_test/virtual_columns b/pkg/sql/logictest/testdata/logic_test/virtual_columns
index f2158d7a13ba..6ba4cdef180b 100644
--- a/pkg/sql/logictest/testdata/logic_test/virtual_columns
+++ b/pkg/sql/logictest/testdata/logic_test/virtual_columns
@@ -1226,3 +1226,26 @@ INSERT INTO t67528 (s) VALUES ('')
 
 statement ok
 CREATE INDEX ON t67528 (v DESC)
+
+# Regression test for #73372. Test backfills with partial indexes that reference
+# non-null virtual columns.
+subtest 73372
+
+statement ok
+CREATE TABLE t73372 (
+  i INT NOT NULL,
+  s STRING NOT NULL,
+  v INT AS (i) VIRTUAL NOT NULL,
+  INDEX idx (i) WHERE v >= 0
+)
+
+statement ok
+INSERT INTO t73372 (i, s) VALUES (0, 'foo')
+
+statement ok
+ALTER TABLE t73372 ALTER PRIMARY KEY USING COLUMNS (s, i)
+
+query ITI
+SELECT * FROM t73372
+----
+0  foo  0