pandas_udf

Crée une fonction définie par l'utilisateur pandas.

Les UDF Pandas sont des fonctions définies par l'utilisateur exécutées par Spark à l'aide d'Arrow pour transférer des données et de Pandas pour travailler avec les données, ce qui permet les opérations Pandas. Une UDF Pandas est définie à l'aide du pandas_udf comme décorateur ou pour envelopper la fonction, et aucune configuration supplémentaire n'est requise. Une UDF Pandas se comporte généralement comme une API de fonction PySpark standard.

Syntaxe

Python

import pyspark.sql.functions as sf

# As a decorator
@sf.pandas_udf(returnType=<returnType>, functionType=<functionType>)
def function_name(col):
    # function body
    pass

# As a function wrapper
sf.pandas_udf(f=<function>, returnType=<returnType>, functionType=<functionType>)

parameter

parameter	Type	Description
f	function	Facultatif. Fonction définie par l’utilisateur. Une fonction Python si elle est utilisée comme fonction autonome.
returnType	pyspark.sql.types.DataType OU str	Facultatif. Type de renvoi de la fonction définie par l'utilisateur. La valeur peut être soit un objet DataType, soit une chaîne de caractères de type formatée DDL.
functionType	int	Facultatif. Une valeur d'énumération dans PandasUDFType. Default: SCALAR. Ce paramètre existe pour la compatibilité. L'utilisation des annotations de type Python est encouragée.

parameter	Type	Description
f	function	Facultatif. Fonction définie par l’utilisateur. Une fonction Python si elle est utilisée comme fonction autonome.
returnType	pyspark.sql.types.DataType OU str	Facultatif. Type de renvoi de la fonction définie par l'utilisateur. La valeur peut être soit un objet DataType, soit une chaîne de caractères de type formatée DDL.
functionType	int	Facultatif. Une valeur d'énumération dans PandasUDFType. Default: SCALAR. Ce paramètre existe pour la compatibilité. L'utilisation des annotations de type Python est encouragée.

Exemples

Exemple 1 : Série à Série – Convertir les chaînes en majuscules.

Python

import pandas as pd
from pyspark.sql.functions import pandas_udf

@pandas_udf("string")
def to_upper(s: pd.Series) -> pd.Series:
    return s.str.upper()

df = spark.createDataFrame([("John Doe",)], ("name",))
df.select(to_upper("name")).show()

Output

+--------------+
|to_upper(name)|
+--------------+
|      JOHN DOE|
+--------------+

**Exemple 2** : Série à série avec des arguments de mot-clé.

Python

import pandas as pd
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import IntegerType
from pyspark.sql import functions as sf

@pandas_udf(returnType=IntegerType())
def calc(a: pd.Series, b: pd.Series) -> pd.Series:
    return a + 10 * b

spark.range(2).select(calc(b=sf.col("id") * 10, a=sf.col("id"))).show()

Output

+-----------------------------+
|calc(b => (id * 10), a => id)|
+-----------------------------+
|                            0|
|                          101|
+-----------------------------+

Exemple 3 : Itérateur de séries à itérateur de séries.

Python

import pandas as pd
from typing import Iterator
from pyspark.sql.functions import pandas_udf

@pandas_udf("long")
def plus_one(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]:
    for s in iterator:
        yield s + 1

df = spark.createDataFrame(pd.DataFrame([1, 2, 3], columns=["v"]))
df.select(plus_one(df.v)).show()

Output

+-----------+
|plus_one(v)|
+-----------+
|          2|
|          3|
|          4|
+-----------+

Exemple 4 : Séries à scalaire - Agrégation groupée.

Python

import pandas as pd
from pyspark.sql.functions import pandas_udf

@pandas_udf("double")
def mean_udf(v: pd.Series) -> float:
    return v.mean()

df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v"))
df.groupby("id").agg(mean_udf(df['v'])).show()

Output

+---+-----------+
| id|mean_udf(v)|
+---+-----------+
|  1|        1.5|
|  2|        6.0|
+---+-----------+

Exemple 5 : séries à scalaire avec fonctions de fenêtre.

Python

import pandas as pd
from pyspark.sql import Window
from pyspark.sql.functions import pandas_udf

@pandas_udf("double")
def mean_udf(v: pd.Series) -> float:
    return v.mean()

df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v"))
w = Window.partitionBy('id').orderBy('v').rowsBetween(-1, 0)
df.withColumn('mean_v', mean_udf("v").over(w)).show()

Output

+---+----+------+
| id|   v|mean_v|
+---+----+------+
|  1| 1.0|   1.0|
|  1| 2.0|   1.5|
|  2| 3.0|   3.0|
|  2| 5.0|   4.0|
|  2|10.0|   7.5|
+---+----+------+

Exemple 6 : Itérateur de série à scalaire – Agrégation groupée et efficace en mémoire.

Python

import pandas as pd
from typing import Iterator
from pyspark.sql.functions import pandas_udf

@pandas_udf("double")
def pandas_mean_iter(it: Iterator[pd.Series]) -> float:
    sum_val = 0.0
    cnt = 0
    for v in it:
        sum_val += v.sum()
        cnt += len(v)
    return sum_val / cnt

df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v"))
df.groupby("id").agg(pandas_mean_iter(df['v'])).show()

Output

+---+-------------------+
| id|pandas_mean_iter(v)|
+---+-------------------+
|  1|                1.5|
|  2|                6.0|
+---+-------------------+