Skip to content

Commit 84b8094

Browse files
tedyumarmbrus
authored andcommitted
[SPARK-11884] Drop multiple columns in the DataFrame API
See the thread Ben started: http://search-hadoop.com/m/q3RTtveEuhjsr7g/ This PR adds drop() method to DataFrame which accepts multiple column names Author: tedyu <yuzhihong@gmail.com> Closes apache#9862 from ted-yu/master.
1 parent 871e85d commit 84b8094

2 files changed

Lines changed: 23 additions & 8 deletions

File tree

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

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1261,16 +1261,24 @@ class DataFrame private[sql](
12611261
* @since 1.4.0
12621262
*/
12631263
def drop(colName: String): DataFrame = {
1264+
drop(Seq(colName) : _*)
1265+
}
1266+
1267+
/**
1268+
* Returns a new [[DataFrame]] with columns dropped.
1269+
* This is a no-op if schema doesn't contain column name(s).
1270+
* @group dfops
1271+
* @since 1.6.0
1272+
*/
1273+
@scala.annotation.varargs
1274+
def drop(colNames: String*): DataFrame = {
12641275
val resolver = sqlContext.analyzer.resolver
1265-
val shouldDrop = schema.exists(f => resolver(f.name, colName))
1266-
if (shouldDrop) {
1267-
val colsAfterDrop = schema.filter { field =>
1268-
val name = field.name
1269-
!resolver(name, colName)
1270-
}.map(f => Column(f.name))
1271-
select(colsAfterDrop : _*)
1272-
} else {
1276+
val remainingCols =
1277+
schema.filter(f => colNames.forall(n => !resolver(f.name, n))).map(f => Column(f.name))
1278+
if (remainingCols.size == this.schema.size) {
12731279
this
1280+
} else {
1281+
this.select(remainingCols: _*)
12741282
}
12751283
}
12761284

sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -378,6 +378,13 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
378378
assert(df.schema.map(_.name) === Seq("value"))
379379
}
380380

381+
test("drop columns using drop") {
382+
val src = Seq((0, 2, 3)).toDF("a", "b", "c")
383+
val df = src.drop("a", "b")
384+
checkAnswer(df, Row(3))
385+
assert(df.schema.map(_.name) === Seq("c"))
386+
}
387+
381388
test("drop unknown column (no-op)") {
382389
val df = testData.drop("random")
383390
checkAnswer(

0 commit comments

Comments
 (0)