another method

Author: zhengruifeng

python/pyspark/sql/tests/test_dataframe.py

@@ -212,7 +212,7 @@ def test_drop_II(self):
         df3 = df1.join(df2, df1["colC"] == df2["colC"]).withColumn(
             "colB",
             when(
-                df1["colB"] == "b", concat(df1["colB"].cast("string"), lit("_newValue"))
+                df1["colB"] == "b", concat(df1["colB"].cast("string"), lit("x"))
             ).otherwise(df1["colB"]),
         )
         df4 = df3.drop(df1["colB"])

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -454,7 +454,6 @@ class Analyzer(override val catalogManager: CatalogManager) extends RuleExecutor
       ResolveNaturalAndUsingJoin ::
       ResolveOutputRelation ::
       new ResolveTableConstraints(catalogManager) ::
-      new ResolveDataFrameDropColumns(catalogManager) ::
       new ResolveSetVariable(catalogManager) ::
       ExtractWindowExpressions ::
       GlobalAggregates ::
@@ -1483,6 +1482,8 @@ class Analyzer(override val catalogManager: CatalogManager) extends RuleExecutor
       new ResolveReferencesInUpdate(catalogManager)
     private val resolveReferencesInSort =
       new ResolveReferencesInSort(catalogManager)
+    private val resolveDataFrameDropColumns =
+      new ResolveDataFrameDropColumns(catalogManager)
 
     /**
      * Return true if there're conflicting attributes among children's outputs of a plan
@@ -1791,6 +1792,9 @@ class Analyzer(override val catalogManager: CatalogManager) extends RuleExecutor
       // Pass for Execute Immediate as arguments will be resolved by [[SubstituteExecuteImmediate]].
       case e : ExecuteImmediateQuery => e
 
+      case d: DataFrameDropColumns if !d.resolved =>
+        resolveDataFrameDropColumns(d)
+
       case q: LogicalPlan =>
         logTrace(s"Attempting to resolve ${q.simpleString(conf.maxToStringFields)}")
         q.mapExpressions(resolveExpressionByPlanChildren(_, q, includeLastResort = true))

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala

@@ -509,6 +509,27 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
       includeLastResort = includeLastResort)
   }
 
+  // Tries to resolve `UnresolvedAttribute` by the children with Plan Ids.
+  // Returns `None` if fail to resolve.
+  private[spark] def tryResolveUnresolvedAttributeByPlanChildren(
+      u: UnresolvedAttribute,
+      q: LogicalPlan,
+      includeLastResort: Boolean = false): Option[Expression] = {
+    resolveDataFrameColumn(u, q.children).map { r =>
+      resolveExpression(
+        r,
+        resolveColumnByName = nameParts => {
+          q.resolveChildren(nameParts, conf.resolver)
+        },
+        getAttrCandidates = () => {
+          assert(q.children.length == 1)
+          q.children.head.output
+        },
+        throws = true,
+        includeLastResort = includeLastResort)
+    }
+  }
+
   /**
    * The last resort to resolve columns. Currently it does two things:
    *  - Try to resolve column names as outer references
@@ -538,7 +559,6 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
       e: Expression,
       q: Seq[LogicalPlan]): Expression = e match {
     case u: UnresolvedAttribute =>
-
       resolveDataFrameColumn(u, q).getOrElse(u)
     case u: UnresolvedDataFrameStar =>
       resolveDataFrameStar(u, q)
@@ -566,14 +586,7 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
       //  df1.select(df2.a)   <-   illegal reference df2.a
       throw QueryCompilationErrors.cannotResolveDataFrameColumn(u)
     }
-
-    if (resolved.nonEmpty) {
-      resolved.map(_._1)
-    } else if (u.getTagValue(LogicalPlan.ALLOW_NON_EXISTENT_COL).nonEmpty) {
-      Some(NonExistentAttribute(u.name))
-    } else {
-      None
-    }
+    resolved.map(_._1)
   }
 
   private def resolveDataFrameColumnByPlanId(

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDataFrameDropColumns.scala

@@ -17,7 +17,6 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
-import org.apache.spark.sql.catalyst.expressions.NonExistentAttribute
 import org.apache.spark.sql.catalyst.plans.logical.{DataFrameDropColumns, LogicalPlan, Project}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.trees.TreePattern.DF_DROP_COLUMNS
@@ -36,8 +35,14 @@ class ResolveDataFrameDropColumns(val catalogManager: CatalogManager)
       // expressions in dropList can be unresolved, e.g.
       //   df.drop(col("non-existing-column"))
       val dropped = d.dropList.flatMap {
-        case u: UnresolvedAttribute => Some(resolveExpressionByPlanChildren(u, d))
-        case n: NonExistentAttribute => None
+        case u: UnresolvedAttribute =>
+          if (u.getTagValue(LogicalPlan.PLAN_ID_TAG).nonEmpty) {
+            // Plan ID comes from Spark Connect,
+            // here we ignore the column if fail to resolve by plan Id.
+            tryResolveUnresolvedAttributeByPlanChildren(u, d)
+          } else {
+            Some(resolveExpressionByPlanChildren(u, d))
+          }
         case e => Some(e)
       }
       val remaining = d.child.output.filterNot(attr => dropped.exists(_.semanticEquals(attr)))

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala

@@ -790,29 +790,3 @@ object FileSourceGeneratedMetadataAttribute {
       .map(attr -> _)
   }
 }
-
-
-case class NonExistentAttribute(name: String) extends Attribute with Unevaluable {
-  override def dataType: DataType = NullType
-  override def canEqual(that: Any): Boolean = that.isInstanceOf[NonExistentAttribute]
-
-  override def withNullability(newNullability: Boolean): Attribute =
-    throw SparkUnsupportedOperationException()
-  override def newInstance(): Attribute =
-    throw SparkUnsupportedOperationException()
-  override def withQualifier(newQualifier: Seq[String]): Attribute =
-    throw SparkUnsupportedOperationException()
-  override def withName(newName: String): Attribute =
-    throw SparkUnsupportedOperationException()
-  override def withMetadata(newMetadata: Metadata): Attribute =
-    throw SparkUnsupportedOperationException()
-  override def qualifier: Seq[String] =
-    throw SparkUnsupportedOperationException()
-  override def exprId: ExprId =
-    throw SparkUnsupportedOperationException()
-  override def withExprId(newExprId: ExprId): Attribute =
-    throw SparkUnsupportedOperationException()
-  override def withDataType(newType: DataType): Attribute =
-    throw SparkUnsupportedOperationException()
-  override def nullable: Boolean = true
-}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala

@@ -217,8 +217,6 @@ object LogicalPlan {
   //    to the old code path.
   private[spark] val PLAN_ID_TAG = TreeNodeTag[Long]("plan_id")
   private[spark] val IS_METADATA_COL = TreeNodeTag[Unit]("is_metadata_col")
-  // Whether an expression can be resolved to a non-existent column.
-  private[spark] val ALLOW_NON_EXISTENT_COL = TreeNodeTag[Unit]("allow_non_existent_col")
 }
 
 /**

sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala

@@ -2420,11 +2420,7 @@ class SparkConnectPlanner(
   private def transformDrop(rel: proto.Drop): LogicalPlan = {
     var output = Dataset.ofRows(session, transformRelation(rel.getInput))
     if (rel.getColumnsCount > 0) {
-      val cols = rel.getColumnsList.asScala.toSeq.map { expr =>
-        val e = transformExpression(expr)
-        e.setTagValue(LogicalPlan.ALLOW_NON_EXISTENT_COL, ())
-        Column(e)
-      }
+      val cols = rel.getColumnsList.asScala.toSeq.map(expr => Column(transformExpression(expr)))
       output = output.drop(cols.head, cols.tail: _*)
     }
     if (rel.getColumnNamesCount > 0) {