Deploy a Python model (more efficiently) over Spark | by Schaun Wheeler | Towards Data Science

UPDATE: I wrote about another way of deploying Scikit-Learn models over Spark that avoids some of the problems of this approach.
I’m posting this so other people can hopefully avoid some of the pain I had to go through. There are times when I want to train a model in scikit-learn, but model outcomes for all records in a large dataset in Spark. Yes, I know I could use Spark MLlib, but I find scikit-learn to have a more robust set of offerings, I understand the code in scikit-learn better than I understand the code in MLlib, and I’m just plain-old more familiar with scikit-learn. However, calling a scikit-learn `predict` method through a PySpark UDF creates a couple problems:
It incurs the overhead of pickling and unpickling the model object for every record of the Spark dataframe.
It fails to take advantage of scikit-learn’s optimizations, which mostly are due to vectorizing function calls over NumPy arrays.
I’ve found I can mitigate some of this overhead by grouping your spark dataframe in order to call my Python object’s methods on multiple records at once. Here’s how you do it:
Set up three columns in your Spark data frame:
Set up three columns in your Spark data frame:
* A unique id. This can be anything. As long as it is unique, you’re good to go.
* All of your predictors. You can use a Structype or MLLib’s VectorAssembler to get all of your predictors into a single column.
* A groups column. You can call row_number() modulo’d by the number of groups you want. I find it generally works well to create enough groups that each group will have 50–100k records in it.
2. Call the Spark SQL function `create_map` to merge your unique id and predictor columns into a single column where each record is a key-value store.
3. Group by your groups column, and call the Spark SQL function `collect_list` on your key-value column. This will aggregate your data set into lists of dictionaries.
4. Broadcast your scikit-learn model.
5. Create a UDF that unpacks a list of dictionaries into a list of keys (your unique ids) and a list of lists (your predictors). You can then feed the list of lists directly into a broadcasted scikit-learn model’s `predict` method. Then zip the result of that function call with your key list and convert to a dictionary. The udf will return a MapType, with the keys and values types set appropriately depending on what format your keys take and what format you want to return from your scikit-learn function call.
6. Call explode on the results of your udf, and include two aliases — one for the keys, and one for the results. You’ll then have a new data frame, the same size as your original (pre-grouped) dataframe, with your results in one column, and keys in the other column that can be used to join the results with the original data.
Here’s an example:
"""assumes the following already exist within the environment:`model`: a scikit-learn class that predicts probabilities for a two-class (0.0, 1.0) model`sdf`: a spark dataframe with at least two columns: "unique_id" and "feature_list""""import pyspark.sql.functions as fimport pyspark.sql.types as timport pyspark.sql.window.Window as wfrom pyspark.context import SparkContextsc = SparkContext.getOrCreate()# broadcast modelmodel_broadcast = sc.broadcast(model)# udf to predict on the clusterdef predict_new(feature_map):    ids, features = zip(*[        (k,  v) for d in feature_map for k, v in d.items()    ])    ind = model_broadcast.value.classes_.tolist().index(1.0)    probs = [        float(v) for v in         model_broadcast.value.predict_proba(features)[:, ind]    ]    return dict(zip(ids, probs))predict_new_udf = f.udf(    predict_new,     t.MapType(t.LongType(), t.FloatType())# set the number of prediction groups to createnparts = 5000# put everything togetheroutcome_sdf = (    sdf    .select(        f.create_map(            f.col('unique_id'),                 f.col('feature_list')        ).alias('feature_map'),         (            f.row_number().over(                w.partitionBy(f.lit(1)).orderBy(f.lit(1))            ) % nparts        ).alias('grouper')    )    .groupby(f.col('grouper'))    .agg(        f.collect_list(f.col('feature_map')).alias('feature_map')    )    .select(        predict_new_udf(f.col('feature_map')).alias('results')    )    .select(        f.explode(f.col('results'))        .alias('unique_id', 'probability_estimate')    ))
I’d like to be able to show some graphs comparing performance with the above method to the performance of simply wrapping a call to the model’s predict method in a PySpark UDF, but I can’t: I figured this method out because I couldn’t get the naive method to finish. I had a dataset with over 100 million records. I had trained a Naive Bayes (using scikit-learn’s MultinomialNB) classifier to distinguish between two classes — 0 vs. 1 — based on a hashed term-document matrix (using scikit-learn’s HashingVectorizer). At first, I simply broadcast the trained model, and then wrote a UDF that took a pre-processed string as an argument. The UDF ran that string through HashingVectorizer, then fed those results into the model’s predict method. I then ran the script and monitored YARN. After two days, the process was about 10% complete. I then rewrote the process to follow the approach I’ve outlined here: the UDF took a list of strings as an argument, and that list was processed by HashingVectorizer and the results fed to the model’s predict method. The entire run completed in 30 minutes.
One last note: different scikit-learn models have vastly different footprints in memory. A Naive Bayes model only needs to keep a few values for each parameter. A random forest needs to keep every tree in the forest. Because you need to broadcast the model to each executor, you could easily find that a model trained on a lot of data requires a whole lot of memory.