cd_Cuaderno de datos
← Artículos
Testing

Probar transformaciones de PySpark con Spark 4.3

Christian Herrera, 8 min de lectura

Un pipeline puede fallar por dos motivos distintos: porque llegan datos que no esperabas o porque tu transformación hace algo distinto de lo que crees. Las reglas de calidad en ejecución detectan lo primero. Lo segundo se detecta antes de desplegar, con tests.

Este artículo monta un proyecto mínimo para probar transformaciones de PySpark en tu máquina. Todos los ejemplos se han ejecutado con PySpark 4.3.0, Python 3.12, Java 17 y pytest 9.2.

Qué necesitas

Spark 4.3 funciona con Java 17, 21 o 25 y con Python 3.11 o superior. Si tu máquina tiene Java 11, Spark no arranca.

python3 -m venv .venv
.venv/bin/pip install "pyspark==4.3.0" pytest

No hace falta un cluster ni una cuenta de Databricks. pip instala un Spark completo que se ejecuta en local.

Escribe transformaciones que se puedan probar

La decisión que más facilita los tests se toma antes de escribirlos: cada transformación es una función que recibe un DataFrame y devuelve otro. No lee tablas, no escribe y no crea la sesión.

# transformaciones.py
from pyspark.sql import DataFrame
from pyspark.sql import functions as F


def limpiar_tickets(df: DataFrame) -> DataFrame:
    return (
        df.withColumn("ticket", F.upper(F.trim("ticket")))
        .filter(F.col("amount") > 0)
        .dropDuplicates(["ticket"])
    )


def ticket_medio(df: DataFrame) -> DataFrame:
    return df.groupBy("tienda").agg(F.avg("amount").alias("ticket_medio"))


def importe_a_numero(df: DataFrame) -> DataFrame:
    return df.withColumn("amount", F.col("amount_txt").cast("double"))

El código que lee de Kafka o escribe en una tabla llama a estas funciones. Así el test puede darles tres filas escritas a mano en lugar de una fuente real.

Una sesión de Spark para todos los tests

Arrancar Spark tarda unos segundos. Con una sesión por test, una suite de cincuenta tests tardaría minutos. La fixture con scope="session" crea una sola sesión y la comparte.

# tests/conftest.py
import os
import sys

import pytest
from pyspark.sql import SparkSession

# Los workers de Python deben usar el mismo intérprete que el driver.
os.environ["PYSPARK_PYTHON"] = sys.executable


@pytest.fixture(scope="session")
def spark():
    spark = (
        SparkSession.builder.master("local[1]")
        .appName("tests")
        .config("spark.sql.shuffle.partitions", "1")
        .config("spark.sql.session.timeZone", "UTC")
        .config("spark.ui.enabled", "false")
        .getOrCreate()
    )
    yield spark
    spark.stop()

Cada línea de configuración tiene un motivo:

Configuración Motivo
PYSPARK_PYTHON Sin ella, los workers usan el python3 del sistema. Si su versión es distinta a la del entorno virtual, los tests fallan con PYTHON_VERSION_MISMATCH.
local[1] Un solo hilo. Suficiente para pocas filas y más predecible.
spark.sql.shuffle.partitions a 1 El valor por defecto es 200. Con tres filas, cada groupBy crearía 200 tareas casi vacías.
spark.sql.session.timeZone a UTC Los tests con fechas dan el mismo resultado en tu portátil y en el servidor de CI.
spark.ui.enabled a false Los tests no necesitan la interfaz web y así no ocupan un puerto.

Comparar DataFrames

PySpark incluye dos funciones para tests en pyspark.testing: assertDataFrameEqual y assertSchemaEqual. Existen desde Spark 3.5, así que no necesitas una librería externa.

# tests/test_transformaciones.py
from pyspark.testing import assertDataFrameEqual

from transformaciones import limpiar_tickets


def test_limpiar_tickets_descarta_importes_no_positivos_y_duplicados(spark):
    entrada = spark.createDataFrame(
        [(" t-1 ", "Madrid", 12.5), ("T-1", "Madrid", 12.5), ("T-2", "Madrid", -4.0)],
        "ticket string, tienda string, amount double",
    )
    esperado = spark.createDataFrame(
        [("T-1", "Madrid", 12.5)],
        "ticket string, tienda string, amount double",
    )

    assertDataFrameEqual(limpiar_tickets(entrada), esperado)

Las tres filas de entrada cubren las tres reglas de la función: un identificador con espacios y minúsculas, un duplicado y un importe negativo.

Por defecto, assertDataFrameEqual se comporta así:

  • No compara el orden de las filas. Spark no garantiza un orden si no hay un orderBy. Si el orden forma parte del resultado, usa checkRowOrder=True.
  • Compara los decimales con tolerancia. rtol vale 1e-5 y atol vale 1e-8, de modo que 10.05 y 10.050000001 se consideran iguales.
  • No distingue columnas que admiten nulos de las que no. ignoreNullable vale True.
  • Sí compara el nombre, el tipo y el orden de las columnas.

Cuando un test falla, el mensaje indica qué filas difieren:

PySparkAssertionError: [DIFFERENT_ROWS] Results do not match: ( 50.00000 % )
*** actual ***


*** expected ***
  Row(ticket='T-1', tienda='Madrid', amount=12.5)
+ Row(ticket='T-2', tienda='Madrid', amount=-4.0)

Probar el esquema

Un cambio de tipo en una columna puede romper a quien lee tu tabla aunque los valores sean correctos. assertSchemaEqual lo comprueba sin necesitar datos:

from pyspark.sql.types import DoubleType, StringType, StructField, StructType
from pyspark.testing import assertSchemaEqual

from transformaciones import ticket_medio


def test_ticket_medio_mantiene_el_esquema(spark):
    entrada = spark.createDataFrame([], "ticket string, tienda string, amount double")
    esperado = StructType(
        [
            StructField("tienda", StringType()),
            StructField("ticket_medio", DoubleType()),
        ]
    )

    assertSchemaEqual(ticket_medio(entrada).schema, esperado)

Probar que algo falla

Desde Spark 4.0, el modo ANSI está activado por defecto. Una conversión inválida ya no devuelve NULL en silencio: lanza un error. Un test puede fijar ese comportamiento para que nadie lo cambie sin darse cuenta.

import pytest
from pyspark.errors import NumberFormatException

from transformaciones import importe_a_numero


def test_importe_no_numerico_falla_con_ansi(spark):
    entrada = spark.createDataFrame([("12,50",)], "amount_txt string")

    with pytest.raises(NumberFormatException):
        importe_a_numero(entrada).collect()

El collect() es necesario. Las transformaciones de Spark son perezosas: sin una acción, la conversión no se ejecuta y el error no aparece.

Si prefieres que un importe inválido se convierta en NULL, usa try_cast en la transformación y cambia el test para que espere ese NULL.

Ejecutar los tests

PYTHONPATH=. .venv/bin/python -m pytest tests -q

....                                                                     [100%]
4 passed in 5.60s

De esos cinco segundos, la mayor parte es el arranque de Spark. Cada test adicional añade décimas.

Qué no cubre un test unitario

Estos tests comprueban la lógica de una transformación con datos pequeños y controlados. Quedan fuera:

  • Los datos reales. Un test no sabe qué llegará mañana por el origen. Para eso están las reglas de calidad en ejecución, como las Expectations de Lakeflow.
  • El streaming con estado. Los watermarks, los datos tardíos y la recuperación desde un checkpoint dependen del tiempo y de varias ejecuciones. Necesitan una prueba de integración.
  • La plataforma. Los permisos, Unity Catalog y la configuración del pipeline se validan desplegando en un entorno de desarrollo.
  • El rendimiento. Tres filas no dicen nada sobre cómo se comporta un join con millones.

Un reparto razonable: muchos tests unitarios para la lógica, unas pocas pruebas de integración para el pipeline completo y reglas de calidad para los datos.

Abrir el curso de Lakeflow