I am using Spark MLLib to make prediction and I would like to know if it is possible to create your custom Estimators.

Here is a reproducible of what I would like my model to do with the Spark api

from sklearn.datasets import load_diabetes
import pandas as  pd
import pyspark 
from pyspark.ml.feature import VectorAssembler, SQLTransformer
from pyspark.ml.classification import LogisticRegression
from pyspark.ml import Pipeline
# Query diabetes data 
diab = load_diabetes()
df = pd.DataFrame(diab.data, columns=diab.feature_names)
df['is_male'] = df.sex > 0
df.drop('sex', inplace=True, axis=1)
df['label'] = diab.target
# Model made with Spark

spark = pyspark.sql.SparkSession.builder.master('local[10]').appName('A random spark context').getOrCreate()
def create_gender_model_male(): 
    return Pipeline(stages=[SQLTransformer(statement='SELECT * FROM __THIS__ WHERE is_male'),
                                VectorAssembler(inputCols=['age', 'bmi', 'bp', 's1'],outputCol='features'),
                                LogisticRegression(featuresCol='features', labelCol='label', maxIter=100,
                                                   elasticNetParam=1, regParam=0.)
def create_gender_model_female(): 
    return Pipeline(stages=[SQLTransformer(statement='SELECT * FROM __THIS__ WHERE not is_male'),
                                VectorAssembler(inputCols=['age', 'bmi', 'bp', 's1'],outputCol='features'),
                                LogisticRegression(featuresCol='features', labelCol='label', maxIter=100,
                                                   elasticNetParam=1, regParam=0.)
df = spark.createDataFrame(df)
class MixedModel():
    def __init__(self):
        self.models = {'male': create_gender_model_male(), 'female': create_gender_model_female()}
        self.fitted_models = {'male': None, 'female': None}
    def fit(self, df): 
        self.fitted_models['male'] = self.models['male'].fit(df)
        self.fitted_models['female'] = self.models['female'].fit(df)
    def predict(self, df): 
        return self.fitted_models['male'].transform(df).union(self.fitted_models['female'].transform(df))
mm = MixedModel()

Here, for example I have one logistic regression per sex but I would also like to be able to have prediction with a tree for males and prediction with Logistic regression for females if I want.

In a perfect world there would be a function:

ModelAggregation(('is_male is true, male, model_for_male), ('is_male is false', model_for_female)))

which would return me an object like my model aggregation

  • $\begingroup$ is this "Ensemble Model" you are looking for? i.stack.imgur.com/ksEkK.png $\endgroup$ Apr 17, 2019 at 21:58
  • $\begingroup$ I am not sure I understand your drawing, is it about one hot encoding. It would be about having one model per Country, sex pair and then mix all those models in a common interface which might enable to take advantage of the parallelization capabilities of Spark $\endgroup$ Apr 17, 2019 at 22:08
  • $\begingroup$ I am trying to understand what you are trying to do, this image is what I got from your question. I think I got it wrong hahaha $\endgroup$ Apr 17, 2019 at 22:20
  • 3
    $\begingroup$ with MLlib you could just one-hot-encode the categorical variables and throw them all in one model?! but that seems too easy? am I missing something? $\endgroup$
    – oW_
    Apr 17, 2019 at 22:53
  • 1
    $\begingroup$ I updated my question with a reproducible example $\endgroup$ Apr 18, 2019 at 10:36

2 Answers 2


As you said in your question, there is no way to do that with the baseline algorithms provided in MLLib.

Two ways you could do that is by either :

  1. Creating a function to generate a Pipeline
  2. Creating a Meta Estimator that would take your base learners and the forking column.

The first one is what you have stated in your question and the second one has been explained by Brian Spiering.

As you said making a custom Estimator and Transformer will make it works nicely with all the MLLib Tools as Tuning Tools.

If you want a really precise example how to do so, there is my library which implements meta algorithms for Ensemble Learning with Spark.

There is a way to train multiple estimators in parallel quite simply in the Stacking Classifier, and for all of them you have an example on how to ask for an estimator as the parameter (to chose either your logistic regression or decision tree).

You can pick ideas out of that !

Edit : Similar question on StackOverflow


You could just pass gender as a parameter to create a pipeline

def create_pipeline(gender): 

    select_statement = "SELECT * FROM __THIS__ WHERE {predicate}".format(predicate = "is_male" if gender == "male" else "not is_male") 

    pipeline = Pipeline(stages=[transformer = SQLTransformer(statement=select_statement),
                                VectorAssembler(inputCols=['age', 'bmi', 'bp', 's1'],outputCol='features'),
                                LogisticRegression(featuresCol='features', labelCol='label', maxIter=100,
                                                   elasticNetParam=1, regParam=0.)
    return pipeline

Then store all the models

model_types = {'model_male': create_model(gender='male'),
               'model_not_male': create_model(gender='not_male')}

Fit each model

for model_types in models:
    pipeline = models[model_types]
    model = pipeline.fit(df)
  • $\begingroup$ That's what my MixedModel does, but as your example, it is running two fit methods one after the other. I was wondering if those two fitting could be parallelized. $\endgroup$ Apr 18, 2019 at 16:07
  • $\begingroup$ Yes - submit each model fitting as different threads stackoverflow.com/questions/49568940/… $\endgroup$ Apr 18, 2019 at 16:45
  • $\begingroup$ Thanks, that's what I am looking for, would it be possible to wrap all of that in a model object which we could use with the cross validator provided by SparkMLlib ? $\endgroup$ Apr 19, 2019 at 17:21
  • $\begingroup$ MLlib's CrossValidator has a parallelism parameter which allows multiple models to be trained at the same time spark.apache.org/docs/2.3.0/ml-tuning.html#cross-validation $\endgroup$ Apr 20, 2019 at 6:31
  • $\begingroup$ Thanks, what about using pipeline to turn the aggregate of the models in an actual estimator as you woud do with pipeline ? $\endgroup$ Apr 21, 2019 at 22:08

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.