web-dev-qa-db-fra.com

TypeError: La colonne n'est pas itérable - Comment itérer sur ArrayType ()?

Tenez compte du DataFrame suivant:

+------+-----------------------+
|type  |names                  |
+------+-----------------------+
|person|[john, sam, jane]      |
|pet   |[whiskers, rover, fido]|
+------+-----------------------+

Qui peut être créé avec le code suivant:

import pyspark.sql.functions as f
data = [
    ('person', ['john', 'sam', 'jane']),
    ('pet', ['whiskers', 'rover', 'fido'])
]

df = sqlCtx.createDataFrame(data, ["type", "names"])
df.show(truncate=False)

Existe-t-il un moyen de modifier directement la colonne ArrayType()"names" en appliquant une fonction à chaque élément, sans utiliser de udf?

Par exemple, supposons que je veuille appliquer la fonction foo à la colonne "names". (J'utiliserai l'exemple où foo est str.upper juste à des fins d'illustration, mais ma question concerne toute fonction valide qui peut être appliquée aux éléments d'un itérable.)

foo = lambda x: x.upper()  # defining it as str.upper as an example
df.withColumn('X', [foo(x) for x in f.col("names")]).show()

TypeError: la colonne n'est pas itérable

Je pourrais le faire en utilisant un udf:

foo_udf = f.udf(lambda row: [foo(x) for x in row], ArrayType(StringType()))
df.withColumn('names', foo_udf(f.col('names'))).show(truncate=False)
#+------+-----------------------+
#|type  |names                  |
#+------+-----------------------+
#|person|[JOHN, SAM, JANE]      |
#|pet   |[WHISKERS, ROVER, FIDO]|
#+------+-----------------------+

Dans cet exemple spécifique, je pourrais éviter le udf en éclatant la colonne, appelez pyspark.sql.functions.upper(), puis groupBy et collect_list:

df.select('type', f.explode('names').alias('name'))\
    .withColumn('name', f.upper(f.col('name')))\
    .groupBy('type')\
    .agg(f.collect_list('name').alias('names'))\
    .show(truncate=False)
#+------+-----------------------+
#|type  |names                  |
#+------+-----------------------+
#|person|[JOHN, SAM, JANE]      |
#|pet   |[WHISKERS, ROVER, FIDO]|
#+------+-----------------------+

Mais c'est beaucoup de code pour faire quelque chose de simple. Existe-t-il un moyen plus direct d'itérer sur les éléments d'un ArrayType() en utilisant des fonctions spark-dataframe?

11
pault

Dans Spark <2.4 , vous pouvez utiliser une fonction définie par l'utilisateur:

from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, DataType, StringType

def transform(f, t=StringType()):
    if not isinstance(t, DataType):
       raise TypeError("Invalid type {}".format(type(t)))
    @udf(ArrayType(t))
    def _(xs):
        if xs is not None:
            return [f(x) for x in xs]
    return _

foo_udf = transform(str.upper)

df.withColumn('names', foo_udf(f.col('names'))).show(truncate=False)
+------+-----------------------+
|type  |names                  |
+------+-----------------------+
|person|[JOHN, SAM, JANE]      |
|pet   |[WHISKERS, ROVER, FIDO]|
+------+-----------------------+

Compte tenu du coût élevé de explode + collect_list idiome, cette approche est presque exclusivement préférée, malgré son coût intrinsèque.

Dans Spark 2.4 ou version ultérieure, vous pouvez utiliser transform * avec upper (voir SPARK-23909 ):

from pyspark.sql.functions import expr

df.withColumn(
    'names', expr('transform(names, x -> upper(x))')
).show(truncate=False)
+------+-----------------------+
|type  |names                  |
+------+-----------------------+
|person|[JOHN, SAM, JANE]      |
|pet   |[WHISKERS, ROVER, FIDO]|
+------+-----------------------+

Il est également possible d'utiliser pandas_udf

from pyspark.sql.functions import pandas_udf, PandasUDFType

def transform_pandas(f, t=StringType()):
    if not isinstance(t, DataType):
       raise TypeError("Invalid type {}".format(type(t)))
    @pandas_udf(ArrayType(t), PandasUDFType.SCALAR)
    def _(xs):
        return xs.apply(lambda xs: [f(x) for x in xs] if xs is not None else xs)
    return _

foo_udf_pandas = transform_pandas(str.upper)

df.withColumn('names', foo_udf(f.col('names'))).show(truncate=False)
+------+-----------------------+
|type  |names                  |
+------+-----------------------+
|person|[JOHN, SAM, JANE]      |
|pet   |[WHISKERS, ROVER, FIDO]|
+------+-----------------------+

bien que seules les dernières combinaisons Arrow/PySpark prennent en charge la gestion des colonnes ArrayType ( SPARK-24259 , SPARK-21187 ). Néanmoins, cette option devrait être plus efficace que l'UDF standard (en particulier avec une surcharge de serde inférieure) tout en prenant en charge les fonctions arbitraires Python.


* n certain nombre d'autres fonctions d'ordre supérieur sont également prises en charge , y compris, mais sans s'y limiter filter et aggregate . Voir par exemple

Oui, vous pouvez le faire en le convertissant en RDD puis en arrière en DF.

>>> df.show(truncate=False)
+------+-----------------------+
|type  |names                  |
+------+-----------------------+
|person|[john, sam, jane]      |
|pet   |[whiskers, rover, fido]|
+------+-----------------------+

>>> df.rdd.mapValues(lambda x: [y.upper() for y in x]).toDF(["type","names"]).show(truncate=False)
+------+-----------------------+
|type  |names                  |
+------+-----------------------+
|person|[JOHN, SAM, JANE]      |
|pet   |[WHISKERS, ROVER, FIDO]|
+------+-----------------------+
1
Bala