[Python] How can I prevent this code in polars to consume all the RAM in my computer?

I have some Polars code designed to calculate performance metrics for multiple academic terms (both past and present) using a for loop. However, this implementation is consuming all available RAM (32 GB) and doesn't seem faster than Pandas. Could you suggest improvements to optimize memory usage and efficiency? Is there an alternative approach I could use to rewrite this code more effectively?

This is my current solution:

def calculate_performance_criteria(grades_history: pl.LazyFrame,
enrolled_students: pl.LazyFrame,
student_performance: pl.LazyFrame,
period: int,
index_period: int,
undergrad_periods: list[int]) -> pl.LazyFrame:

period_min = min(undergrad_periods[max(0, index_period - 3):index_period + 1])

resumen_riesgo_acumulado = (
student_performance
.filter(pl.col("PERIODO") <= period)
.group_by(["CODIGO_ESTUDIANTE", "NIVEL_PROGRAMA_1"])
.agg([
pl.sum("EN_SUSPENSION").alias("SUSPENSIONES_ACUMULADAS"),
pl.sum("EN_PRUEBA_ACADEMICA").alias("PRUEBAS_ACADEMICAS_ACUMULADAS")
])
)

resumen_riesgo_ultimos_4 = (
student_performance
.filter(
(pl.col("PERIODO") >= period_min) &
(pl.col("PERIODO") <= period) &
(pl.col("PERIODO") % 100).is_in([10, 20])
)
.group_by("CODIGO_ESTUDIANTE")
.agg([
pl.sum("EN_SUSPENSION").alias("SUSP_ULT_4_PERIODOS"),
pl.sum("EN_PRUEBA_ACADEMICA").alias("PA_ULT_4_PERIODOS")
])
)
historial_academico_filtrado = (
grades_history
.filter(pl.col("PERIODO") <= period)
.with_columns([
(pl.col("ESTADO_MATERIA") == "INCOMPLETO").cast(pl.Int8).alias("ES_INCOMPLETO"),
pl.lit(1).alias("MATERIA_CONTEO")
])
)
resumen_incompletos_por_periodo = (
historial_academico_filtrado
.group_by(["PERIODO", "CODIGO_ESTUDIANTE", "DESCRIPCION_NIVEL_MATERIA"])
.agg([
pl.sum("ES_INCOMPLETO").alias("TOTAL_INCOMPLETOS"),
pl.sum("MATERIA_CONTEO").alias("TOTAL_MATERIAS")
])
.with_columns([
(pl.col("TOTAL_INCOMPLETOS") == pl.col("TOTAL_MATERIAS")).cast(pl.Int8).alias("INCOMPLETO_TOTAL")
])
)
resumen_incompletos_acumulado = (
resumen_incompletos_por_periodo
.group_by(["CODIGO_ESTUDIANTE", "DESCRIPCION_NIVEL_MATERIA"])
.agg([
pl.sum("INCOMPLETO_TOTAL").alias("INCOMPLETO_TOTAL_ACUMULADO"),
])
)
resumen_creditos_acumulados = (
historial_academico_filtrado
.group_by(["CODIGO_ESTUDIANTE", "DESCRIPCION_NIVEL_MATERIA"])
.agg([
pl.sum("CREDITOS_APROBADOS").alias("CREDITOS_APROBADOS"),
pl.sum("CREDITOS_REPROBADOS").alias("CREDITOS_REPROBADOS"),
pl.sum("CREDITOS_INCOMPLETOS").alias("CREDITOS_INCOMPLETOS"),
pl.sum("CREDITOS_RETIRADOS").alias("CREDITOS_RETIRADOS"),
pl.sum("CREDITOS_PENDIENTES").alias("CREDITOS_PENDIENTES")
])
.with_columns([
(
pl.col("CREDITOS_APROBADOS") / (
pl.col("CREDITOS_APROBADOS") +
pl.col("CREDITOS_REPROBADOS") +
pl.col("CREDITOS_INCOMPLETOS") +
pl.col("CREDITOS_RETIRADOS") +
pl.col("CREDITOS_PENDIENTES")
)
).alias("PORCENTAJE_CREDITOS_APROBADOS")
])
)
semestres_matriculados = (
historial_academico_filtrado
.filter(~pl.col("PERIODO").is_in(CURRENT_PERIODS))
.group_by(["PERIODO", "CODIGO_ESTUDIANTE", "DESCRIPCION_NIVEL_MATERIA", "CODIGO_PROGRAMA_1"])
.agg([
pl.sum("HORAS_CREDITO").alias("TOTAL_HORAS_CREDITO")
])
.with_columns([
pl.when(
(pl.col("CODIGO_PROGRAMA_1") == "MEDI") & (pl.col("TOTAL_HORAS_CREDITO") > 12)
).then(1.0)
.when(
(pl.col("CODIGO_PROGRAMA_1") == "MEDI") & (pl.col("TOTAL_HORAS_CREDITO") > 7)
).then(0.5)
.when(
(pl.col("CODIGO_PROGRAMA_1") != "MEDI") & (pl.col("TOTAL_HORAS_CREDITO") > 9)
).then(1.0)
.when(
(pl.col("CODIGO_PROGRAMA_1") != "MEDI") & (pl.col("TOTAL_HORAS_CREDITO") > 4)
).then(0.5)
.otherwise(0.25)
.alias("SEMESTRE_MATRICULADO")
])
.group_by(["CODIGO_ESTUDIANTE", "DESCRIPCION_NIVEL_MATERIA"])
.agg([
pl.sum("SEMESTRE_MATRICULADO").alias("TOTAL_SEMESTRES_MATRICULADOS")
])
)
estudiantes_periodo_actual = (
student_performance
.filter(pl.col("PERIODO") == period)
)

student_risk_criteria_summary = (
resumen_riesgo_acumulado
.join(resumen_riesgo_ultimos_4, on="CODIGO_ESTUDIANTE", how="left")
.join(
resumen_incompletos_acumulado,
left_on=["CODIGO_ESTUDIANTE", "NIVEL_PROGRAMA_1"],
right_on=["CODIGO_ESTUDIANTE", "DESCRIPCION_NIVEL_MATERIA"],
how="left"
)
.join(
resumen_creditos_acumulados,
left_on=["CODIGO_ESTUDIANTE", "NIVEL_PROGRAMA_1"],
right_on=["CODIGO_ESTUDIANTE", "DESCRIPCION_NIVEL_MATERIA"],
how="left"
)
.join(
semestres_matriculados,
left_on=["CODIGO_ESTUDIANTE", "NIVEL_PROGRAMA_1"],
right_on=["CODIGO_ESTUDIANTE", "DESCRIPCION_NIVEL_MATERIA"],
how="left"
)
)
return student_risk_criteria_summary


And then this function is being used right here:

def compute_all_performance_indicators(grades_history: pl.LazyFrame,
periods_to_evaluate: list[int],
enrolled_students: pl.LazyFrame,
student_performance: pl.LazyFrame
) -> pl.LazyFrame:
undergrad_periods,postgrad_periods = periods_to_evaluate
undergrad_evaluated_periods = undergrad_periods[3:]
postgrad_evaluated_periods = postgrad_periods[2:]
resultados: list[pl.DataFrame] = []

for index_period, period in enumerate(undergrad_evaluated_periods):
criterios = (
calculate_performance_criteria(
grades_history,
enrolled_students,
student_performance,
period,
index_period,
undergrad_periods
)
.with_columns(pl.lit(period).alias("PERIODO_EVALUACION"))
.cache()
)
resultados.append(criterios)


criterios_unificados = pl.concat(resultados, how="vertical_relaxed")

estudiantes_en_periodos = (
student_performance
.filter(pl.col("PERIODO").is_in(undergrad_evaluated_periods))
)

resultado_final = estudiantes_en_periodos.join(
criterios_unificados,
left_on=["CODIGO_ESTUDIANTE", "NIVEL_PROGRAMA_1", "PERIODO"],
right_on=["CODIGO_ESTUDIANTE", "NIVEL_PROGRAMA_1", "PERIODO_EVALUACION"],
how="left"
)


return resultado_final

Continue reading...