[SPARK-12988][SQL] Can't drop columns that contain dots

sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala

@@ -143,13 +143,27 @@ class DataFrame private[sql](
       queryExecution.analyzed
   }
 
+  /**
+   * Resolves a column path i.e column name may contain "." or "`". .
+   */
   protected[sql] def resolve(colName: String): NamedExpression = {
     queryExecution.analyzed.resolveQuoted(colName, sqlContext.analyzer.resolver).getOrElse {
       throw new AnalysisException(
         s"""Cannot resolve column name "$colName" among (${schema.fieldNames.mkString(", ")})""")
     }
   }
 
+  /**
+   * Resolves a column name. This is called when it is required to resolve a column by its
+   * name only and not as a column path..
+   */
+  private[sql] def resolveColName(colName: String, userSuppliedName: String): Boolean = {
+    // First remove any user supplied quotes.
+    val unquotedColName = userSuppliedName.stripPrefix("`").stripSuffix("`")
+    sqlContext.analyzer.resolver(colName, unquotedColName)
+
+  }
+
   protected[sql] def numericColumns: Seq[Expression] = {
     schema.fields.filter(_.dataType.isInstanceOf[NumericType]).map { n =>
       queryExecution.analyzed.resolveQuoted(n.name, sqlContext.analyzer.resolver).get
@@ -1175,12 +1189,11 @@ class DataFrame private[sql](
    * @since 1.3.0
    */
   def withColumn(colName: String, col: Column): DataFrame = {
-    val resolver = sqlContext.analyzer.resolver
     val output = queryExecution.analyzed.output
-    val shouldReplace = output.exists(f => resolver(f.name, colName))
+    val shouldReplace = output.exists(f => resolveColName(f.name, colName))
     if (shouldReplace) {
       val columns = output.map { field =>
-        if (resolver(field.name, colName)) {
+        if (resolveColName(field.name, colName)) {
           col.as(colName)
         } else {
           Column(field)
@@ -1196,12 +1209,11 @@ class DataFrame private[sql](
    * Returns a new [[DataFrame]] by adding a column with metadata.
    */
   private[spark] def withColumn(colName: String, col: Column, metadata: Metadata): DataFrame = {
-    val resolver = sqlContext.analyzer.resolver
     val output = queryExecution.analyzed.output
-    val shouldReplace = output.exists(f => resolver(f.name, colName))
+    val shouldReplace = output.exists(f => resolveColName(f.name, colName))
     if (shouldReplace) {
       val columns = output.map { field =>
-        if (resolver(field.name, colName)) {
+        if (resolveColName(field.name, colName)) {
           col.as(colName, metadata)
         } else {
           Column(field)
@@ -1220,12 +1232,11 @@ class DataFrame private[sql](
    * @since 1.3.0
    */
   def withColumnRenamed(existingName: String, newName: String): DataFrame = {
-    val resolver = sqlContext.analyzer.resolver
     val output = queryExecution.analyzed.output
-    val shouldRename = output.exists(f => resolver(f.name, existingName))
+    val shouldRename = output.exists(f => resolveColName(f.name, existingName))
     if (shouldRename) {
       val columns = output.map { col =>
-        if (resolver(col.name, existingName)) {
+        if (resolveColName(col.name, existingName)) {
           Column(col).as(newName)
         } else {
           Column(col)
@@ -1255,9 +1266,9 @@ class DataFrame private[sql](
    */
   @scala.annotation.varargs
   def drop(colNames: String*): DataFrame = {
-    val resolver = sqlContext.analyzer.resolver
+    val output = queryExecution.analyzed.output
     val remainingCols =
-      schema.filter(f => colNames.forall(n => !resolver(f.name, n))).map(f => Column(f.name))
+      output.filter(f => colNames.forall(n => !resolveColName(f.name, n))).map(f => Column(f))
     if (remainingCols.size == this.schema.size) {
       this
     } else {
@@ -1274,16 +1285,20 @@ class DataFrame private[sql](
    * @since 1.4.1
    */
   def drop(col: Column): DataFrame = {
-    val expression = col match {
+    def expression(attr: Attribute): Expression = col match {
       case Column(u: UnresolvedAttribute) =>
-        queryExecution.analyzed.resolveQuoted(u.name, sqlContext.analyzer.resolver).getOrElse(u)
+        if (resolveColName(attr.name, u.name)) attr else u
       case Column(expr: Expression) => expr
     }
     val attrs = this.logicalPlan.output
     val colsAfterDrop = attrs.filter { attr =>
-      attr != expression
+      attr != expression(attr)
     }.map(attr => Column(attr))
-    select(colsAfterDrop : _*)
+    if (colsAfterDrop.size == this.schema.size) {
+      this
+    } else {
+      select(colsAfterDrop: _*)
+    }
   }
 
   /**

sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala

@@ -1270,4 +1270,29 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
       Seq(1 -> "a").toDF("i", "j").filter($"i".cast(StringType) === "1"),
       Row(1, "a"))
   }
+
+  test("SPARK-12988: drop columns with ` in column name") {
+    val src = Seq((1, 2, 3)).toDF("a_b", "a.b", "a.c")
+    val df = src.drop("a_b")
+    checkAnswer(df, Row(2, 3))
+    assert(df.schema.map(_.name) === Seq("a.b", "a.c"))
+    val df1 = src.drop("a.b")
+    checkAnswer(df1, Row(1, 3))
+    assert(df1.schema.map(_.name) === Seq("a_b", "a.c"))
+    val df2 = src.drop("`a.c`")
+    checkAnswer(df2, Row(1, 2))
+    assert(df2.schema.map(_.name) === Seq("a_b", "a.b"))
+    val col1 = new Column("a_b")
+    val df4 = src.drop(col1)
+    checkAnswer(df4, Row(2, 3))
+    assert(df4.schema.map(_.name) === Seq("a.b", "a.c"))
+    val col2 = new Column("a.b")
+    val df5 = src.drop(col2)
+    checkAnswer(df5, Row(1, 3))
+    assert(df5.schema.map(_.name) === Seq("a_b", "a.c"))
+    val col3 = new Column("`a.c`")
+    val df6 = src.drop(col3)
+    checkAnswer(df6, Row(1, 2))
+    assert(df6.schema.map(_.name) === Seq("a_b", "a.b"))
+  }
 }