Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
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
9 changes: 8 additions & 1 deletion sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,8 @@ private[sql] object DataFrame {


/**
* A collection of rows that have the same columns.
* :: Experimental ::
* A distributed collection of data organized into named columns.
*
* A [[DataFrame]] is equivalent to a relational table in Spark SQL, and can be created using
* various functions in [[SQLContext]].
Expand Down Expand Up @@ -72,6 +73,7 @@ private[sql] object DataFrame {
* }}}
*/
// TODO: Improve documentation.
@Experimental
trait DataFrame extends RDDApi[Row] {

val sqlContext: SQLContext
Expand Down Expand Up @@ -425,6 +427,11 @@ trait DataFrame extends RDDApi[Row] {
*/
def addColumn(colName: String, col: Column): DataFrame

/**
* Returns a new [[DataFrame]] with a column renamed.
*/
def renameColumn(existingName: String, newName: String): DataFrame

/**
* Returns the first `n` rows.
*/
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -238,6 +238,14 @@ private[sql] class DataFrameImpl protected[sql](
select(Column("*"), col.as(colName))
}

override def renameColumn(existingName: String, newName: String): DataFrame = {
val colNames = schema.map { field =>
val name = field.name
if (name == existingName) Column(name).as(newName) else Column(name)
}
select(colNames :_*)
}

override def head(n: Int): Array[Row] = limit(n).collect()

override def head(): Row = head(1).head
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,8 @@ private[sql] class IncomputableColumn(protected[sql] val expr: Expression) exten

override def addColumn(colName: String, col: Column): DataFrame = err()

override def renameColumn(existingName: String, newName: String): DataFrame = err()

override def head(n: Int): Array[Row] = err()

override def head(): Row = err()
Expand Down
21 changes: 21 additions & 0 deletions sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
Original file line number Diff line number Diff line change
Expand Up @@ -313,6 +313,27 @@ class DataFrameSuite extends QueryTest {
)
}

test("addColumn") {
val df = testData.toDataFrame.addColumn("newCol", col("key") + 1)
checkAnswer(
df,
testData.collect().map { case Row(key: Int, value: String) =>
Row(key, value, key + 1)
}.toSeq)
assert(df.schema.map(_.name).toSeq === Seq("key", "value", "newCol"))
}

test("renameColumn") {
val df = testData.toDataFrame.addColumn("newCol", col("key") + 1)
.renameColumn("value", "valueRenamed")
checkAnswer(
df,
testData.collect().map { case Row(key: Int, value: String) =>
Row(key, value, key + 1)
}.toSeq)
assert(df.schema.map(_.name).toSeq === Seq("key", "valueRenamed", "newCol"))
}

test("apply on query results (SPARK-5462)") {
val df = testData.sqlContext.sql("select key from testData")
checkAnswer(df("key"), testData.select('key).collect().toSeq)
Expand Down