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()
mm.fit(df)
mm.transform(df)
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