Использование модели sklearn в pyspark для прогнозирования

#python #pyspark #scikit-learn

#python #pyspark #scikit-learn

Вопрос:

У меня есть фрейм данных pyspark и обученная модель sklearn. Я хочу использовать его для прогнозирования с использованием метода mapPartitions из pyspark.

Как я могу написать оболочку вокруг моей модели для прогнозирования?

 
X_test = spark.read.csv('path') # Dataframe with columns for testing model

with open('/Users/skv/Desktop/model.pickle', 'rb') as file:
    model = pickle.load(file)
model_bc = sc.broadcast(model)

# I wrote function, but it predicts row-by-row

def predict_wrapper():
    def predict(rows):
            yield model_bc.value.predict(rows)[0]
    return predict

X_test_small.rdd.mapPartitions(predict_wrapper()).collect()

 

Что я должен изменить для прогнозирования один раз для каждого раздела?