SPARK-47320. Fix failure after branch merge

ahshahid · Jan 16, 2025 · 385e412 · 385e412
1 parent 7c0d037
commit 385e412
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -931,11 +931,16 @@ class Dataset[T] private[sql](
 
       case other => other
     }
-    val inputForProj = logicalPlan.outputSet
+    val inputForProjOpt = if (this.queryExecution.isLazyAnalysis) {
+      None
+    } else {
+      Some(logicalPlan.outputSet)
+    }
     val namedExprs = untypedCols.map(ne => (ne.named transformUp {
-      case attr: AttributeReference if attr.metadata.contains(DATASET_ID_KEY) &&
-        (!inputForProj.contains(attr) ||
-          isIncorrectlyResolved(attr, inputForProj, HashSet(id))) =>
+      case attr: AttributeReference if inputForProjOpt.isDefined &&
+        attr.metadata.contains(DATASET_ID_KEY) &&
+        (!inputForProjOpt.get.contains(attr) ||
+          isIncorrectlyResolved(attr, inputForProjOpt.get, HashSet(id))) =>
         val ua = UnresolvedAttribute(Seq(attr.name))
         ua.copyTagsFrom(attr)
         ua.setTagValue(LogicalPlan.UNRESOLVED_ATTRIBUTE_MD_TAG, attr)