我认为您在达到预期产出方面处于正确的轨道上。
我设法找到解决此类问题的两种可能的解决方案:一种使用火花UDF,其他用途熊猫 UDF.
火花UDF
from pyspark.sql.functions import udf
@udf('integer')
def predict_udf(*cols):
return int(braodcast_model.value.predict((cols,)))
list_of_columns = df.columns
df_prediction = df.withColumn('prediction', predict_udf(*list_of_columns))
熊猫 UDF
import pandas as pd
from pyspark.sql.functions import pandas_udf
@pandas_udf('integer')
def predict_pandas_udf(*cols):
X = pd.concat(cols, axis=1)
return pd.Series(braodcast_model.value.predict(X))
list_of_columns = df.columns
df_prediction = df.withColumn('prediction', predict_pandas_udf(*list_of_columns))
可重现的例子
在这里,我使用了 Databricks Community 集群Spark 3.1.2
, pandas==1.2.4
and pyarrow==4.0.0
.
broadcasted_model
是来自 scikit-learn 的简单逻辑回归,经过训练乳腺癌数据集 https://scikit-learn.org/stable/modules/generated/sklearn.datasets.load_breast_cancer.html.
import pandas as pd
import joblib
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline
from pyspark.sql.functions import udf, pandas_udf
# load dataset
X, y = load_breast_cancer(return_X_y=True, as_frame=True)
# split in training and test set
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=28)
# create a small pipeline with standardization and model
pipe = make_pipeline(StandardScaler(), LogisticRegression())
# save and reload the model
path = '/databricks/driver/test_model.joblib'
joblib.dump(model, path)
loaded_model = joblib.load(path)
# sample of unseen data
df = spark.createDataFrame(X_test.sample(50, random_state=42))
# create broadcasted model
sc = spark.sparkContext
braodcast_model = sc.broadcast(loaded_model)
然后我使用了上面说明的两种方法,你会看到输出df_prediction
在两种情况下都是相同的。