UDFとして関数を登録する

Python

def squared(s):
  return s * s
spark.udf.register("squaredWithPython", squared)

オプションとしてUDFの戻り値の型を設定することができます。デフォルトの戻り値はStringTypeとなります。

Python

from pyspark.sql.types import LongType
def squared_typed(s):
  return s * s
spark.udf.register("squaredWithPython", squared_typed, LongType())

Spark SQLでUDFを呼び出す

Python

spark.range(1, 20).createOrReplaceTempView("test")

SQL

%sql select id, squaredWithPython(id) as id_squared from test

データフレームでUDFを使う

Python

from pyspark.sql.functions import udf
from pyspark.sql.types import LongType
squared_udf = udf(squared, LongType())
df = spark.table("test")
display(df.select("id", squared_udf("id").alias("id_squared")))

あるいは、アノテーション構文を用いて同じUDFを宣言することができます。

Python

from pyspark.sql.functions import udf
@udf("long")
def squared_udf(s):
  return s * s
df = spark.table("test")
display(df.select("id", squared_udf("id").alias("id_squared")))

評価順序とnullチェック

Spark SQL(SQL、データフレーム、データセットAPIを含む)は、サブ表現の評価順序を保証しません。特に、オペレーターや関数の入力は必ずしも左から右、あるいは固定の順序で評価されません。例えば、論理演算のANDやOR表現には、左から右の「短絡」のセマンティクスはありません。

そのため、WHEREやHAVING句のようなブール演算の評価順序はクエリー最適化やクエリー計画の間で並び替えられる可能性があるためこの順序や副作用に依存するのは危険です。特にUDFがnullチェックで短絡セマンティクスに依存している場合、以下のようにUDFを呼び出す前にnullチェックが行われることは保証されません。

Python

spark.udf.register("strlen", lambda s: len(s), "int")
spark.sql("select s from test1 where s is not null and strlen(s) > 1") # no guarantee

このWHERE句はnullを除外する前にstrlenUDFが呼び出されることを保証しません。

適切にnullチェックを行うためには、以下のいずれかを行うことをお勧めします。

UDF自身でnullに対応するようにし、UDF内でnullチェックを行うようにします。
nullチェックにIFやCASE WHEN表現を使用し、条件分岐でUDFを呼び出します。

Python

spark.udf.register("strlen_nullsafe", lambda s: len(s) if not s is None else -1, "int")
spark.sql("select s from test1 where s is not null and strlen_nullsafe(s) > 1") // ok
spark.sql("select s from test1 where if(s is not null, strlen(s), null) > 1")   // ok

Databricks 無料トライアル

Databricks 無料トライアル