examples - python spark dataframe

¿Cómo agregar una columna constante en un Spark DataFrame?

Question

Quiero agregar una columna en un DataFrame con algún valor arbitrario (que es lo mismo para cada fila). Me sale un error cuando uso withColumn siguiente manera:

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

Parece que puedo engañar a la función para que funcione como quiero al sumar y restar una de las otras columnas (para que sumen a cero) y luego sumar el número que quiero (10 en este 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)]

Esto es supremamente hacky, ¿verdad? ¿Asumo que hay una manera más legítima de hacer esto?

Answer 1

En spark 2.2 hay dos formas de agregar valor constante en una columna en DataFrame:

1) Usando lit

2) Usando typedLit .

La diferencia entre los dos es que typedLit también puede manejar tipos de typedLit parametrizados, por ejemplo, List, Seq y Map

Marco de datos de muestra:

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

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

1) Usando lit : agregando valor de cadena constante en una nueva columna llamada newcol:

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

Resultado:

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

Resultado:

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

Answer 2

Spark 2.2+

Spark 2.2 presenta typedLit para admitir Seq , Map y Tuples ( SPARK-19254 ) y se deben SPARK-19254 siguientes llamadas (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 ):

El segundo argumento para DataFrame.withColumn debe ser una Column por lo que debe usar un literal:

from pyspark.sql.functions import lit

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

Si necesita columnas complejas, puede construirlas usando bloques como 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)))

Se pueden usar exactamente los mismos métodos en 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)))

Para proporcionar nombres para structs use cualquiera de los alias en cada campo:

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

o cast sobre todo el objeto

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

También es posible, aunque más lento, usar un UDF.

Nota :

Se pueden usar las mismas construcciones para pasar argumentos constantes a UDF o funciones SQL.