PySpark UDF vs Pandas UDF vs mapInPandas: Qual Você Deve Usar?

Quando você precisa de lógica personalizada no PySpark, normalmente vai recorrer a uma de três ferramentas:

• UDF Python regular (udf(...))
• Pandas UDF (@pandas_udf)
• mapInPandas (DataFrame → iterador de Pandas DataFrames)

Todas conseguem “rodar Python no Spark”, mas se comportam de forma bem diferente em performance, flexibilidade e no quanto você preserva das otimizações do Spark.

Este guia traz um framework de decisão prático, além de exemplos que você pode copiar.

Modelo mental: o que muda entre as três?

1) UDF regular (Python linha a linha)

• O Spark envia colunas para processos worker de Python.
• Sua função roda uma linha por vez.
• Geralmente é a mais lenta.
• Pode bloquear o otimizador do Spark e a geração de código.

Use quando: a lógica é simples, os dados são pequenos ou velocidade não importa.

2) Pandas UDF (lotes vetorizados via Arrow)

• O Spark envia dados para o Python em lotes colunares usando Apache Arrow.
• Sua função roda sobre Pandas Series / DataFrames (vetorizado).
• Normalmente é muito mais rápida do que UDF regular.

Use quando: você precisa de lógica personalizada por coluna e quer melhor performance.

3) mapInPandas (controle total por lote de partição)

• O Spark chama sua função uma vez por chunk de partição, entregando um iterador de Pandas DataFrames.
• Você pode fazer lógica multi-coluna, transformações complexas e até expansão de linhas.
• Ótimo para etapas de “mini ETL” em Python ainda paralelizadas pelo Spark.

Use quando: você precisa de transformações complexas que não se encaixam no formato “uma coluna entra → uma coluna sai”.

Tabela rápida de decisão

Regra: funções nativas do Spark > Pandas UDF > mapInPandas > UDF regular (típico, não absoluto).

Dataset de exemplo

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import StringType, DoubleType, StructType, StructField, LongType
 
spark = SparkSession.builder.getOrCreate()
 
df = spark.createDataFrame(
    [(" Alice  ", "US", 10),
     ("bob", "UK", 3),
     (None, "US", 7)],
    ["name", "country", "visits"]
)

Exemplo de UDF regular (simples, mas lenta)

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType
 
def clean_name(x):
    if x is None:
        return None
    return x.strip().lower()
 
clean_name_udf = udf(clean_name, StringType())
 
df_udf = df.withColumn("name_clean", clean_name_udf("name"))
df_udf.show()

Prós

• Mais fácil de entender
• Funciona em todo lugar

Contras

• Overhead do Python linha a linha
• Frequentemente impede o Spark de fazer otimizações agressivas

Exemplo de Pandas UDF (vetorizada, geralmente mais rápida)

import pandas as pd
from pyspark.sql.functions import pandas_udf
 
@pandas_udf("string")
def clean_name_vec(s: pd.Series) -> pd.Series:
    return s.str.strip().str.lower()
 
df_pandas_udf = df.withColumn("name_clean", clean_name_vec("name"))
df_pandas_udf.show()

Prós

• Processamento em lote + vetorização
• Vazão (throughput) muito melhor em colunas grandes

Contras

• Requer suporte a Arrow e ambiente compatível
• Ainda roda no lado do Python e ainda é menos otimizável do que funções nativas

Exemplo de mapInPandas (mais flexível)

Caso de uso: gerar múltiplas colunas derivadas + regras customizadas

Talvez você queira:

• nome limpo
• score baseado em país e visitas
• faixas (buckets) de rótulos

import pandas as pd
 
def transform(pdf_iter):
    for pdf in pdf_iter:
        pdf["name_clean"] = pdf["name"].astype("string").str.strip().str.lower()
        pdf["visits"] = pdf["visits"].fillna(0).astype("float64")
        pdf["score"] = pdf["visits"] * pdf["country"].eq("US").astype("float64").add(1.0)  # US -> 2.0x, else 1.0x
        pdf["bucket"] = pd.cut(pdf["visits"], bins=[-1, 0, 5, 999999], labels=["none", "low", "high"])
        yield pdf
 
out_schema = "name string, country string, visits long, name_clean string, score double, bucket string"
 
df_map = df.mapInPandas(transform, schema=out_schema)
df_map.show()

Prós

• Extremamente flexível
• Ótimo para “pipelines pandas por partição”
• Pode expandir linhas, calcular múltiplas saídas, chamar libs externas (com cuidado)

Contras

• Você precisa definir o schema corretamente
• Mais chance de criar skew / partições grandes sem querer
• Ainda há overhead de Python/Arrow

E quanto à performance?

Você não precisa de benchmarks perfeitos para escolher certo. Use estas heurísticas práticas:

UDF regular geralmente é a pior quando:

• dezenas de milhões de linhas
• transformações simples que o Spark poderia fazer nativamente
• usada dentro de filtros/joins

Pandas UDF brilha quando:

• você está transformando uma ou mais colunas com operações vetorizadas
• você consegue escrever a lógica com pandas/numpy de forma eficiente

mapInPandas é melhor quando:

• você precisa de transformações em múltiplas etapas que seriam dolorosas em Spark SQL
• você quer criar múltiplas colunas de uma vez
• você precisa de expansão de linhas ou lógica condicional complexa

Pegadinhas de corretude e schema

Pandas UDF

• A saída deve bater exatamente com o tipo declarado.
• Nulos/NaNs podem aparecer; trate-os explicitamente.

mapInPandas

• O DataFrame de saída deve corresponder ao schema: nomes de colunas + dtypes + ordem.
• Cuidado com dtype object do Python; faça cast explícito para string/float.

Lista do “evite isso” (anti-patterns comuns)

• Usar UDF regular para operações básicas de string (lower, trim, regex) → use built-ins do Spark.
• Chamar APIs de rede dentro de UDFs → será lento, instável e difícil de retentar com segurança.
• Loops pesados por linha em Python dentro de Pandas UDF/mapInPandas → destrói os benefícios da vetorização.
• Retornar tipos inconsistentes (às vezes int, às vezes string) → falhas em runtime / nulos.

Fluxo de decisão recomendado

1. Dá para fazer com built-ins do Spark? → Use built-ins.

2. É uma transformação por coluna (mesmo número de linhas entrando/saindo)? → Use Pandas UDF.

3. Você precisa de lógica multi-coluna, múltiplas colunas de saída ou expansão de linhas? → Use mapInPandas.

4. É dado pequeno / protótipo rápido? → UDF regular é aceitável.

Dica final: valide com planos do Spark

Mesmo sem benchmarking completo, você aprende bastante com:

df_pandas_udf.explain(True)

Se você perceber que o Spark faz menos otimização depois de adicionar sua função, isso é um sinal para tentar built-ins ou reestruturar.