From 6b69e948b34a28233f814462747b324c58f9a301 Mon Sep 17 00:00:00 2001
From: Marcus Gartner <marcus@cockroachlabs.com>
Date: Thu, 29 Apr 2021 17:02:11 -0700
Subject: [PATCH] opt: do not provide project ordering for columns removed when
 simplified

Previously, `projectCanProvideOrdering` did not simplify an ordering
base on an FD set if the original ordering could be provided, while
`projectBuildChildReqOrdering` always attempted to simplify an ordering.
In #64342 the behavior of `OrderingChoice.Simplify` changed  to remove
columns in groups that are not known to be equivalent by the FD. As a
result, `projectCanProvideOrdering` could return true in cases where the
simplified ordering, with some columns removed, could not be provided,
causing `ordering.projectBuildChildReqOrdering` to panic.

This commit updates `ordering.projectCanProvideOrdering` so that the
ordering is always simplified, matching the behavior of
`ordering.projectBuildChildReqOrdering`.

Fixes #64399

Release note: None
---
 .../logictest/testdata/logic_test/sqlsmith    | 29 +++++++++++++++++++
 pkg/sql/opt/ordering/project.go               | 17 +++++------
 pkg/sql/opt/ordering/project_test.go          | 10 ++++++-
 3 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/pkg/sql/logictest/testdata/logic_test/sqlsmith b/pkg/sql/logictest/testdata/logic_test/sqlsmith
index bea5a85ef6aa..26a314bf0d74 100644
--- a/pkg/sql/logictest/testdata/logic_test/sqlsmith
+++ b/pkg/sql/logictest/testdata/logic_test/sqlsmith
@@ -157,3 +157,32 @@ query T
 SELECT ARRAY[1][NULL]
 ----
 NULL
+
+# Regression: #64399 (panic due to optimizer ordering bug)
+statement ok
+CREATE TABLE ab (
+  a INT,
+  b INT AS (a % 2) STORED,
+  INDEX (b)
+);
+CREATE TABLE cd (
+  c INT,
+  d INT AS (c % 2) VIRTUAL
+);
+
+statement ok
+SELECT NULL
+FROM ab AS ab1
+  JOIN cd AS cd1
+    JOIN cd AS cd2 ON cd1.c = cd2.c
+    JOIN ab AS ab2 ON
+      cd2.c = ab2.a
+      AND cd2.c = ab2.crdb_internal_mvcc_timestamp
+      AND cd1.c = ab2.a
+      AND cd1.c = ab2.b
+      AND cd1.d = ab2.b
+    ON ab1.b = cd1.d
+  JOIN cd AS cd3 ON
+    ab1.b = cd3.c
+    AND cd2.c = cd3.d
+    AND cd1.d = cd3.c
diff --git a/pkg/sql/opt/ordering/project.go b/pkg/sql/opt/ordering/project.go
index 3a87400f2bb8..927c8dacf779 100644
--- a/pkg/sql/opt/ordering/project.go
+++ b/pkg/sql/opt/ordering/project.go
@@ -22,18 +22,15 @@ func projectCanProvideOrdering(expr memo.RelExpr, required *physical.OrderingCho
 	proj := expr.(*memo.ProjectExpr)
 	inputCols := proj.Input.Relational().OutputCols
 
-	if required.CanProjectCols(inputCols) {
-		return true
-	}
-
-	// We may be able to "remap" columns using the internal FD set.
-	if fdSet := proj.InternalFDs(); required.CanSimplify(fdSet) {
-		simplified := required.Copy()
+	// Use a simplified ordering if it exists. This must be kept consistent with
+	// projectBuildChildReqOrdering, which always simplifies the ordering if
+	// possible.
+	simplified := *required
+	if fdSet := proj.InternalFDs(); simplified.CanSimplify(fdSet) {
+		simplified = required.Copy()
 		simplified.Simplify(fdSet)
-		return simplified.CanProjectCols(inputCols)
 	}
-
-	return false
+	return simplified.CanProjectCols(inputCols)
 }
 
 func projectBuildChildReqOrdering(
diff --git a/pkg/sql/opt/ordering/project_test.go b/pkg/sql/opt/ordering/project_test.go
index 73f043765856..7001bbb65b03 100644
--- a/pkg/sql/opt/ordering/project_test.go
+++ b/pkg/sql/opt/ordering/project_test.go
@@ -39,7 +39,7 @@ func TestProject(t *testing.T) {
 
 	input := &testexpr.Instance{
 		Rel: &props.Relational{
-			OutputCols: opt.MakeColSet(1, 2, 3, 4),
+			OutputCols: opt.MakeColSet(1, 2, 3, 4, 6),
 			FuncDeps:   fds,
 		},
 	}
@@ -73,6 +73,14 @@ func TestProject(t *testing.T) {
 			req: "+5",
 			exp: "no",
 		},
+		{
+			// Regression test for #64399. projectCanProvideOrdering should not
+			// return true when the columns remaining in the ordering after
+			// simplification cannot be provided. This causes
+			// projectBuildChildReqOrdering to panic.
+			req: "+(5|6)",
+			exp: "no",
+		},
 	}
 	for _, tc := range testCases {
 		req := physical.ParseOrderingChoice(tc.req)