Apply sklearn trained model on a dataframe with PySpark

user_5706548 0
Viewed 5503

I trained a random forest algorithm with Python and would like to apply it on a big dataset with PySpark.

I first loaded the trained sklearn RF model (with joblib), loaded my data that contains the features into a Spark dataframe and then I add a column with the predictions, with a user-defined function like that:

def predictClass(features):
    return rf.predict(features)
udfFunction = udf(predictClass, StringType())
new_dataframe = dataframe.withColumn('prediction', 
udfFunction('features'))

It takes so much time to run though, is there a more efficient way to do the same thing? (without using Spark ML)

scikit-learn pyspark apache-spark python
3 Answers

I had to do same thing in recent project. The bad thing about applying udf for each row that pyspark has to read sklearn model each time so that's why it takes ages to finish. The best solution I have found was to use .mapPartitions or foreachPartition method on rdd, really good explanation is here

https://github.com/mahmoudparsian/pyspark-tutorial/blob/master/tutorial/map-partitions/README.md

It works fast because it ensures you that there is no shuffling and for each partition pyspark has to read the model and predict only once. So, the flow would be:

  • convert DF to RDD
  • broadcast model to nodes so it will be accessible for workers
  • write an udf function which takes interator (which contains all rows within a partition) as an argument
  • iterate through rows and create a proper matrix with your features (order matters)
  • call .predict only once
  • return predictions
  • transform rdd to df if needed
user_72483970

Now you can also use pandas_udf introduced in spark 2.3 to achieve high processing speed and distributed computation. It is based on pyarrow a python implementation of Apache Arrow use for in-memory computation.

user_74372550
Related