UDF

ユーザー定義関数 (UDF) を作成します。

構文

Python

import pyspark.sql.functions as sf

# As a decorator
@sf.udf
def function_name(col):
    # function body
    pass

# As a decorator with return type
@sf.udf(returnType=<returnType>, useArrow=<useArrow>)
def function_name(col):
    # function body
    pass

# As a function wrapper
sf.udf(f=<function>, returnType=<returnType>, useArrow=<useArrow>)

パラメーター

パラメーター	Type	説明
f	function	オプション。スタンドアロン関数として使用される場合の Python 関数。
returnType	pyspark.sql.types.DataType または str	オプション。ユーザー定義関数の戻り値の型。値は、DataType オブジェクトまたは DDL 形式の型文字列のいずれかになります。デフォルトは StringType です。
useArrow	bool	オプション。Arrow を使用して (デ) シリアル化を最適化するかどうか。Noneの場合、Spark構成「spark.sql.execution.pythonUDF.arrow.enabled」有効になります。

例

例 1 : ラムダ、デコレータ、戻り値の型を持つデコレータを使用して UDF を作成する。

Python

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf

slen = udf(lambda s: len(s), IntegerType())

@udf
def to_upper(s):
    if s is not None:
        return s.upper()

@udf(returnType=IntegerType())
def add_one(x):
    if x is not None:
        return x + 1

df = spark.createDataFrame([(1, "John Doe", 21)], ("id", "name", "age"))
df.select(slen("name").alias("slen(name)"), to_upper("name"), add_one("age")).show()

Output

+----------+--------------+------------+
|slen(name)|to_upper(name)|add_one(age)|
+----------+--------------+------------+
|         8|      JOHN DOE|          22|
+----------+--------------+------------+

例 2 : キーワード引数を持つ UDF。

Python

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf, col

@udf(returnType=IntegerType())
def calc(a, b):
    return a + 10 * b

spark.range(2).select(calc(b=col("id") * 10, a=col("id"))).show()

Output

+-----------------------------+
|calc(b => (id * 10), a => id)|
+-----------------------------+
|                            0|
|                          101|
+-----------------------------+

例 3 : Pandasシリーズ型ヒントを使用してベクトル化されたUDF 。

Python

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf, col, PandasUDFType
import pandas as pd

@udf(returnType=IntegerType())
def pd_calc(a: pd.Series, b: pd.Series) -> pd.Series:
    return a + 10 * b

pd_calc.evalType == PandasUDFType.SCALAR
spark.range(2).select(pd_calc(b=col("id") * 10, a="id")).show()

Output

+--------------------------------+
|pd_calc(b => (id * 10), a => id)|
+--------------------------------+
|                               0|
|                             101|
+--------------------------------+

例 4 : PyArrow 配列型ヒントを使用してベクトル化された UDF。

Python

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf, col, ArrowUDFType
import pyarrow as pa

@udf(returnType=IntegerType())
def pa_calc(a: pa.Array, b: pa.Array) -> pa.Array:
    return pa.compute.add(a, pa.compute.multiply(b, 10))

pa_calc.evalType == ArrowUDFType.SCALAR
spark.range(2).select(pa_calc(b=col("id") * 10, a="id")).show()

Output

+--------------------------------+
|pa_calc(b => (id * 10), a => id)|
+--------------------------------+
|                               0|
|                             101|
+--------------------------------+

例 5 : Arrow に最適化された Python UDF (Spark 4.2 以降のデフォルト)。

Python

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf

# Arrow optimization is enabled by default since Spark 4.2
@udf(returnType=IntegerType())
def my_udf(x):
    return x + 1

# To explicitly disable Arrow optimization and use pickle-based serialization:
@udf(returnType=IntegerType(), useArrow=False)
def legacy_udf(x):
    return x + 1

例 6 : 非決定論的 UDF の作成。

Python

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf
import random

random_udf = udf(lambda: int(random.random() * 100), IntegerType()).asNondeterministic()