scala - example - spark sql tutorial

Crear una nueva columna con función en Spark Dataframe

Question

Estoy intentando descubrir la nueva API de marco de datos en Spark. parece un buen paso adelante pero tiene problemas para hacer algo que debería ser bastante simple. Tengo un marco de datos con 2 columnas, "ID" y "Cantidad". Como un ejemplo genérico, digamos que quiero devolver una nueva columna llamada "código" que devuelve un código basado en el valor de "Amt". Puedo escribir una función como esta:

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

Cuando trato de usarlo así:

val myDF = sqlContext.parquetFile("hdfs:/to/my/file.parquet")
myDF.withColumn("Code", coder(myDF("Amt")))

Obtengo errores de desajuste de tipo

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

He intentado cambiar el tipo de entrada en mi función a org.apache.spark.sql.Column pero luego empiezo a obtener los retornos sin la función de compilación porque quiere un valor booleano en la sentencia if.

¿Estoy haciendo esto mal? ¿Hay una forma mejor / otra de hacerlo que usar conColumna?

Gracias por tu ayuda.

Answer 1

Deberíamos evitar definir udf funciones de udf tanto como sea posible debido a su sobrecarga de serialization y deserialization de columnas.

Puede lograr la solución con simple when función de chispa es la siguiente

val myDF = sqlContext.parquetFile("hdfs:/to/my/file.parquet")
myDF.withColumn("Code", when(myDF("Amt") < 100, "Little").otherwise("Big"))

Answer 2

Otra forma de hacerlo: puede crear cualquier función, pero de acuerdo con el error anterior, debe definir la función como una variable

Ejemplo:

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

Ahora esta afirmación funciona perfectamente:

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

Answer 3

Supongamos que tiene una columna "Amt" en su esquema:

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

Creo que conColumn es la forma correcta de agregar una columna