udf

Crea una función definida por el usuario (UDF).

Syntax

import pyspark.sql.functions as sf

# As a decorator
@sf.udf
def function_name(col):
    # function body
    pass

# As a decorator with return type
@sf.udf(returnType=<returnType>, useArrow=<useArrow>)
def function_name(col):
    # function body
    pass

# As a function wrapper
sf.udf(f=<function>, returnType=<returnType>, useArrow=<useArrow>)

Parámetros

Parámetro	Tipo	Description
f	function	Optional. Función de Python si se usa como función independiente.
returnType	pyspark.sql.types.DataType o str	Optional. Tipo de valor devuelto de la función definida por el usuario. El valor puede ser un objeto DataType o una cadena de tipo con formato DDL. El valor predeterminado es StringType.
useArrow	bool	Optional. Indica si se debe usar Arrow para optimizar la serialización (de). Cuando es None, la configuración de Spark "spark.sql.execution.pythonUDF.arrow.enabled" surte efecto.

Examples

Ejemplo 1: Creación de UDF mediante lambda, decorador y decorador con tipo de valor devuelto.

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

slen = udf(lambda s: len(s), IntegerType())

@udf
def to_upper(s):
    if s is not None:
        return s.upper()

@udf(returnType=IntegerType())
def add_one(x):
    if x is not None:
        return x + 1

df = spark.createDataFrame([(1, "John Doe", 21)], ("id", "name", "age"))
df.select(slen("name").alias("slen(name)"), to_upper("name"), add_one("age")).show()

+----------+--------------+------------+
|slen(name)|to_upper(name)|add_one(age)|
+----------+--------------+------------+
|         8|      JOHN DOE|          22|
+----------+--------------+------------+

Ejemplo 2: UDF con argumentos de palabra clave.

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf, col

@udf(returnType=IntegerType())
def calc(a, b):
    return a + 10 * b

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

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

Ejemplo 3: UDF vectorizada mediante sugerencias de tipo de serie pandas.

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf, col, PandasUDFType
import pandas as pd

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

pd_calc.evalType == PandasUDFType.SCALAR
spark.range(2).select(pd_calc(b=col("id") * 10, a="id")).show()

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

Ejemplo 4: UDF vectorizada mediante sugerencias de tipo de matriz PyArrow.

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf, col, ArrowUDFType
import pyarrow as pa

@udf(returnType=IntegerType())
def pa_calc(a: pa.Array, b: pa.Array) -> pa.Array:
    return pa.compute.add(a, pa.compute.multiply(b, 10))

pa_calc.evalType == ArrowUDFType.SCALAR
spark.range(2).select(pa_calc(b=col("id") * 10, a="id")).show()

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

Ejemplo 5: UDF de Python optimizada para flechas (valor predeterminado desde Spark 4.2).

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

# Arrow optimization is enabled by default since Spark 4.2
@udf(returnType=IntegerType())
def my_udf(x):
    return x + 1

# To explicitly disable Arrow optimization and use pickle-based serialization:
@udf(returnType=IntegerType(), useArrow=False)
def legacy_udf(x):
    return x + 1

Ejemplo 6: Crear una UDF no determinista.

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf
import random

random_udf = udf(lambda: int(random.random() * 100), IntegerType()).asNondeterministic()