How to perform a Linear Regression by group in PySpark? - machine-learning

Question

The goal is to perform linear regression for each user in a scalable way in PySpark. Features: x1 and x2. Output: y
Regression equation (zero intercept): y = m(x1) + n(x2)
Example:
pdf = pd.DataFrame(
{
"user": [1, 1, 1, 2, 2, 2],
"x1": [1, 2, 3, 1, 2, 3],
"x2": [2, 3, 4, 5, 6, 7],
"y": [2, 4, 6, 3, 6, 9],
}
)
df = sc.createDataFrame(pdf)
df.show()
Data looks like:
+----+---+---+---+
|user| x1| x2| y|
+----+---+---+---+
| 1| 1| 2| 2|
| 1| 2| 3| 4|
| 1| 3| 4| 6|
| 2| 1| 5| 3|
| 2| 2| 6| 6|
| 2| 3| 7| 9|
+----+---+---+---+

Answer 1

I have used PandasUDF, that works for my use case. Ben Webber explains it in his post
from his post we can take the following approach: (BDR:6.4;Spark:2.4.5)
from pyspark.sql.functions import pandas_udf, PandasUDFType
from pyspark.sql.types import *
import statsmodels.api as sm
import pandas as pd
pdf = pd.DataFrame(
{
"user": [1, 1, 1, 2, 2, 2],
"x1": [1, 2, 3, 1, 2, 3],
"x2": [2, 3, 4, 5, 6, 7],
"y": [2, 4, 6, 3, 6, 9],
}
)
df = spark.createDataFrame(pdf)
schema = StructType([StructField('user', DoubleType(), True),
StructField('r_squared', DoubleType(), True)])
#pandas_udf(schema, PandasUDFType.GROUPED_MAP)
def train_LR(input_pd):
usr = input_pd.iloc[0]['user']
# Implement lineal regression, as per your needs
model = sm.OLS(input_pd['y'], input_pd[['x1','x2']]).fit()
R_sq = model.rsquared
return pd.DataFrame({'user': usr, 'r_squared': R_sq }, index=[0])
results = df.groupby('user').apply(train_LR)
display(results)
Note that the UDF is getting a pandas dataframe with only the data for the groupby