[SPARK-12512][SQL] support column name with dot in withColumn()

xguo27 · marmbrus · commit 573ac55d7469 · 2016-01-04T12:34:04.000-08:00
Author: Xiu Guo <xguo27@gmail.com> Closes apache#10500 from xguo27/SPARK-12512.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -1171,13 +1171,17 @@ class DataFrame private[sql](
    */
   def withColumn(colName: String, col: Column): DataFrame = {
     val resolver = sqlContext.analyzer.resolver
-    val replaced = schema.exists(f => resolver(f.name, colName))
-    if (replaced) {
-      val colNames = schema.map { field =>
-        val name = field.name
-        if (resolver(name, colName)) col.as(colName) else Column(name)
+    val output = queryExecution.analyzed.output
+    val shouldReplace = output.exists(f => resolver(f.name, colName))
+    if (shouldReplace) {
+      val columns = output.map { field =>
+        if (resolver(field.name, colName)) {
+          col.as(colName)
+        } else {
+          Column(field)
+        }
       }
-      select(colNames : _*)
+      select(columns : _*)
     } else {
       select(Column("*"), col.as(colName))
     }
@@ -1188,13 +1192,17 @@ class DataFrame private[sql](
    */
   private[spark] def withColumn(colName: String, col: Column, metadata: Metadata): DataFrame = {
     val resolver = sqlContext.analyzer.resolver
-    val replaced = schema.exists(f => resolver(f.name, colName))
-    if (replaced) {
-      val colNames = schema.map { field =>
-        val name = field.name
-        if (resolver(name, colName)) col.as(colName, metadata) else Column(name)
+    val output = queryExecution.analyzed.output
+    val shouldReplace = output.exists(f => resolver(f.name, colName))
+    if (shouldReplace) {
+      val columns = output.map { field =>
+        if (resolver(field.name, colName)) {
+          col.as(colName, metadata)
+        } else {
+          Column(field)
+        }
       }
-      select(colNames : _*)
+      select(columns : _*)
     } else {
       select(Column("*"), col.as(colName, metadata))
     }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
@@ -1221,4 +1221,11 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
           " _2: bigint ... 2 more fields> ... 2 more fields> ... 2 more fields]")
 
   }
+
+  test("SPARK-12512: support `.` in column name for withColumn()") {
+    val df = Seq("a" -> "b").toDF("col.a", "col.b")
+    checkAnswer(df.select(df("*")), Row("a", "b"))
+    checkAnswer(df.withColumn("col.a", lit("c")), Row("c", "b"))
+    checkAnswer(df.withColumn("col.c", lit("c")), Row("a", "b", "c"))
+  }
 }

Original file line number	Diff line number	Diff line change
`@@ -1221,4 +1221,11 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {`
`1221`	`1221`	`" _2: bigint ... 2 more fields> ... 2 more fields> ... 2 more fields]")`
`1222`	`1222`
`1223`	`1223`	`}`
	`1224`	`+`
	`1225`	+ test("SPARK-12512: support `.` in column name for withColumn()") {
	`1226`	`+ val df = Seq("a" -> "b").toDF("col.a", "col.b")`
	`1227`	`+ checkAnswer(df.select(df("*")), Row("a", "b"))`
	`1228`	`+ checkAnswer(df.withColumn("col.a", lit("c")), Row("c", "b"))`
	`1229`	`+ checkAnswer(df.withColumn("col.c", lit("c")), Row("a", "b", "c"))`
	`1230`	`+ }`
`1224`	`1231`	`}`