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?
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]|
+------+-----------------------+