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

Author: dilipbiswal

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

@@ -1256,8 +1256,12 @@ class DataFrame private[sql](
   @scala.annotation.varargs
   def drop(colNames: String*): DataFrame = {
     val resolver = sqlContext.analyzer.resolver
+    val output = queryExecution.analyzed.output
+    val unquotedCols = colNames.map (
+      f => f.stripPrefix("`").stripSuffix("`")
+    )
     val remainingCols =
-      schema.filter(f => colNames.forall(n => !resolver(f.name, n))).map(f => Column(f.name))
+      output.filter(f => unquotedCols.forall(n => !resolver(f.name, n))).map(f => Column(f))
     if (remainingCols.size == this.schema.size) {
       this
     } else {
@@ -1276,7 +1280,9 @@ class DataFrame private[sql](
   def drop(col: Column): DataFrame = {
     val expression = col match {
       case Column(u: UnresolvedAttribute) =>
-        queryExecution.analyzed.resolveQuoted(u.name, sqlContext.analyzer.resolver).getOrElse(u)
+        val unquotedName = u.name.stripPrefix("`").stripSuffix("`")
+        queryExecution.analyzed.resolveQuoted(s"`${unquotedName}`",
+          sqlContext.analyzer.resolver).getOrElse(u)
       case Column(expr: Expression) => expr
     }
     val attrs = this.logicalPlan.output

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"))
+  }
 }