databricksでsklearnの学習済みモデルをspark DataFrameに適用する方法が、意外と検索しても出てこなかったので備忘録。サンプルコードはmlflowからsklearnの学習済みモデルを読み込む想定なので、再利用される場合は適宜pickle.load()なりに変更してください。
pandas_udf
Spark 2.3以降であればpandas_udfが使えます。行ごとに展開するudfよりもpandas_udfの方が高速らしいです。
Spark 3.0以降ではpandas_udfで関数を定義する際、型ヒントの記述が必須。
from pyspark.sql.functions import pandas_udf
import pandas as pd
import mlflow
@pandas_udf('double')
def predict(*cols: pd.Series) -> pd.Series:
model_path = f'runs://{run.info.run_id}/model'
model = mlflow.sklearn.load_model(model_path)
X = pd.concat(cols, axis=1)
return pd.Series(model.predict_proba(X)[:, 1])
prediction_df = spark_df.withColumn('pred', predict(spark_df.columns))
@pandas_udf('double')でリターン値のデータ型を指定しています。モデルによって予測値のデータ型に合わせる必要があります。
pandas_udf with iterator
pandas_udfは指定されたバッチサイズごとにループ処理をしています。この方法だとループする度にモデルの読み込みが発生するため、ファイルサイズが大きいモデルの場合はオーバーヘッドが高くなります。mlflow.sklearn.load_model()を関数の外ですれば良いだけなんですが、それはそれでイケテナイです。
pandas_udfをイテレータで定義すると、モデルの読み込みが１度で済みます。ただしSpark 3.0以降でのみ対応。
from typing import Iterator, Tuple
@pandas_udf('double')
def predict(iterator: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]:
model_path = f'runs://{run.info.run_id}/model'
model = mlflow.sklearn.load_model(model_path)
for features in iterator:
pdf = pd.concat([features], axis = 1)
yield pd.Series(model.predict(pdf))
prediction_df = spark_df.withColumn('pred', predict(spark_df.columns))
なお、pandas_udfのバッチサイズはspark.conf.get(‘spark.sql.execution.arrow.maxRecordsPerBatch’)で確認可能です。。デフォルトは10,000。
mapInPandas
Spark 3.0以降に登場した、pandas DataFrameを入力値にした関数をそのままspark DataFrameに適用するメソッド。ただしリターン値となるspark DataFrameのスキーマを引数として渡す必要があるため、若干面倒。spark(静的言語のJAVA)にそのまま渡すイメージっぽい（あくまで個人の主観です）。
from typing import Iterator, Tuple
from pyspark.sql.functions import lit
from pyspark.sql.types import DoubleType
@pandas_udf('double')
def predict(iterator: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]:
model_path = f'runs://{run.info.run_id}/model'
model = mlflow.sklearn.load_model(model_path)
for features in iterator:
yield pd.concat([features, pd.Series(model.predict(features), name = 'prediction'))], axis = 1)
# 予測値のカラムをデータ値をNoneとして定義し、そのスキーマ情報を取得
schema = spark_df.withColumn('prediction', lit(None).cast(DoubleType())).schema
spark_df = spark_df.mapInPandas(predict, schema)