Posterior a consultar especialistas en la materia, programadores de diversas ramas y profesores dimos con la solución al dilema y la plasmamos en esta publicación.
Solución:
Tuve que hacer lo mismo en un proyecto reciente. Lo malo de aplicar udf para cada fila que pyspark tiene que leer el modelo sklearn cada vez, es por eso que tarda años en terminar. La mejor solución que encontré fue usar el método .mapPartitions o foreachPartition en rdd, aquí hay una muy buena explicación.
https://github.com/mahmoudparsian/pyspark-tutorial/blob/master/tutorial/map-partitions/README.md
Funciona rápido porque le asegura que no hay mezcla y para cada partición, pyspark tiene que leer el modelo y predecir solo una vez. Entonces, el flujo sería:
- convertir DF a RDD
- transmitir el modelo a los nodos para que sea accesible para los trabajadores
- escriba una función udf que tome interator (que contiene todas las filas dentro de una partición) como argumento
- iterar a través de filas y crear una matriz adecuada con sus características (el orden es importante)
- llama a .predict solo una vez
- predicciones de retorno
- transformar rdd a df si es necesario
Recuerda que puedes compartir este artículo si te fue de ayuda.