UDFs ou funções nativas

O Spark roda na JVM. Uma UDF (*user-defined function*) em Python obriga cada linha a sair da JVM, ser serializada para um processo Python, voltar e ser desserializada. Além disso, o otimizador não enxerga o que a UDF faz: não empurra filtro através dela nem a combina com outras etapas.

A mesma regra, dos dois jeitos

from pyspark.sql import functions as F
from pyspark.sql.types import StringType

@F.udf(returnType=StringType())
def faixa_udf(valor):
    if valor is None:
        return "sem valor"
    return "alto" if valor >= 500 else "baixo"

com_udf = pedidos.withColumn("faixa", faixa_udf("valor"))

nativa = pedidos.withColumn(
    "faixa",
    F.when(F.col("valor").isNull(), "sem valor")
     .when(F.col("valor") >= 500, "alto")
     .otherwise("baixo"),
)

As duas dão o mesmo resultado; a nativa vira código gerado na JVM. Quase tudo tem função pronta em pyspark.sql.functions: texto (regexp_extract, split), datas (to_date, date_trunc), JSON (from_json), arrays (transform, filter) e mapas.

Se precisar de Python: pandas UDF

Uma *pandas UDF* (vetorizada) recebe um lote inteiro como pandas.Series, trocado pelo Apache Arrow. Paga a ida e volta uma vez por lote, não por linha.

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

@pandas_udf("double")
def com_imposto(valor: pd.Series) -> pd.Series:
    return valor * 1.1

pedidos.withColumn("com_imposto", com_imposto("valor"))

• UDF em Python também esconde o tipo: declare returnType, senão o padrão é texto.
• Teste a lógica da UDF como função Python comum, fora do Spark: é o teste mais barato que existe.