MMRアルゴリズムをPysparkで実装する

レコメンド結果にMMRを適用して多様性を考慮したかった。
以下のようにPythonで実装されているコードはあったが、Pysparkで実装れているサンプルがなかったので実装してみた。

MMR実装を参考にしたページ

実装コード

モジュールインポート

from pyspark.sql import DataFrame as SDF
from typing import Set, Callable, List
import pyspark.sql.functions as F

検索クエリとアイテムリストの類似度結果を取得する関数

def sim_func(df: SDF, item_id: int, rec_item_id: int) -> float:
    min_score = df.filter(F.col('item_id') == item_id).groupBy().min('score').collect()[0][0]
    try:
        score = df.filter((F.col('item_id') == item_id) & (F.col('rec_item_id') == rec_item_id)).collect()[0]['score']
    except:
        score = min_score
   return score

選択済みアイテムがオリジナルのアイテムと同じ集合になるまでLOOP

def mmr(df: SDF, items: Set[int], item_id: int, lambda_: float, sim_func1: Callable[[SDF, int, int], float], sim_func2: Callable[[SDF, int, int]) -> List[int]:
    def _argmax(keys, f):
        return max(keys, key=f)

    selected = []
    while set(selected) != items:
        remaining = items - set(selected)
        mmr_score = lambda x: labmda_ * sim_func1(df, item_id, x) - (1 - lambda_) * max([sim_func2(df, x, y) for y in set(selected)-{x}] or [0])
        next_selected = _argmax(remaining, mmr_score)
        selected.append(next_selected)
    return selected

MMR実行

mmr(
    df, 
    set(list(df.filter(F.col('item_id') == 12345678).select('rec_item_id').toPandas()['rec_item_id'])),
    12345678,
    0.8,
    sim_func,
    sim_func
)