From d66bd159c6db093a85d81dbfefb4f4f8ed4e21dc Mon Sep 17 00:00:00 2001
From: Marcus Gartner <marcus@cockroachlabs.com>
Date: Wed, 29 Jan 2025 10:55:51 -0500
Subject: [PATCH] opt: do not project extra columns above partial index scans

The `GeneratePartialIndexScans` rule no longer projects all non-indexed
columns held constant by the partial index predicate. Now, only columns
produced by the original scan are produced. This fixes an issue where
the generated Project expression could have more output columns than the
other expressions in its group, causing a test-only assertion to fail.

Fixes #140019

There is no release note because this has caused no known correctness
bugs.

Release note: None
---
 pkg/sql/opt/xform/scan_index_iter.go    |  5 ++--
 pkg/sql/opt/xform/testdata/rules/select | 37 +++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/pkg/sql/opt/xform/scan_index_iter.go b/pkg/sql/opt/xform/scan_index_iter.go
index 4f28052b7cdf..8153556fc440 100644
--- a/pkg/sql/opt/xform/scan_index_iter.go
+++ b/pkg/sql/opt/xform/scan_index_iter.go
@@ -307,8 +307,9 @@ func (it *scanIndexIter) ForEachStartingAfter(ord int, f enumerateIndexFunc) {
 				isCovering = true
 
 				// Build a projection only for constant columns not in the
-				// index.
-				constCols = constCols.Difference(indexCols)
+				// index and produced by the original scan.
+				constCols.DifferenceWith(indexCols)
+				constCols.IntersectionWith(it.scanPrivate.Cols)
 				constProj = it.buildConstProjectionsFromPredicate(predFilters, constCols)
 			}
 		}
diff --git a/pkg/sql/opt/xform/testdata/rules/select b/pkg/sql/opt/xform/testdata/rules/select
index f796aca2ee4f..727d1f2e520f 100644
--- a/pkg/sql/opt/xform/testdata/rules/select
+++ b/pkg/sql/opt/xform/testdata/rules/select
@@ -557,6 +557,43 @@ project
       ├── columns: k:1!null
       └── key: (1)
 
+exec-ddl
+CREATE TABLE t140019 (
+  k INT PRIMARY KEY,
+  i INT,
+  j INT,
+  INDEX (k) WHERE i = 0 AND j = 1
+)
+----
+
+# Regression test for #140019. Do not project unnecessary columns, like i=0,
+# that are held constant by the partial index predicate.
+opt disable=SimplifyZeroCardinalityGroup
+SELECT NULL FROM t140019 WHERE false GROUP BY j
+----
+project
+ ├── columns: "?column?":6
+ ├── cardinality: [0 - 0]
+ ├── fd: ()-->(6)
+ ├── distinct-on
+ │    ├── columns: j:3
+ │    ├── grouping columns: j:3
+ │    ├── cardinality: [0 - 0]
+ │    ├── key: (3)
+ │    └── select
+ │         ├── columns: j:3
+ │         ├── cardinality: [0 - 0]
+ │         ├── project
+ │         │    ├── columns: j:3!null
+ │         │    ├── fd: ()-->(3)
+ │         │    ├── scan t140019@t140019_k_idx,partial
+ │         │    └── projections
+ │         │         └── 1 [as=j:3]
+ │         └── filters
+ │              └── false [constraints=(contradiction; tight)]
+ └── projections
+      └── NULL [as="?column?":6]
+
 # --------------------------------------------------
 # GenerateConstrainedScans
 # --------------------------------------------------