One-hot encoding in Spark MLlib

<aside> 💡

To one-hot encode categorical features in Spark MLlib, we are going to use two classes: the StringIndexer class and the OneHotEncoder class.

• The StringIndexer class indexes string-type columns to a numerical index. Each unique value in the string-type column is mapped to a unique integer.
• The OneHotEncoder class accepts indexed columns and converts them to a one-hot encoded vector-type feature. </aside>

from pyspark.ml.feature import StringIndexer
from pyspark.sql.functions import col

# StringIndexer
string_cols = ["Contract"]
index_cols = [column + "_index" for column in string_cols]

string_indexer = StringIndexer(inputCols=string_cols, outputCols=index_cols, handleInvalid="skip")
string_indexer_model = string_indexer.fit(sample_df)
indexed_df = string_indexer_model.transform(sample_df)

indexed_df.show()

from pyspark.ml.feature import OneHotEncoder

# Create a list of one-hot encoded feature names
ohe_cols = [column + "_ohe" for column in string_cols]

# Instantiate the OneHotEncoder with the column lists
ohe = OneHotEncoder(inputCols=index_cols, outputCols=ohe_cols, handleInvalid="keep")

# Fit the OneHotEncoder on the indexed data
ohe_model = ohe.fit(indexed_df)

# Transform indexed_df using the ohe_model
ohe_df = ohe_model.transform(indexed_df)
ohe_df.show()

from pyspark.ml.feature import VectorAssembler

selected_ohe_cols = ["Contract_ohe"]

# Use VectorAssembler to assemble the selected one-hot encoded columns into a dense vector
assembler = VectorAssembler(inputCols=selected_ohe_cols, outputCol="features")
result_df_dense = assembler.transform(ohe_df)

# Select relevant columns for display
result_df_display = result_df_dense.select("Contract", "features")

result_df_display.show(truncate=False)