spark grid search

spark_grid_search.py

from pyspark.ml.regression import LinearRegression
from pyspark.ml.tuning import CrossValidator, ParamGridBuilder
from pyspark.ml.evaluation import RegressionEvaluator
from pyspark.ml.feature import OneHotEncoder, StringIndexer, VectorAssembler, StandardScaler
from pyspark.ml.pipeline import Pipeline


indexers = [
    StringIndexer(
    inputCol=c, 
    outputCol=f'{c}_idx', handleInvalid='keep')
    for c in categorical_feat
]
encoders = [
    OneHotEncoder(
        inputCol=f'{c}_idx',
        outputCol=f'{c}_onehot',
    ) 
    for c in categorical_feat
]
num_assembler = VectorAssembler(
    inputCols=numeric_feat,
    outputCol='num_features',
)
scaler = StandardScaler(inputCol='num_features', outputCol='num_features_scaled')
assembler = VectorAssembler(
    inputCols=[f'{c}_onehot' for c in categorical_feat] + ['num_features_scaled'],
    outputCol='features',
)

lr = LinearRegression(standardization=False, maxIter=100)
pipeline = Pipeline(
    stages=indexers + encoders + [num_assembler, scaler, assembler, lr])
    
# this is our grid
grid = (
    ParamGridBuilder()
    .addGrid(lr.elasticNetParam, np.arange(0, 1.01, 0.01))
    .addGrid(lr.regParam, [0, 0.5, 1, 2])
    .build()
)

crossval = CrossValidator(estimator=pipeline,
                          estimatorParamMaps=grid,
                          evaluator=RegressionEvaluator(),
                          numFolds=3)