From ff187c38234686ede0f859f4fe00a8013d8dc86f Mon Sep 17 00:00:00 2001 From: Kevin Yu Date: Mon, 1 Feb 2016 16:08:01 -0800 Subject: [PATCH 1/2] Fixing the name resolution in drop column --- .../src/main/scala/org/apache/spark/sql/DataFrame.scala | 3 ++- .../test/scala/org/apache/spark/sql/DataFrameSuite.scala | 6 ++++++ 2 files changed, 8 insertions(+), 1 deletion(-) 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 index 518f9dcf94a70..86804c28a54c3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala @@ -1257,7 +1257,8 @@ class DataFrame private[sql]( def drop(colNames: String*): DataFrame = { val resolver = sqlContext.analyzer.resolver val remainingCols = - schema.filter(f => colNames.forall(n => !resolver(f.name, n))).map(f => Column(f.name)) + schema.filter(f => colNames.forall(n => !resolver(f.name, n))) + .map(f => Column(UnresolvedAttribute(Seq(f.name)))) if (remainingCols.size == this.schema.size) { this } else { 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 index 4ff99bdf2937d..b2aea483b3047 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala @@ -36,6 +36,12 @@ import org.apache.spark.sql.types._ class DataFrameSuite extends QueryTest with SharedSQLContext { import testImplicits._ + test("SPARK-12987") { + val df = Seq((1, 1)).toDF("a_b", "a.c") + val df1 = df.drop("a_b") + assert(df1.toString() === "[a.c: int]") + } + test("analysis error should be eagerly reported") { // Eager analysis. withSQLConf(SQLConf.DATAFRAME_EAGER_ANALYSIS.key -> "true") { From fb59a8a43ce124eb494f840d967f89810fed8f7f Mon Sep 17 00:00:00 2001 From: Kevin Yu Date: Wed, 3 Feb 2016 13:20:54 -0800 Subject: [PATCH 2/2] fixing code review comments --- .../main/scala/org/apache/spark/sql/DataFrame.scala | 5 ++++- .../scala/org/apache/spark/sql/DataFrameSuite.scala | 12 +++++++----- 2 files changed, 11 insertions(+), 6 deletions(-) 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 index 86804c28a54c3..ac20cbfdcdbe7 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala @@ -1256,9 +1256,12 @@ class DataFrame private[sql]( @scala.annotation.varargs def drop(colNames: String*): DataFrame = { val resolver = sqlContext.analyzer.resolver + val remainingCols = schema.filter(f => colNames.forall(n => !resolver(f.name, n))) - .map(f => Column(UnresolvedAttribute(Seq(f.name)))) + .map(f => Column(UnresolvedAttribute.quoted(f.name))) + + if (remainingCols.size == this.schema.size) { this } else { 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 index b2aea483b3047..6d6e60bcf60e4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala @@ -36,11 +36,6 @@ import org.apache.spark.sql.types._ class DataFrameSuite extends QueryTest with SharedSQLContext { import testImplicits._ - test("SPARK-12987") { - val df = Seq((1, 1)).toDF("a_b", "a.c") - val df1 = df.drop("a_b") - assert(df1.toString() === "[a.c: int]") - } test("analysis error should be eagerly reported") { // Eager analysis. @@ -1297,4 +1292,11 @@ class DataFrameSuite extends QueryTest with SharedSQLContext { Seq(1 -> "a").toDF("i", "j").filter($"i".cast(StringType) === "1"), Row(1, "a")) } + + test("SPARK-12987: drop column ") { + val df = Seq((1, 2)).toDF("a_b", "a.c") + val df1 = df.drop("a_b") + checkAnswer(df1, Row(2)) + assert(df1.schema.map(_.name) === Seq("a.c")) + } }