web-dev-qa-db-ja.com

Spark：条件付きでデータフレームに列を追加します

入力データを取得しようとしています：

A    B       C
--------------
4    blah    2
2            3
56   foo     3

そして、Bが空かどうかに基づいて列を最後に追加します。

A    B       C     D
--------------------
4    blah    2     1
2            3     0
56   foo     3     1

入力データフレームを一時テーブルとして登録し、SQLクエリを入力することで、これを簡単に行うことができます。

しかし、Scalaメソッドだけでこれを行う方法を知りたいので、Scala内でSQLクエリを入力する必要はありません。

.withColumnを試してみましたが、私が望むことをすることができません。

scalaapache-sparkapache-spark-sqlspark-dataframe
30
mcmcmc

次のように、withColumn関数でwhenを試してください。

val sqlContext = new SQLContext(sc)
import sqlContext.implicits._ // for `toDF` and $""
import org.Apache.spark.sql.functions._ // for `when`

val df = sc.parallelize(Seq((4, "blah", 2), (2, "", 3), (56, "foo", 3), (100, null, 5)))
    .toDF("A", "B", "C")

val newDf = df.withColumn("D", when($"B".isNull or $"B" === "", 0).otherwise(1))

newDf.show()ショー

+---+----+---+---+
|  A|   B|  C|  D|
+---+----+---+---+
|  4|blah|  2|  1|
|  2|    |  3|  0|
| 56| foo|  3|  1|
|100|null|  5|  0|
+---+----+---+---+

isNullケースをテストするために(100, null, 5)行を追加しました。

このコードをSpark 1.6.0で試しましたが、whenのコードでコメントされているように、1.4.0以降のバージョンで動作します。

76
emeth

悪いのは、質問の一部を見逃していたことです。

最善の、最もクリーンな方法は、UDFを使用することです。コード内の説明。

// create some example data...BY DataFrame
// note, third record has an empty string
case class Stuff(a:String,b:Int)
val d= sc.parallelize(Seq( ("a",1),("b",2),
     ("",3) ,("d",4)).map { x => Stuff(x._1,x._2)  }).toDF

// now the good stuff.
import org.Apache.spark.sql.functions.udf
// function that returns 0 is string empty 
val func = udf( (s:String) => if(s.isEmpty) 0 else 1 )
// create new dataframe with added column named "notempty"
val r = d.select( $"a", $"b", func($"a").as("notempty") )

    scala> r.show
+---+---+--------+
|  a|  b|notempty|
+---+---+--------+
|  a|  1|    1111|
|  b|  2|    1111|
|   |  3|       0|
|  d|  4|    1111|
+---+---+--------+
3
Roberto Congiu

このようなものはどうですか？

val newDF = df.filter($"B" === "").take(1) match {
  case Array() => df
  case _ => df.withColumn("D", $"B" === "")
}

take(1)を使用すると、最小限のヒットがあります。

1
Justin Pihony