Address some comments

Signed-off-by: Karen Feng <karen.feng@databricks.com>

Author: karenfeng

sql/catalyst/src/main/scala-2.12/org/apache/spark/sql/catalyst/expressions/AttributeMap.scala

@@ -23,6 +23,10 @@ package org.apache.spark.sql.catalyst.expressions
  * of the name, or the expected nullability).
  */
 object AttributeMap {
+  def apply[A](kvs: Map[Attribute, A]): AttributeMap[A] = {
+    new AttributeMap(kvs.map(kv => (kv._1.exprId, kv)))
+  }
+
   def apply[A](kvs: Seq[(Attribute, A)]): AttributeMap[A] = {
     new AttributeMap(kvs.map(kv => (kv._1.exprId, kv)).toMap)
   }

sql/catalyst/src/main/scala-2.13/org/apache/spark/sql/catalyst/expressions/AttributeMap.scala

@@ -23,6 +23,10 @@ package org.apache.spark.sql.catalyst.expressions
  * of the name, or the expected nullability).
  */
 object AttributeMap {
+  def apply[A](kvs: Map[Attribute, A]): AttributeMap[A] = {
+    new AttributeMap(kvs.map(kv => (kv._1.exprId, kv)))
+  }
+
   def apply[A](kvs: Seq[(Attribute, A)]): AttributeMap[A] = {
     new AttributeMap(kvs.map(kv => (kv._1.exprId, kv)).toMap)
   }
@@ -37,6 +41,8 @@ class AttributeMap[A](val baseMap: Map[ExprId, (Attribute, A)])
 
   override def get(k: Attribute): Option[A] = baseMap.get(k.exprId).map(_._2)
 
+  override def getOrElse[B1 >: A](k: Attribute, default: => B1): B1 = get(k).getOrElse(default)
+
   override def contains(k: Attribute): Boolean = get(k).isDefined
 
   override def updated[B1 >: A](key: Attribute, value: B1): Map[Attribute, B1] =

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/NestedColumnAliasing.scala

@@ -146,11 +146,7 @@ object NestedColumnAliasing {
     val nestedFieldToAlias = attributeToExtractValuesAndAliases.values.flatten.toMap
 
     // A reference attribute can have multiple aliases for nested fields.
-    val attrToAliases = new AttributeMap(
-      attributeToExtractValuesAndAliases.map { case (attr, evAliasSeq) =>
-        attr.exprId -> (attr, evAliasSeq.map(_._2))
-      }
-    )
+    val attrToAliases = AttributeMap(attributeToExtractValuesAndAliases.mapValues(_.map(_._2)))
 
     plan match {
       case Project(projectList, child) =>
@@ -245,7 +241,9 @@ object NestedColumnAliasing {
     val otherRootReferences = new mutable.ArrayBuffer[AttributeReference]()
     exprList.foreach { e =>
       collectRootReferenceAndExtractValue(e).foreach {
-        case ev: ExtractValue => nestedFieldReferences.append(ev)
+        case ev: ExtractValue =>
+          assert(ev.references.size == 1, s"$ev should have one reference")
+          nestedFieldReferences.append(ev)
         case ar: AttributeReference => otherRootReferences.append(ar)
       }
     }
@@ -306,7 +304,7 @@ object GeneratorNestedColumnAliasing {
     // when `nestedSchemaPruningEnabled` is on, nested columns will be pruned further at
     // file format readers if it is supported.
     case Project(projectList, g: Generate) if (SQLConf.get.nestedPruningOnExpressions ||
-      SQLConf.get.nestedSchemaPruningEnabled) && canPruneGenerator(g.generator) =>
+        SQLConf.get.nestedSchemaPruningEnabled) && canPruneGenerator(g.generator) =>
       // On top on `Generate`, a `Project` that might have nested column accessors.
       // We try to get alias maps for both project list and generator's children expressions.
       val attrToExtractValues = NestedColumnAliasing.getAttributeToExtractValues(
@@ -373,9 +371,7 @@ object GeneratorNestedColumnAliasing {
             val updatedGenerate = rewrittenG.copy(generatorOutput = updatedGeneratorOutput)
 
             // Replace nested column accessor with generator output.
-            val attrExprIdsOnGenerator = attrToExtractValuesOnGenerator.map { case (attr, _) =>
-              attr.exprId
-            }.toSet
+            val attrExprIdsOnGenerator = attrToExtractValuesOnGenerator.keys.map(_.exprId).toSet
             val updatedProject = p.withNewChildren(Seq(updatedGenerate)).transformExpressions {
               case f: ExtractValue if nestedFieldsOnGenerator.contains(f) =>
                 updatedGenerate.output