Merge pull request #139872 from michae2/backport24.2-139388

release-24.2: opt/memo: silence error in factorOutVirtualCols for mutation columns

pkg/sql/logictest/testdata/logic_test/distsql_stats

@@ -3337,8 +3337,14 @@ mutation         {crdb_internal_a_shard_3}    10         true
 statement ok
 SET CLUSTER SETTING jobs.debug.pausepoints = ''
 
+let $job
+SELECT job_id FROM crdb_internal.jobs WHERE description LIKE 'ALTER TABLE %t118537 ALTER PRIMARY KEY USING COLUMNS (a) USING HASH' AND status LIKE 'pause%' FETCH FIRST 1 ROWS ONLY
+
+statement ok
+RESUME JOB $job
+
 statement ok
-RESUME JOB (SELECT job_id FROM crdb_internal.jobs WHERE description LIKE 'ALTER TABLE %t118537 ALTER PRIMARY KEY USING COLUMNS (a) USING HASH' AND status = 'paused' FETCH FIRST 1 ROWS ONLY)
+SHOW JOB WHEN COMPLETE $job
 
 # Test optimizer_use_virtual_computed_column_stats.
 
@@ -3507,6 +3513,78 @@ upper_bound   range_rows  distinct_range_rows  equal_rows
 'dispatched'  0           0                    1
 'delivered'   0           0                    1
 
+# Regression test for #138809: check that we can still use stats on virtual
+# computed columns after dropping the column.
+
+statement ok
+CREATE TABLE t138809 (
+  a INT PRIMARY KEY,
+  b INT AS (a + 1) VIRTUAL,
+  c INT AS (a + 2) VIRTUAL
+)
+
+statement ok
+INSERT INTO t138809 (a) VALUES (1), (2)
+
+statement ok
+ANALYZE t138809
+
+query TIIIIB colnames
+SELECT column_names, row_count, null_count, distinct_count, avg_size, histogram_id IS NOT NULL AS has_histogram
+FROM [SHOW STATISTICS FOR TABLE t138809]
+ORDER BY column_names::STRING
+----
+column_names  row_count  null_count  distinct_count  avg_size  has_histogram
+{a}           2          0           2               8         true
+{b}           2          0           2               8         true
+{c}           2          0           2               8         true
+
+query T
+EXPLAIN SELECT count(*) FROM t138809 WHERE b > 1
+----
+distribution: full
+vectorized: true
+·
+• group (scalar)
+│ estimated row count: 1
+│
+└── • scan
+      estimated row count: 2 (100% of the table; stats collected <hidden> ago)
+      table: t138809@t138809_pkey
+      spans: [/1 - ]
+
+statement ok
+SET CLUSTER SETTING jobs.debug.pausepoints = 'schemachanger.root..1'
+
+statement error job \d+ was paused before it completed with reason: pause point "schemachanger\.root\.\.1" hit
+ALTER TABLE t138809 DROP COLUMN c
+
+query T
+EXPLAIN SELECT count(*) FROM t138809 WHERE b > 1
+----
+distribution: full
+vectorized: true
+·
+• group (scalar)
+│ estimated row count: 1
+│
+└── • scan
+      estimated row count: 2 (100% of the table; stats collected <hidden> ago)
+      table: t138809@t138809_pkey
+      spans: [/1 - ]
+
+statement ok
+SET CLUSTER SETTING jobs.debug.pausepoints = DEFAULT
+
+let $job
+SELECT job_id FROM crdb_internal.jobs WHERE description LIKE 'ALTER TABLE %t138809 DROP COLUMN c' AND status LIKE 'pause%' FETCH FIRST 1 ROWS ONLY
+
+statement ok
+RESUME JOB $job
+
+statement ok
+SHOW JOB WHEN COMPLETE $job
+
 # Regression test for #67050: make sure we skip over enum values that have been
 # dropped when decoding histograms.
 

pkg/sql/opt/memo/statistics_builder.go

@@ -5150,9 +5150,16 @@ func (sb *statisticsBuilder) factorOutVirtualCols(
 		}
 		expr, ok := tab.ComputedCols[colID]
 		if !ok {
-			panic(errors.AssertionFailedf(
-				"could not find computed column expression for column %v in table %v", colID, tab.Alias,
-			))
+			// If we can't find the computed column expression, this is probably a
+			// mutation column. Whatever the reason, it's safe to skip: we simply
+			// won't factor out matching expressions.
+			if buildutil.CrdbTestBuild &&
+				!tab.Table.Column(tab.MetaID.ColumnOrdinal(colID)).IsMutation() {
+				panic(errors.AssertionFailedf(
+					"could not find computed column expression for column %v in table %v", colID, tab.Alias,
+				))
+			}
+			return
 		}
 		virtExprs = append(virtExprs, virtExpr{colID: colID, expr: expr})
 	})