From d6524f902c390afd3bf1ba4f95bae83aca9e7b2d Mon Sep 17 00:00:00 2001
From: Aaron Son <aaron@dolthub.com>
Date: Wed, 24 Feb 2021 18:41:33 -0800
Subject: [PATCH 1/2] sql/analyzer/prune_columns.go: Turn of pruneColumns when
 a Subquery expression exists.

---
 sql/analyzer/prune_columns.go      | 30 ++++++++++++++++++++++++++++++
 sql/analyzer/prune_columns_test.go |  3 ++-
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/sql/analyzer/prune_columns.go b/sql/analyzer/prune_columns.go
index 6beb45c669..c165b5dd2d 100644
--- a/sql/analyzer/prune_columns.go
+++ b/sql/analyzer/prune_columns.go
@@ -56,6 +56,11 @@ func pruneColumns(ctx *sql.Context, a *Analyzer, n sql.Node, scope *Scope) (sql.
 		return n, nil
 	}
 
+	if !pruneColumnsIsSafe(n) {
+		a.Log("not pruning columns because it is not safe.")
+		return n, nil
+	}
+
 	columns := columnsUsedByNode(n)
 	findUsedColumns(columns, n)
 
@@ -72,6 +77,31 @@ func pruneColumns(ctx *sql.Context, a *Analyzer, n sql.Node, scope *Scope) (sql.
 	return fixRemainingFieldsIndexes(n, scope)
 }
 
+func pruneColumnsIsSafe(n sql.Node) bool {
+	isSafe := true
+	// We do not run pruneColumns if there is a Subquery
+	// expression, because the field rewrites and scope handling
+	// in here are not principled.
+	plan.Inspect(n, func(n sql.Node) bool {
+		if exp, ok := n.(sql.Expressioner); ok {
+			for _, e := range exp.Expressions() {
+				sql.Inspect(e, func(e sql.Expression) bool {
+					if _, ok := e.(*plan.Subquery); ok {
+						isSafe = false
+						return false
+					}
+					return true
+				})
+				if !isSafe {
+					return false
+				}
+			}
+		}
+		return isSafe
+	})
+	return isSafe
+}
+
 func columnsUsedByNode(n sql.Node) usedColumns {
 	columns := make(usedColumns)
 
diff --git a/sql/analyzer/prune_columns_test.go b/sql/analyzer/prune_columns_test.go
index 37b6403944..457c115814 100644
--- a/sql/analyzer/prune_columns_test.go
+++ b/sql/analyzer/prune_columns_test.go
@@ -357,7 +357,7 @@ func TestPruneColumns(t *testing.T) {
 			),
 		},
 		{
-			name: "Drop projected columns not in subquery",
+			name: "Keep projected columns when there is a subquery",
 			node: plan.NewProject(
 				[]sql.Expression{
 					gf(0, "t1", "foo"),
@@ -405,6 +405,7 @@ func TestPruneColumns(t *testing.T) {
 				plan.NewProject(
 					[]sql.Expression{
 						gf(0, "t1", "foo"),
+						gf(1, "t1", "bar"),
 					},
 					plan.NewResolvedTable(t1),
 				),

From 742973bcb7e34adec6be91430f7052eaddd5b949 Mon Sep 17 00:00:00 2001
From: Aaron Son <aaron@dolthub.com>
Date: Wed, 24 Feb 2021 20:16:40 -0800
Subject: [PATCH 2/2] PR feedback.

---
 sql/analyzer/prune_columns.go | 17 +++--------------
 1 file changed, 3 insertions(+), 14 deletions(-)

diff --git a/sql/analyzer/prune_columns.go b/sql/analyzer/prune_columns.go
index c165b5dd2d..06e841aa28 100644
--- a/sql/analyzer/prune_columns.go
+++ b/sql/analyzer/prune_columns.go
@@ -82,20 +82,9 @@ func pruneColumnsIsSafe(n sql.Node) bool {
 	// We do not run pruneColumns if there is a Subquery
 	// expression, because the field rewrites and scope handling
 	// in here are not principled.
-	plan.Inspect(n, func(n sql.Node) bool {
-		if exp, ok := n.(sql.Expressioner); ok {
-			for _, e := range exp.Expressions() {
-				sql.Inspect(e, func(e sql.Expression) bool {
-					if _, ok := e.(*plan.Subquery); ok {
-						isSafe = false
-						return false
-					}
-					return true
-				})
-				if !isSafe {
-					return false
-				}
-			}
+	plan.InspectExpressions(n, func(e sql.Expression) bool {
+		if _, ok := e.(*plan.Subquery); ok {
+			isSafe = false
 		}
 		return isSafe
 	})