web-dev-qa-db-fra.com

Créer une nouvelle colonne avec une fonction dans Spark Dataframe

J'essaie de comprendre la nouvelle API Dataframe dans Spark. semble être un bon pas en avant mais avoir du mal à faire quelque chose qui devrait être assez simple. J'ai un dataframe avec 2 colonnes, "ID" et "Montant". En tant qu'exemple générique, supposons que je souhaite renvoyer une nouvelle colonne appelée "code" qui renvoie un code basé sur la valeur de "Amt". Je peux écrire une fonction qui ressemble à ceci:

def coder(myAmt:Integer):String {
  if (myAmt > 100) "Little"
  else "Big"
}

Quand j'essaye de l'utiliser comme ça:

val myDF = sqlContext.parquetFile("hdfs:/to/my/file.parquet")

myDF.withColumn("Code", coder(myDF("Amt")))

Je reçois des erreurs de correspondance de type

found   : org.Apache.spark.sql.Column
required: Integer

J'ai essayé de changer le type d'entrée de ma fonction en org.Apache.spark.sql.Column, mais je commence alors à obtenir des erreurs avec la fonction de compilation car elle veut un booléen dans l'instruction if.

Est-ce que je fais mal? Existe-t-il un moyen/meilleur moyen de faire cela que d'utiliser withColumn?

Merci de votre aide.

scalaapache-sparkdataframe
36
J Calbreath

Disons que vous avez la colonne "Amt" dans votre schéma:

import org.Apache.spark.sql.functions._
val myDF = sqlContext.parquetFile("hdfs:/to/my/file.parquet")
val coder: (Int => String) = (arg: Int) => {if (arg < 100) "little" else "big"}
val sqlfunc = udf(coder)
myDF.withColumn("Code", sqlfunc(col("Amt")))

Je pense que withColumn est la bonne façon d’ajouter une colonne

53
yjshen

Nous devrions éviter de définir autant que possible les fonctions udf en raison de la surcharge de serialization et deserialization de colonnes.

Vous pouvez obtenir la solution avec when spark) comme ci-dessous

val myDF = sqlContext.parquetFile("hdfs:/to/my/file.parquet")

myDF.withColumn("Code", when(myDF("Amt") < 100, "Little").otherwise("Big"))
11
Ramesh Maharjan

Une autre façon de faire: Vous pouvez créer n’importe quelle fonction, mais en fonction de l’erreur ci-dessus, vous devez définir la fonction comme variable.

Exemple:

val coder = udf((myAmt:Integer) => {
  if (myAmt > 100) "Little"
  else "Big"
})

Maintenant, cette déclaration fonctionne parfaitement:

myDF.withColumn("Code", coder(myDF("Amt")))
1
imran