From c01dad46c773146b27b8ffc09fdf83b09edefec1 Mon Sep 17 00:00:00 2001
From: Gengliang Wang <gengliang@apache.org>
Date: Wed, 12 Apr 2023 22:04:31 -0700
Subject: [PATCH] [MINOR][SQL] Simplify the method
 resolveExprsAndAddMissingAttrs

### What changes were proposed in this pull request?

The method `resolveExprsAndAddMissingAttrs` contains redundant code: getting the `newExprs` and `newChild` shows up 4 times in different branches.
This PR is to simplify the implementation of the method.

### Why are the changes needed?

Code clean up and remove redundant code.
### Does this PR introduce _any_ user-facing change?

No

### How was this patch tested?

Existing tests

Closes #40761 from gengliangwang/cleanup.

Lead-authored-by: Gengliang Wang <gengliang@apache.org>
Co-authored-by: Gengliang Wang <ltnwgl@gmail.com>
Signed-off-by: Gengliang Wang <gengliang@apache.org>
---
 .../analysis/ColumnResolutionHelper.scala     | 59 +++++++++----------
 1 file changed, 29 insertions(+), 30 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
index ba550bce791bb..c5634278490db 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala
@@ -50,40 +50,39 @@ trait ColumnResolutionHelper extends Logging {
       (exprs, plan)
     } else {
       plan match {
-        case p: Project =>
-          // Resolving expressions against current plan.
-          val maybeResolvedExprs = exprs.map(resolveExpressionByPlanOutput(_, p))
-          // Recursively resolving expressions on the child of current plan.
-          val (newExprs, newChild) = resolveExprsAndAddMissingAttrs(maybeResolvedExprs, p.child)
-          // If some attributes used by expressions are resolvable only on the rewritten child
-          // plan, we need to add them into original projection.
-          val missingAttrs = (AttributeSet(newExprs) -- p.outputSet).intersect(newChild.outputSet)
-          (newExprs, Project(p.projectList ++ missingAttrs, newChild))
-
-        case a @ Aggregate(groupExprs, aggExprs, child) =>
-          val maybeResolvedExprs = exprs.map(resolveExpressionByPlanOutput(_, a))
-          val (newExprs, newChild) = resolveExprsAndAddMissingAttrs(maybeResolvedExprs, child)
-          val missingAttrs = (AttributeSet(newExprs) -- a.outputSet).intersect(newChild.outputSet)
-          if (missingAttrs.forall(attr => groupExprs.exists(_.semanticEquals(attr)))) {
-            // All the missing attributes are grouping expressions, valid case.
-            (newExprs, a.copy(aggregateExpressions = aggExprs ++ missingAttrs, child = newChild))
-          } else {
-            // Need to add non-grouping attributes, invalid case.
-            (exprs, a)
-          }
-
-        case g: Generate =>
-          val maybeResolvedExprs = exprs.map(resolveExpressionByPlanOutput(_, g))
-          val (newExprs, newChild) = resolveExprsAndAddMissingAttrs(maybeResolvedExprs, g.child)
-          (newExprs, g.copy(unrequiredChildIndex = Nil, child = newChild))
-
         // For `Distinct` and `SubqueryAlias`, we can't recursively resolve and add attributes
         // via its children.
         case u: UnaryNode if !u.isInstanceOf[Distinct] && !u.isInstanceOf[SubqueryAlias] =>
-          val maybeResolvedExprs = exprs.map(resolveExpressionByPlanOutput(_, u))
-          val (newExprs, newChild) = resolveExprsAndAddMissingAttrs(maybeResolvedExprs, u.child)
-          (newExprs, u.withNewChildren(Seq(newChild)))
+          val (newExprs, newChild) = {
+            // Resolving expressions against current plan.
+            val maybeResolvedExprs = exprs.map(resolveExpressionByPlanOutput(_, u))
+            // Recursively resolving expressions on the child of current plan.
+            resolveExprsAndAddMissingAttrs(maybeResolvedExprs, u.child)
+          }
+          // If some attributes used by expressions are resolvable only on the rewritten child
+          // plan, we need to add them into original projection.
+          lazy val missingAttrs =
+            (AttributeSet(newExprs) -- u.outputSet).intersect(newChild.outputSet)
+          u match {
+            case p: Project =>
+              (newExprs, Project(p.projectList ++ missingAttrs, newChild))
+
+            case a @ Aggregate(groupExprs, aggExprs, child) =>
+              if (missingAttrs.forall(attr => groupExprs.exists(_.semanticEquals(attr)))) {
+                // All the missing attributes are grouping expressions, valid case.
+                (newExprs,
+                  a.copy(aggregateExpressions = aggExprs ++ missingAttrs, child = newChild))
+              } else {
+                // Need to add non-grouping attributes, invalid case.
+                (exprs, a)
+              }
+
+            case g: Generate =>
+              (newExprs, g.copy(unrequiredChildIndex = Nil, child = newChild))
 
+            case _ =>
+              (newExprs, u.withNewChildren(Seq(newChild)))
+          }
         // For other operators, we can't recursively resolve and add attributes via its children.
         case other =>
           (exprs.map(resolveExpressionByPlanOutput(_, other)), other)