From aeda903f4c8e741b4c4a3bf9b83484e156289195 Mon Sep 17 00:00:00 2001
From: Drew Kimball <drewk@cockroachlabs.com>
Date: Tue, 6 Jan 2026 18:21:40 -0600
Subject: [PATCH] opt: fix PruneUnionAllCols panic with outer scope columns

Before this change, the PruneUnionAllCols normalization rule would panic
in crdb-test builds when the projection above a UnionAll referenced
columns from an outer scope (e.g., due to an apply-join or routine).
This occurred because the rule computed the needed column set by
combining ProjectionOuterCols and passthrough columns, which could
include outer scope columns not present in the UnionAll's output. These
outer columns were then passed to NeededColMapLeft/Right, which call
TranslateColSetStrict and panic when given unknown columns.

This change fixes the issue by intersecting the needed column set with
the UnionAll's output columns before passing it to NeededColMapLeft/
Right. This ensures only columns actually present in the UnionAll are
translated, preventing the panic.

Fixes #159793

Release note: None

Co-Authored-By: Claude <noreply@anthropic.com>
---
 pkg/sql/opt/norm/rules/prune_cols.opt      | 13 +++-
 pkg/sql/opt/norm/testdata/rules/prune_cols | 90 ++++++++++++++++++++++
 2 files changed, 100 insertions(+), 3 deletions(-)

diff --git a/pkg/sql/opt/norm/rules/prune_cols.opt b/pkg/sql/opt/norm/rules/prune_cols.opt
index aca0d20c7d3f..1bf30d677c57 100644
--- a/pkg/sql/opt/norm/rules/prune_cols.opt
+++ b/pkg/sql/opt/norm/rules/prune_cols.opt
@@ -604,6 +604,10 @@
 # filter, leading to an additional column being kept on just the right side).
 # If extraneous, these Projects may be cleaned up later by rules like
 # EliminateProject.
+#
+# Note: The projections could reference columns from an outer scope, e.g. due
+# to an apply-join or routine. We intersect with the UnionAll's output to ensure
+# that $needed only contains columns from the UnionAll.
 [PruneUnionAllCols, Normalize]
 (Project
     $union:(UnionAll $left:* $right:* $colmap:*)
@@ -611,9 +615,12 @@
     $passthrough:* &
         (CanPruneCols
             $union
-            $needed:(UnionCols
-                (ProjectionOuterCols $projections)
-                $passthrough
+            $needed:(IntersectionCols
+                (UnionCols
+                    (ProjectionOuterCols $projections)
+                    $passthrough
+                )
+                (OutputCols $union)
             )
         )
 )
diff --git a/pkg/sql/opt/norm/testdata/rules/prune_cols b/pkg/sql/opt/norm/testdata/rules/prune_cols
index a596b7993329..2ee663792d31 100644
--- a/pkg/sql/opt/norm/testdata/rules/prune_cols
+++ b/pkg/sql/opt/norm/testdata/rules/prune_cols
@@ -5528,3 +5528,93 @@ distinct-on
       │    └── columns: c100478.region:5!null p_id:7!null
       └── filters
            └── p_id:7 = p.id:2 [outer=(2,7), constraints=(/2: (/NULL - ]; /7: (/NULL - ]), fd=(2)==(7), (7)==(2)]
+
+# Regression test for #159793. PruneUnionAllCols must intersect the needed
+# columns with the UnionAll's output columns before calling NeededColMapLeft/
+# Right, since the projections may reference outer columns.
+exec-ddl
+CREATE TABLE t159793 (k INT PRIMARY KEY)
+----
+
+norm format=hide-all
+SELECT (
+  WITH cte (col) AS (
+    SELECT * FROM (VALUES (false))
+    UNION
+    SELECT * FROM (VALUES ((NOT true) AND false), (false))
+    UNION ALL
+    SELECT * FROM (VALUES (NULL::BOOL IN (SELECT true FROM t159793 AS t1)))
+  )
+  SELECT t.k::TEXT FROM cte WHERE col
+) FROM t159793 AS t
+----
+project
+ ├── ensure-distinct-on
+ │    ├── left-join-apply
+ │    │    ├── scan t159793 [as=t]
+ │    │    ├── project
+ │    │    │    ├── union-all
+ │    │    │    │    ├── project
+ │    │    │    │    │    └── limit
+ │    │    │    │    │         ├── select
+ │    │    │    │    │         │    ├── values
+ │    │    │    │    │         │    │    ├── (false,)
+ │    │    │    │    │         │    │    └── (false,)
+ │    │    │    │    │         │    └── filters
+ │    │    │    │    │         │         └── column1
+ │    │    │    │    │         └── 1
+ │    │    │    │    └── project
+ │    │    │    │         └── select
+ │    │    │    │              ├── values
+ │    │    │    │              │    └── tuple
+ │    │    │    │              │         └── any: eq
+ │    │    │    │              │              ├── project
+ │    │    │    │              │              │    ├── scan t159793 [as=t1]
+ │    │    │    │              │              │    └── projections
+ │    │    │    │              │              │         └── true
+ │    │    │    │              │              └── CAST(NULL AS BOOL)
+ │    │    │    │              └── filters
+ │    │    │    │                   └── column1
+ │    │    │    └── projections
+ │    │    │         └── t.k::STRING
+ │    │    └── filters (true)
+ │    └── aggregations
+ │         └── const-agg
+ │              └── k
+ └── projections
+      └── k
+
+# Regression test for #159793 with PL/pgSQL. The projection references a
+# PL/pgSQL variable which is an outer scope column.
+norm
+DO $$
+DECLARE
+  decl OIDVECTOR;
+BEGIN
+  WHILE true LOOP
+  CONTINUE WHEN false;
+  WITH cte (col) AS
+    (
+      SELECT * FROM (VALUES (false))
+    UNION
+      SELECT * FROM (VALUES ((NOT true) AND false), (false))
+    UNION ALL
+      SELECT * FROM (VALUES (NULL::BOOL IN (SELECT true FROM t159793 AS t1)))
+    )
+    SELECT decl FROM cte WHERE col INTO decl;
+  END LOOP;
+END;
+$$
+----
+call
+ ├── cardinality: [0 - 0]
+ ├── volatile
+ └── procedure: inline_code_block
+      ├── strict
+      └── body
+           └── values
+                ├── columns: stmt_loop_2:30
+                ├── cardinality: [1 - 1]
+                ├── key: ()
+                ├── fd: ()-->(30)
+                └── (stmt_loop_2(CAST(NULL AS OIDVECTOR)),)