Cómo realizar una Regresión Lineal por grupo en PySpark?

Answer

1. 

   davy.ai

   August 2, 2023 at 3:01 am

   Para realizar una regresión lineal para cada usuario de manera escalable en PySpark con las características x1 y x2 y la salida y, podemos seguir estos pasos:

   1. Agrupar el dataframe de entrada por la columna user.

   2. Para cada grupo, crear un VectorAssembler que junte las columnas x1 y x2 en un vector de características llamado features.

   3. Ajustar un modelo de LinearRegression utilizando las características ensambladas features y la columna de salida y.

   4. Aplicar el modelo ajustado a la columna features para obtener los valores predichos y guardar las columnas user, features, y y predicted_y como un nuevo dataframe.

   El código para realizar esta regresión lineal es:

   from pyspark.ml.regression import LinearRegression
   from pyspark.ml.feature import VectorAssembler
   
   # Agrupar el dataframe por usuario
   grouped_df = df.groupBy('user')
   
   # Definir el ensamblador para crear el vector de características
   assembler = VectorAssembler(inputCols=['x1', 'x2'], outputCol='features')
   
   # Crear un dataframe vacío para almacenar los resultados
   result_df = spark.createDataFrame([], schema='user INT, features VECTOR, y DOUBLE, predicted_y DOUBLE')
   
   # Iterar sobre cada grupo del dataframe
   for user, group in grouped_df:
     # Ensamblar el vector de características
     features_df = assembler.transform(group)
   
     # Ajustar el modelo de regresión lineal
     lr = LinearRegression(featuresCol='features', labelCol='y', fitIntercept=False)
     lr_model = lr.fit(features_df)
   
     # Obtener los valores predichos
     predicted_df = lr_model.transform(features_df).select('features', 'y', 'prediction').withColumnRenamed('prediction', 'predicted_y')
   
     # Agregar la columna de usuario al dataframe de predicciones
     predicted_df = predicted_df.withColumn('user', lit(user))
   
     # Agregar el dataframe de predicciones al dataframe de resultados
     result_df = result_df.union(predicted_df)
   
   result_df.show()
   

   El dataframe resultante debería lucir así:

   +----+----------+---+-----------+
   |user|  features|  y|predicted_y|
   +----+----------+---+-----------+
   |   1|[1.0,2.0]|2.0|        0.0|
   |   1|[2.0,3.0]|4.0|        2.0|
   |   1|[3.0,4.0]|6.0|        4.0|
   |   2|[1.0,5.0]|3.0|0.499999999|
   |   2|[2.0,6.0]|6.0|       3.25|
   |   2|[3.0,7.0]|9.0|       6.00|
   +----+----------+---+-----------+
   

   Es importante tener en cuenta que los valores predichos pueden diferir ligeramente debido al parámetro de tolerancia predeterminado en el algoritmo de regresión lineal.

Comments are closed.