web-dev-qa-db-fra.com

Apache Spark - Attribuer le résultat de l'UDF à plusieurs colonnes de trame de données

J'utilise pyspark, je charge un gros fichier csv dans une trame de données avec spark-csv, et en tant qu'étape de prétraitement, je dois appliquer une variété d'opérations aux données disponibles dans l'une des colonnes (qui contient une chaîne json) . Cela renverra des valeurs X, dont chacune doit être stockée dans leur propre colonne séparée.

Cette fonctionnalité sera implémentée dans un UDF. Cependant, je ne sais pas comment renvoyer une liste de valeurs à partir de cette UDF et les insérer dans des colonnes individuelles. Voici un exemple simple:

(...)
from pyspark.sql.functions import udf
def udf_test(n):
    return [n/2, n%2]

test_udf=udf(udf_test)


df.select('amount','trans_date').withColumn("test", test_udf("amount")).show(4)

Cela produit ce qui suit:

+------+----------+--------------------+
|amount|trans_date|                test|
+------+----------+--------------------+
|  28.0|2016-02-07|         [14.0, 0.0]|
| 31.01|2016-02-07|[15.5050001144409...|
| 13.41|2016-02-04|[6.70499992370605...|
| 307.7|2015-02-17|[153.850006103515...|
| 22.09|2016-02-05|[11.0450000762939...|
+------+----------+--------------------+
only showing top 5 rows

Quelle serait la meilleure façon de stocker les deux valeurs (dans cet exemple) renvoyées par le udf sur des colonnes distinctes? À l'heure actuelle, ils sont saisis sous forme de chaînes:

df.select('amount','trans_date').withColumn("test", test_udf("amount")).printSchema()

root
 |-- amount: float (nullable = true)
 |-- trans_date: string (nullable = true)
 |-- test: string (nullable = true)
pythonapache-sparkpysparkapache-spark-sqluser-defined-functions
38
Everaldo Aguiar

Il n'est pas possible de créer plusieurs colonnes de niveau supérieur à partir d'un seul appel UDF mais vous pouvez créer un nouveau struct. Il nécessite un UDF avec returnType spécifié:

from pyspark.sql.functions import udf
from pyspark.sql.types import *

schema = StructType([
    StructField("foo", FloatType(), False),
    StructField("bar", FloatType(), False)
])

def udf_test(n):
    return (n / 2, n % 2) if n and n != 0.0 else (float('nan'), float('nan'))

test_udf = udf(udf_test, schema)
df = sc.parallelize([(1, 2.0), (2, 3.0)]).toDF(["x", "y"])

foobars = df.select(test_udf("y").alias("foobar"))
foobars.printSchema()
## root
##  |-- foobar: struct (nullable = true)
##  |    |-- foo: float (nullable = false)
##  |    |-- bar: float (nullable = false)

Vous aplatissez davantage le schéma avec un simple select:

foobars.select("foobar.foo", "foobar.bar").show()
## +---+---+
## |foo|bar|
## +---+---+
## |1.0|0.0|
## |1.5|1.0|
## +---+---+

Voir aussi Dériver plusieurs colonnes d'une seule colonne dans un Spark DataFrame

58
zero323