diff options
Diffstat (limited to 'sql/core')
-rw-r--r-- | sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala | 7 | ||||
-rw-r--r-- | sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala | 6 |
2 files changed, 9 insertions, 4 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 174bc6f42a..2f10aa9f3c 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 @@ -1237,9 +1237,14 @@ class DataFrame private[sql]( * @since 1.4.1 */ def drop(col: Column): DataFrame = { + val expression = col match { + case Column(u: UnresolvedAttribute) => + queryExecution.analyzed.resolveQuoted(u.name, sqlContext.analyzer.resolver).getOrElse(u) + case Column(expr: Expression) => expr + } val attrs = this.logicalPlan.output val colsAfterDrop = attrs.filter { attr => - attr != col.expr + attr != expression }.map(attr => Column(attr)) select(colsAfterDrop : _*) } 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 6424f1f1d9..f4c7aa34e5 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 @@ -388,13 +388,13 @@ class DataFrameSuite extends QueryTest with SharedSQLContext { assert(df.schema.map(_.name) === Seq("key", "value")) } - test("drop unknown column with same name (no-op) with column reference") { + test("drop unknown column with same name with column reference") { val col = Column("key") val df = testData.drop(col) checkAnswer( df, - testData.collect().toSeq) - assert(df.schema.map(_.name) === Seq("key", "value")) + testData.collect().map(x => Row(x.getString(1))).toSeq) + assert(df.schema.map(_.name) === Seq("value")) } test("drop column after join with duplicate columns using column reference") { |