Skip to content

[SPARK-11884] Drop multiple columns in the DataFrame API #9862

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 17 commits into from
Closed

[SPARK-11884] Drop multiple columns in the DataFrame API #9862

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
f2ca6d0
Drop multiple columns in the DataFrame API
tedyu Nov 20, 2015
01686ad
Add varargs annotation
tedyu Nov 20, 2015
e1612ca
have the single column version delegate to the new method
tedyu Nov 20, 2015
2c23f90
have the single column version delegate to the new method
tedyu Nov 20, 2015
64a959b
Correct syntax for passing varargs
tedyu Nov 20, 2015
4541231
Add test for dropping multiple columns
tedyu Nov 20, 2015
6bbe12f
Formatting
tedyu Nov 21, 2015
de0ec01
Fix scalastyle warning
tedyu Nov 21, 2015
c512662
Minor change
tedyu Nov 21, 2015
34ccee0
Address Michael's review comments
tedyu Dec 4, 2015
1e4555a
Address Scalastyle check warnings
tedyu Dec 4, 2015
394632f
Address Scalastyle check warnings
tedyu Dec 4, 2015
86c68bc
Address Scalastyle check warnings
tedyu Dec 4, 2015
b18fc6e
Address compilation error
tedyu Dec 4, 2015
331b892
Address Wenchen's comment
tedyu Dec 4, 2015
04adf76
Add missing parameter for resolver
tedyu Dec 4, 2015
a28313b
Address Wenchen's comment
tedyu Dec 5, 2015
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 16 additions & 8 deletions sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1241,16 +1241,24 @@ class DataFrame private[sql](
* @since 1.4.0
*/
def drop(colName: String): DataFrame = {
drop(Seq(colName) : _*)
}

/**
* Returns a new [[DataFrame]] with columns dropped.
* This is a no-op if schema doesn't contain column name(s).
* @group dfops
* @since 1.6.0
*/
@scala.annotation.varargs
def drop(colNames: String*): DataFrame = {
Copy link
Contributor

@marmbrus marmbrus Nov 20, 2015

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would need a varargs annotation and I don't think that we should duplicate the column resolution logic. Otherwise it might fall out of sync.

val resolver = sqlContext.analyzer.resolver
val shouldDrop = schema.exists(f => resolver(f.name, colName))
if (shouldDrop) {
val colsAfterDrop = schema.filter { field =>
val name = field.name
!resolver(name, colName)
}.map(f => Column(f.name))
select(colsAfterDrop : _*)
} else {
val remainingCols =
schema.filter(f => colNames.forall(n => !resolver(f.name, n))).map(f => Column(f.name))
if (remainingCols.size == this.schema.size) {
this
} else {
this.select(remainingCols: _*)
}
}

Expand Down
7 changes: 7 additions & 0 deletions sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -378,6 +378,13 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
assert(df.schema.map(_.name) === Seq("value"))
}

test("drop columns using drop") {
val src = Seq((0, 2, 3)).toDF("a", "b", "c")
val df = src.drop("a", "b")
checkAnswer(df, Row(3))
assert(df.schema.map(_.name) === Seq("c"))
}

test("drop unknown column (no-op)") {
val df = testData.drop("random")
checkAnswer(
Expand Down