web-dev-qa-db-fra.com

Spark: Ajouter une colonne à la structure de données conditionnellement

J'essaie de prendre mes données d'entrée:

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

Et ajoutez une colonne à la fin en fonction du fait que B soit vide ou non:

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

Je peux le faire facilement en enregistrant le dataframe en entrée en tant que table temporaire, puis en tapant une requête SQL.

Mais j'aimerais vraiment savoir comment faire cela avec juste des méthodes Scala et ne pas avoir à taper une requête SQL dans Scala.

J'ai essayé .withColumn, mais je ne peux pas que ça fasse ce que je veux.

scalaapache-sparkapache-spark-sqlspark-dataframe
30
mcmcmc

Essayez withColumn avec la fonction when comme suit:

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() montre

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

J'ai ajouté la ligne (100, null, 5) Pour tester le cas isNull.

J'ai essayé ce code avec Spark 1.6.0, Mais comme indiqué dans le code de when, il fonctionne sur les versions après 1.4.0.

76
emeth

Ma mauvaise, j'avais raté une partie de la question.

La meilleure façon de nettoyer consiste à utiliser un UDF. Explication dans le code.

// 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

Que diriez-vous quelque chose comme ça?

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

Utiliser take(1) devrait avoir un impact minimal

1
Justin Pihony