PysparkでUDFを書く

いつも忘れちゃうので。

UDFの定義の仕方

書き方は２通り。

udf関数に取り込む

lambda関数を書くときに便利。

from pyspark.sql.functions import udf, col
from pyspark.sql.types import StringType
from urllib.parse import urlparse

extract_domain = udf(lambda url: urlparse(url).netloc, StringType())

デコレータを使う

やりたい処理が長くなっちゃうときに便利。

@udf(StringType)
def extract_domain(url):
  uri = urlparse(url)
  domain = uri.netloc
  return domain

定義したUDFの使い方

関数として普通に使える。

df = df.withColumn('domain', extract_domain(col('url')))