Pyspark : Model Loading multiple times with udf

Question

Trying to apply udf on a large csv file that makes a model prediction based on some conditions, but for some reason the model is being loaded multiple number of times. Below is a sample snippet of how the flow looks like:

# main.py loads predict.py

from predict import predict_udf

data = spark.read("csv_path")
data.show()

| Column1  | Column2 |
| -------- | ------- |
|   Class1 |         |
|   Class2 |         |

data.withColumn("Column2",predict_udf(col("Column1"))

# predict.py

model = load_model() # Initialising model object

def predict(class_name):
    if class_name == "Class1":
        # Do something

    elif class_name == "Class2":
        # Do something else

predict_udf = udf(predict, StringType())

Ideally I was expecting the model to be loaded once since it is defined as a global variable, but the model is being loaded multiple number of times.

user_16919760 · Answer

You can try broadcasting the model to avoid unnecessary repetitive reads.

model = load_model() # Initialising model object
model_bc = spark.sparkContext.broadcast(model)

then you will be able to access initial model in udf via model_bc.value attribute.

This is a good example that illustrates complete solution.

Pyspark : Model Loading multiple times with udf

1 Answers