web-dev-qa-db-fra.com

Comment ajouter une colonne constante dans un DataFrame Spark?

Je veux ajouter une colonne dans une DataFrame avec une valeur arbitraire (identique pour chaque ligne). J'obtiens une erreur lorsque j'utilise withColumn comme suit:

dt.withColumn('new_column', 10).head(5)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-50-a6d0257ca2be> in <module>()
      1 dt = (messages
      2     .select(messages.fromuserid, messages.messagetype, floor(messages.datetime/(1000*60*5)).alias("dt")))
----> 3 dt.withColumn('new_column', 10).head(5)

/Users/evanzamir/spark-1.4.1/python/pyspark/sql/dataframe.pyc in withColumn(self, colName, col)
   1166         [Row(age=2, name=u'Alice', age2=4), Row(age=5, name=u'Bob', age2=7)]
   1167         """
-> 1168         return self.select('*', col.alias(colName))
   1169 
   1170     @ignore_unicode_prefix

AttributeError: 'int' object has no attribute 'alias'

Il semble que je puisse faire en sorte que la fonction fonctionne comme je le souhaite en ajoutant et en soustrayant une des autres colonnes (pour qu’elles s’ajoutent à zéro), puis en ajoutant le nombre que je veux (10 dans ce cas):

dt.withColumn('new_column', dt.messagetype - dt.messagetype + 10).head(5)
[Row(fromuserid=425, messagetype=1, dt=4809600.0, new_column=10),
 Row(fromuserid=47019141, messagetype=1, dt=4809600.0, new_column=10),
 Row(fromuserid=49746356, messagetype=1, dt=4809600.0, new_column=10),
 Row(fromuserid=93506471, messagetype=1, dt=4809600.0, new_column=10),
 Row(fromuserid=80488242, messagetype=1, dt=4809600.0, new_column=10)]

C'est suprêmement hacky, non? Je suppose qu'il existe un moyen plus légitime de le faire?

79
Evan Zamir

Spark 2.2+

Spark 2.2 introduit typedLit pour prendre en charge Seq, Map et Tuples ( SPARK-19254 ) et les appels suivants doivent être pris en charge (Scala):

import org.Apache.spark.sql.functions.typedLit

df.withColumn("some_array", typedLit(Seq(1, 2, 3)))
df.withColumn("some_struct", typedLit(("foo", 1, .0.3)))
df.withColumn("some_map", typedLit(Map("key1" -> 1, "key2" -> 2)))

Spark 1.3+ (lit), 1.4+ (array, struct), 2.0+ (map):

Le deuxième argument pour DataFrame.withColumn devrait être une Column, vous devez donc utiliser un littéral:

from pyspark.sql.functions import lit

df.withColumn('new_column', lit(10))

Si vous avez besoin de colonnes complexes, vous pouvez les construire en utilisant des blocs tels que array:

from pyspark.sql.functions import array, create_map, struct

df.withColumn("some_array", array(lit(1), lit(2), lit(3)))
df.withColumn("some_struct", struct(lit("foo"), lit(1), lit(.3)))
df.withColumn("some_map", create_map(lit("key1"), lit(1), lit("key2"), lit(2)))

Les mêmes méthodes peuvent être utilisées dans Scala.

import org.Apache.spark.sql.functions.{array, lit, map, struct}

df.withColumn("new_column", lit(10))
df.withColumn("map", map(lit("key1"), lit(1), lit("key2"), lit(2)))

Pour fournir des noms pour structs, utilisez soit alias sur chaque champ:

df.withColumn(
    "some_struct",
    struct(lit("foo").alias("x"), lit(1).alias("y"), lit(0.3).alias("z"))
 )

ou cast sur l'objet entier

df.withColumn(
    "some_struct", 
    struct(lit("foo"), lit(1), lit(0.3)).cast("struct<x: string, y: integer, z: double>")
 )

Il est également possible, bien que plus lent, d’utiliser un fichier UDF.

Remarque:

Les mêmes constructions peuvent être utilisées pour passer des arguments constants à des fonctions UDF ou SQL.

152
zero323

Dans spark 2.2, il existe deux manières d'ajouter une valeur constante dans une colonne de DataFrame: 

1) Utilisation de lit 

2) Utilisation de typedLit

La différence entre les deux est que typedLit peut également gérer des types de scala paramétrés, par exemple. Liste, Seq et Carte

Exemple de DataFrame:

val df = spark.createDataFrame(Seq((0,"a"),(1,"b"),(2,"c"))).toDF("id", "col1")

+---+----+
| id|col1|
+---+----+
|  0|   a|
|  1|   b|
+---+----+

1) Utilisation de lit: Ajout de valeur constante de chaîne dans la nouvelle colonne nommée newcol:

import org.Apache.spark.sql.functions.lit
val newdf = df.withColumn("newcol",lit("myval"))

Résultat:

+---+----+------+
| id|col1|newcol|
+---+----+------+
|  0|   a| myval|
|  1|   b| myval|
+---+----+------+

2) Utilisation de typedLit:

import org.Apache.spark.sql.functions.typedLit
df.withColumn("newcol", typedLit(("sample", 10, .044)))

Résultat:

+---+----+-----------------+
| id|col1|           newcol|
+---+----+-----------------+
|  0|   a|[sample,10,0.044]|
|  1|   b|[sample,10,0.044]|
|  2|   c|[sample,10,0.044]|
+---+----+-----------------+
11
Ayush Vatsyayan