Guide to implementing Linear Regression in Pyspark and R | by Kieran Tan Kah Wang | Towards Data Science

Which programming tool should I use to build my Linear Regression model? Is there any differences in terms of the prediction accuracy? Which is simpler to implement with just a few lines of codes? Does the size of my training data plays a difference? Let’s explore all these questions in this article.
The dataset I’ve used is a dummy dataset which contains just two variables, YearsExperience and Salary. We are trying to find out if the salary of an employee is linearly dependent on the number of years of work experience the employee has, and logically, we will expect such a relationship to exists.
We will make use of Pyspark to train our Linear Regression model in Python as Pyspark has the ability to scale up data processing speed which is highly valued in the world of big data. As this article focused more on implementing Linear Regression model, I will not be touching on the technical aspects in setting up Pyspark (i.e. SparkContext, SparkSession, SparkConf) which you can find in my previous post.
from pyspark import SQLContext, SparkConf, SparkContextfrom pyspark.sql import SparkSessionsc = SparkContext.getOrCreate()if (sc is None):    sc = SparkContext(master="local[*]", appName="Linear Regression")spark = SparkSession(sparkContext=sc)sqlcontext = SQLContext(sc)data = sqlcontext.read.csv('./Salary_Data.csv', header = True, inferSchema = True)
Once we get the data into Pyspark, we will need to tell it which is our feature variables and which is our target variable. For this dataset, it’s quite straightforward that YearsExperience is our feature variable while Salary is our target variable. We have our data, and we are going to build our model, but how do we know if our model performs well or not? Splitting our dataset into training and testing data is one way to do it among data scientists, and the usual ratio is 70:30 or 80:20. Another thing to note in implementing machine learning algorithms in Pyspark is that we can make use of vector assembler, which combine all features into a sparse single vector that greatly reduce the training time for machine learning models.
from pyspark.ml.feature import VectorAssembler# defining Salary as our label/predictor variabledataset = data.select(data.YearsExperience, data.Salary.alias('label'))# split data into 70% training and 30% testing datatraining, test = dataset.randomSplit([0.7, 0.3], seed = 100)# assembler to assemble the features into vector formassembler = VectorAssembler().setInputCols(['YearsExperience',]).setOutputCol('features')trainingSet = assembler.transform(training)# select only features column and label column since we have already vectorised our featurestrainingSet = trainingSet.select("features","label")
Finally, we are ready to train a Linear Regression model using our training dataset. And as mentioned above, we will also make use of a testing dataset to know how well our model performed. As seen in the output below, our model now contains a prediction column which contains the predicted salary when we fit our feature column (i.e. YearsExperience) as testing dataset into our Linear Regression model. We can see that some predicted values are close to the actual values, while some are not.
from pyspark.ml.regression import LinearRegression# fit the training set to linear regression modellr = LinearRegression()lr_Model = lr.fit(trainingSet)# assembler to assemble the features into vector formtestSet = assembler.transform(test)# select only features column and label column since we have already vectorised our featurestestSet = testSet.select("features", "label")# fit the testing data into our linear regression modeltestSet = lr_Model.transform(testSet)testSet.show(truncate=False)
So how do we measure how well our Linear Regression model performed? We can make use of R2, which is a measure of the goodness of fit between a value of 0 to 100%. We can implement in Pyspark easily with just 2 lines of code.
from pyspark.ml.evaluation import RegressionEvaluatorevaluator = RegressionEvaluator()print(evaluator.evaluate(testSet, {evaluator.metricName: "r2"}))
A R2 value of 93.9% suggests that our Linear Regression has performed significantly well in terms of predicting Salary when we fit our trained model using YearsExperience.
Now, let’s build our Linear Regression model in R. We split the data into 70% training data and 30% testing data as what we have did in Pyspark. Whereas, let’s try to use the same testing data as we used in Pyspark to see if there’s any difference in R2 performance in the model’s predictions.
data = read.csv("./Salary_Data.csv")# sample data# set.seed(100)# dt = sort(sample(nrow(data), nrow(data)*.7))dt = c(5,8,9,10,16,19,20,22,26,27)# split data into training and testing datatrainingSet<-data[-dt,]testingSet<-data[dt,]
After obtaining the training and testing dataset, we train our Linear Regression model with the training dataset and predict the Salary using the testing dataset.
# fit linear regression modellm_model = lm(data = trainingSet, Salary~.)# predict the target variable with testing datapredicted_salary = predict(lm_model, testingSet)
Then, we measure the goodness of fit on our predicted values by calculating R2. Note that in R, there’s no libraries or commands we can run to directly get R2. Hence, we have to calculate them manually by obtaining the residual sum of squares and total sum of squares which can further explored here.
# residual sum of squaresrss <- sum((predicted_salary - testingSet$Salary) ^ 2) # total sum of squarestss <- sum((testingSet$Salary - mean(testingSet$Salary)) ^ 2) # obtain r^2rsq <- 1 - rss/tssrsq
We obtained the same R2 as we did in Pyspark, which means in terms of accuracy of models, choosing Pyspark or R to train your model does not matter. And it might seem obvious that training Linear Regression model in R is a lot more simpler than in Pyspark. Whereas, a point to note is that Pyspark is more preferable in terms of training big data and the processing speed will be significantly faster than in R itself.
Here’s my recommendation when it comes to deciding between R or Pyspark in building Linear Regression model:
Pyspark -> Suitable for processing large amount of data
R -> Suitable for easy implementation
Thanks for reading and I will look into several other Machine Learning algorithms over my next few posts, cheers!