Repository: spark Updated Branches: refs/heads/master 555b2086a -> 40a10d767
[SPARK-9392][SQL] Dataframe drop should work on unresolved columns Dataframe drop should work on unresolved columns Author: Yanbo Liang <[email protected]> Closes #8821 from yanboliang/spark-9392. Project: http://git-wip-us.apache.org/repos/asf/spark/repo Commit: http://git-wip-us.apache.org/repos/asf/spark/commit/40a10d76 Tree: http://git-wip-us.apache.org/repos/asf/spark/tree/40a10d76 Diff: http://git-wip-us.apache.org/repos/asf/spark/diff/40a10d76 Branch: refs/heads/master Commit: 40a10d7675578f8370d07e23810d9fc5d58e0550 Parents: 555b208 Author: Yanbo Liang <[email protected]> Authored: Wed Oct 21 17:50:33 2015 -0700 Committer: Reynold Xin <[email protected]> Committed: Wed Oct 21 17:50:33 2015 -0700 ---------------------------------------------------------------------- sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala | 7 ++++++- .../src/test/scala/org/apache/spark/sql/DataFrameSuite.scala | 6 +++--- 2 files changed, 9 insertions(+), 4 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/spark/blob/40a10d76/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala ---------------------------------------------------------------------- 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 174bc6f..2f10aa9 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 : _*) } http://git-wip-us.apache.org/repos/asf/spark/blob/40a10d76/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala ---------------------------------------------------------------------- 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 6424f1f..f4c7aa3 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") { --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
