Cómo agregar datos con groupBy en PySpark: ejemplos

Agrupar filas por una categoría y calcular totales — ventas por región, pedidos por cliente o importe por producto — es uno de los análisis más frecuentes en el trabajo con datos. En PySpark esta tarea se resuelve con el método groupBy, que funciona exactamente igual sobre mil o mil millones de filas. A lo largo de esta guía vas a crear un DataFrame de ejemplo y agregarlo con recuentos, sumas y medias, terminando con una comprobación simple para asegurar que los números son correctos.

Requisitos previos

• Python 3.8 o superior instalado.
• PySpark instalado (pip install pyspark).
• Saber leer y crear un DataFrame básico.
• Un entorno para ejecutar código: un notebook, Databricks o la terminal.

Paso 1: Crear la SparkSession y un DataFrame de ejemplo

Todo en PySpark empieza con una SparkSession, el punto de entrada para trabajar con datos distribuidos. Vamos a crear una sesión y, a partir de una lista de tuplas, un pequeño DataFrame con ventas por región y producto. Así tenemos datos concretos donde aplicar el groupBy.

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("groupby-exemplo").getOrCreate()

dados = [
    ("Norte", "Portatil", 1200),
    ("Norte", "Rato", 25),
    ("Norte", "Portatil", 1100),
    ("Sul", "Portatil", 1300),
    ("Sul", "Rato", 30),
    ("Centro", "Portatil", 1250),
]
colunas = ["regiao", "produto", "valor"]
df = spark.createDataFrame(dados, colunas)
df.show()

El createDataFrame recibe los datos y la lista de nombres de columnas. El df.show() muestra las seis filas originales; fíjate en que hay regiones repetidas — es justo lo que el groupBy va a juntar en un solo grupo.

Paso 2: Contar filas por grupo

La forma más simple de agregar es contar cuántas filas hay en cada grupo. Llama a groupBy con el nombre de la columna y después count().

df.groupBy("regiao").count().show()

El groupBy("regiao") devuelve un objeto GroupedData — todavía no es un resultado, es solo la definición de cómo agrupar. Solo cuando aplicas una función como count() obtienes un nuevo DataFrame, con una fila por región y el número de ventas de cada una.

Paso 3: Sumar y calcular medias con agg

Contar rara vez basta. Para calcular varias métricas a la vez, usa el método agg con funciones del módulo functions, que suele importarse como F. El alias da un nombre legible a cada columna del resultado.

from pyspark.sql import functions as F

df.groupBy("regiao").agg(
    F.sum("valor").alias("total"),
    F.avg("valor").alias("media"),
    F.count("produto").alias("n_vendas"),
).show()

Ahora cada región tiene, en una sola fila, el total vendido, la media por venta y el número de ventas. Sin el alias, las columnas tendrían nombres automáticos como sum(valor), más difíciles de reutilizar después.

Paso 4: Agrupar por varias columnas y ordenar

Puedes pasar más de una columna al groupBy para obtener un mayor detalle — por ejemplo, el total por región y por producto a la vez. Añade orderBy para ver primero los valores más altos.

df.groupBy("regiao", "produto").agg(
    F.sum("valor").alias("total")
).orderBy(F.col("total").desc()).show()

El resultado tiene una fila por cada combinación de región y producto, ordenada de mayor a menor total. Cambiar desc() por asc() invierte el orden.

Verificar el resultado

Una buena forma de confirmar que la agregación es correcta es comparar la suma de los grupos con la suma de todas las filas: las dos deben ser iguales.

df.agg(F.sum("valor").alias("total_geral")).show()

El total general debe dar 4905. Si sumas los totales por región del Paso 3 (Norte 2325, Sul 1330 y Centro 1250) obtienes exactamente 4905 — señal de que ninguna fila quedó fuera ni se contó dos veces.

Conclusión

Con groupBy, agg y el alias ya puedes responder a la mayoría de preguntas de "totales por categoría" en PySpark, sea cual sea el volumen de datos. El siguiente paso natural son las window functions, que calculan agregados sin colapsar las filas — por ejemplo, un total acumulado por región. Pruébalo ahora: ¿puedes adaptar el Paso 3 para mostrar también el valor máximo de cada región con F.max?