Come aggiungere una colonna costante in Spark DataFrame?

Question

Voglio aggiungere una colonna in un DataFrame con un valore arbitrario (che è lo stesso per ogni riga). Ottengo un errore quando uso withColumn come segue:

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' 

Sembra che posso ingannare la funzione a lavorare come voglio aggiungendo e sottraendo una delle altre colonne (in modo da aggiungere a zero) e poi aggiungendo il numero che voglio (10 in questo caso):

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)] 

Questo è estremamente hacky, giusto? Presumo che ci sia un modo più legittimo per farlo?

Answer 1

Spark 2.2+

Spark 2.2 introduce typedLit per sostenere Seq, Map, e Tuples (SPARK-19254) e in seguito a inviti dovrebbero essere sostenuti (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):

Il secondo argomento per DataFrame.withColumn dovrebbe essere un Column quindi bisogna utilizzare un letterale:

from pyspark.sql.functions import lit 

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

Se avete bisogno di colonne complesse si possono costruire questi utilizzando blocchi come 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))) 

esattamente gli stessi metodi possono essere utilizzati in 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))) 

È anche possibile, anche se più lento, utilizzare un UDF.

Answer 2

In scintilla 2.2 ci sono due modi per aggiungere valore costante in una colonna in dataframe:

1) Uso lit

2) Utilizzo typedLit.

La differenza tra i due è che typedLit può anche gestire tipi di scala parametrizzati, ad es.Lista, Seq, e mappa

dataframe Esempio:

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

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

1) Utilizzando lit: aggiunta costante valore stringa nella nuova colonna denominata newcol:

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

Risultato:

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

2) Usando typedLit:

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

Risultato:

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