Fix nested expression

Signed-off-by: Karen Feng <karen.feng@databricks.com>
apache · Feb 27, 2021 · e1719d3 · e1719d3
1 parent 80beda8
commit e1719d3
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
@@ -990,17 +990,18 @@ class Analyzer(override val catalogManager: CatalogManager)
 
     private def getMetadataAttributes(plan: LogicalPlan): Seq[Attribute] = {
       lazy val childMetadataOutput = plan.children.flatMap(_.metadataOutput)
-      plan.expressions.collect {
+      plan.expressions.flatMap(_.collect {
         case a: Attribute if a.isMetadataCol => a
         case a: Attribute if childMetadataOutput.exists(_.exprId == a.exprId) =>
           childMetadataOutput.find(_.exprId == a.exprId).get
-      }
+      })
     }
 
     private def hasMetadataCol(plan: LogicalPlan): Boolean = {
       lazy val childMetadataOutput = plan.children.flatMap(_.metadataOutput)
       plan.expressions.exists(_.find {
-        case a: Attribute => a.isMetadataCol || childMetadataOutput.exists(_.exprId == a.exprId)
+        case a: Attribute =>
+          a.isMetadataCol || childMetadataOutput.exists(_.exprId == a.exprId)
         case _ => false
       }.isDefined)
     }
@@ -1016,9 +1017,7 @@ class Analyzer(override val catalogManager: CatalogManager)
     def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperatorsUp {
       case node if node.children.nonEmpty && node.resolved && hasMetadataCol(node) =>
         val inputAttrs = AttributeSet(node.children.flatMap(_.output))
-        val metaCols = getMetadataAttributes(node).flatMap(_.collect {
-          case a: Attribute if a.isMetadataCol && !inputAttrs.contains(a) => a
-        })
+        val metaCols = getMetadataAttributes(node).filterNot(inputAttrs.contains)
         if (metaCols.isEmpty) {
           node
         } else {