PySpark Random Forest – Building and Evaluating Random Forest Models using PySpark MLlib: A Step-By-Step Guide

Lets discuss how to build and evaluate Random Forest models using PySpark MLlib and cover key aspects such as hyperparameter tuning and variable selection, providing example code to help you along the way.

Random Forest is an ensemble machine learning algorithm that can be used for both classification and regression tasks.

PySpark is the Python library for Apache Spark, an open-source big data processing framework that can process large-scale data in parallel. PySpark MLlib is a machine learning library built on top of PySpark that provides various algorithms and tools for building scalable machine learning models.

We will cover the following topics in this post:

  1. Setting up the environment

  2. Loading and preprocessing the data

  3. Bilding a Random Forest model

  4. Hyperparameter tuning

  5. Evaluating the model

  6. Example code

1. Import required libraries and initialize SparkSession

First, let’s import the necessary libraries and create a SparkSession, the entry point to use PySpark.

import findspark

from pyspark import SparkFiles
from pyspark.sql import SparkSession
from import Pipeline
from import StringIndexer, VectorAssembler
from import RandomForestClassifier
from import MulticlassClassificationEvaluator
from import CrossValidator, ParamGridBuilder

spark = SparkSession.builder.appName("RandomForestExample").getOrCreate()

2. Load the dataset

For this example, we will use the “Iris” dataset. Save the dataset as a CSV file, and then use the following code to load the data into a PySpark DataFrame.

url = ""

df ="Iris.csv"), header=True, inferSchema=True)
| Id|SepalLengthCm|SepalWidthCm|PetalLengthCm|PetalWidthCm|    Species|
|  1|          5.1|         3.5|          1.4|         0.2|Iris-setosa|
|  2|          4.9|         3.0|          1.4|         0.2|Iris-setosa|
|  3|          4.7|         3.2|          1.3|         0.2|Iris-setosa|
|  4|          4.6|         3.1|          1.5|         0.2|Iris-setosa|
|  5|          5.0|         3.6|          1.4|         0.2|Iris-setosa|
only showing top 5 rows

3. Prepare the data

Before building the model, we need to assemble the input features into a single feature vector using the VectorAssembler class. Then, we will split the dataset into a training set (80%) and a testing set (20%).

# Preprocessing: StringIndexer for categorical labels
stringIndexer  = StringIndexer(inputCol="Species", outputCol="label")

# Define the feature and label columns & Assemble the feature vector
assembler = VectorAssembler(inputCols=["SepalLengthCm", "SepalWidthCm", "PetalLengthCm", "PetalWidthCm"], outputCol="features")

rf = RandomForestClassifier(labelCol="label", featuresCol="features")

# Split the data into training and test sets
train_data, test_data = df.randomSplit([0.7, 0.3], seed=42)

4. Create a Pipeline

We will now create a pipeline that includes the feature assembler and the Random Forest classifier.

pipeline = Pipeline(stages=[stringIndexer, assembler, rf])

5. Hyperparameter Tuning and Model Selection

We will use cross-validation to select the best model based on hyperparameter tuning. We will tune the numTrees and maxDepth parameters of the Random Forest model.

# Define the hyperparameter grid
paramGrid = ParamGridBuilder() \
    .addGrid(rf.numTrees, [10, 20, 30]) \
    .addGrid(rf.maxDepth, [5, 10, 15]) \

# Create the cross-validator
cross_validator = CrossValidator(estimator=pipeline,
                          evaluator=MulticlassClassificationEvaluator(labelCol="label", metricName="accuracy"),
                          numFolds=5, seed=42)

# Train the model with the best hyperparameters
cv_model =

6. Analyze feature importance

To understand the importance of each variable in the model, we can examine the featureImportances attribute of the best Random Forest model obtained from cross-validation.

best_rf_model = cv_model.bestModel.stages[-1]
importances = best_rf_model.featureImportances
feature_list = ["SepalLengthCm", "SepalWidthCm", "PetalLengthCm", "PetalWidthCm"]

print("Feature Importances:")
for feature, importance in zip(feature_list, importances):
    print(f"{feature}: {importance:.4f}")
Feature Importances:
SepalLengthCm: 0.0887
SepalWidthCm: 0.0590
PetalLengthCm: 0.3873
PetalWidthCm: 0.4650

This will display the importance of each feature in the best Random Forest model. Features with higher importance values contribute more to the model’s decision-making process.

7. Evaluating the model

To evaluate the performance of our Lasso Regression model, we’ll use the RegressionEvaluator class from PySpark MLlib.

# Make predictions on the test data
predictions = cv_model.transform(test_data)

evaluator = MulticlassClassificationEvaluator(labelCol="label", metricName="accuracy")

# Evaluate the model
accuracy = evaluator.evaluate(predictions)
print("Test set accuracy = {:.2f}".format(accuracy))
Test set accuracy = 0.93

8. Save and load the model (optional)

If you want to reuse the model in the future, you can save it to disk and load it back when needed.

# Save the model"rf_model")

# Load the model
from import RandomForestClassificationModel
loaded_model = RandomForestClassificationModel.load("rf_model")


we have demonstrated how to build and evaluate a Random Forest model using PySpark MLlib. We covered important aspects such as hyperparameter tuning, variable selection, and model evaluation.

With this knowledge, you can now apply the Random Forest algorithm to your own datasets using PySpark and gain valuable insights from your data.

Course Preview

Machine Learning A-Z™: Hands-On Python & R In Data Science

Free Sample Videos:

Machine Learning A-Z™: Hands-On Python & R In Data Science

Machine Learning A-Z™: Hands-On Python & R In Data Science

Machine Learning A-Z™: Hands-On Python & R In Data Science

Machine Learning A-Z™: Hands-On Python & R In Data Science

Machine Learning A-Z™: Hands-On Python & R In Data Science