Create a Spark ML

Pipeline

from pyspark.ml import Pipeline
from pyspark.ml.feature import Imputer, VectorAssembler, RobustScaler, StringIndexer, OneHotEncoder

# Imputer (mean strategy for all double/numeric)
to_impute = num_missing_cols
imputer = Imputer(inputCols=to_impute, outputCols=to_impute, strategy='mode')

# Scale numerical
numerical_assembler = VectorAssembler(inputCols=num_cols, outputCol="numerical_assembled")
numerical_scaler = RobustScaler(inputCol="numerical_assembled", outputCol="numerical_scaled")

# OHE categoricals
ohe_cols = [column + '_ohe' for column in string_cols]
one_hot_encoder = OneHotEncoder(inputCols=string_cols_indexed, outputCols=ohe_cols, handleInvalid="keep")

# Assembler 
feature_cols = ["numerical_scaled"] + ohe_cols
vector_assembler = VectorAssembler(inputCols=feature_cols, outputCol="features")

# Instantiate the pipeline
stages_list = [
	imputer,
	numerical_assembler,
	numerical_scaler,
	one_Hot_encoder,
	vector_assembler
]

pipeline = Pipeline(stages=stages_list)

Fit the Pipeline

<aside> 💡

This pipeline will index the columns we specified and finally create a vector from all input columns

</aside>

pipeline_model = pipeline.fit(train_df)

Transform

# Transform both training_df and test_df
train_transformed_df = pipeline_model.transform(train_df)
test_transformed_df = pipeline_model.transform(test_df)