PySpark에서 Python UDF 사용하기

일반 DataFrame 연산은 Spark가 실행할 계획을 만듭니다. Python UDF는 행의 값을 Python 함수에 전달하여 사용자 로직을 executor에서 실행하게 합니다. 복잡한 문자열 형식을 파싱하는 등 기본 함수로 표현하기 어려운 계산에 사용할 수 있습니다.

가능하면 기본 PySpark 함수를 우선 사용합니다. UDF에는 Python 실행과 데이터 전달 비용이 있고 Spark가 내부 Python 로직을 일반 열 식처럼 최적화할 수 없습니다.

한 열에서 값 추출하기

다음은 Weather report: rain 55-62처럼 끝에 음수가 없는 정수 범위가 있는 가상의 문자열 형식입니다. 마지막 범위의 앞 숫자를 낮은 온도로 읽습니다. NULL 입력은 NULL 결과로 유지하며, 잘못된 비어 있지 않은 형식의 파싱 오류는 전파합니다.

from pyspark.sql import functions as F
from pyspark.sql.types import IntegerType


def extract_low_temperature(weather_report):
    if weather_report is None:
        return None
    return int(weather_report.split(" ")[-1].split("-")[0])


low_temp_udf = F.udf(extract_low_temperature, IntegerType())


def add_low_temperature(df):
    return df.withColumn("low", low_temp_udf(F.col("weather_report")))

반환형은 Spark의 IntegerType으로 명시합니다. 반환하는 Python 값이 선언한 자료형과 맞아야 합니다.

UDF를 정의한 다음에는 기본 열 함수처럼 전체 Column을 인자로 전달합니다. 각 행의 값을 직접 Python 반복문으로 수집하는 호출과 구분합니다.

여러 열을 인자로 받기

UDF의 열 인자 순서는 Python 함수 인자 순서와 대응합니다. 다음 가상 규칙은 온도가 70을 초과하고 바람이 없으면 good, 그 외에는 bad입니다. 입력 온도나 바람 상태가 NULL이면 NULL 결과를 반환합니다.

from pyspark.sql import functions as F
from pyspark.sql.types import StringType


def weather_quality(temperature, windy):
    if temperature is None or windy is None:
        return None
    if temperature > 70 and windy is False:
        return "good"
    return "bad"


weather_udf = F.udf(weather_quality, StringType())


def add_weather_quality(df):
    return df.withColumn("quality", weather_udf(F.col("temp"), F.col("wind")))

여기서 temp는 숫자, wind는 불리언입니다. 이 정도 조건은 기본 함수로도 표현할 수 있으므로 실제 파이프라인에서는 다음 native 표현식을 선택할 수 있습니다.

from pyspark.sql import functions as F


def add_weather_quality_native(df):
    missing = F.col("temp").isNull() | F.col("wind").isNull()
    good = (F.col("temp") > 70) & (F.col("wind") == F.lit(False))
    return df.withColumn(
        "quality",
        F.when(missing, F.lit(None).cast("string"))
        .when(good, "good")
        .otherwise("bad"),
    )

실행 위치와 조사

UDF 내부에서 쓰는 로그는 쿼리를 구성한 driver의 최상위 로그와 같은 위치에 나타나지 않습니다. 반환값을 조사해야 한다면 임시 디버그 UDF가 설명 문자열을 새 열로 반환하게 할 수 있습니다. 로그와 실행 계획에서 driver 로그, executor 계산과 제한된 행 수집을 구분합니다.

UDF 안에서 사용하는 라이브러리와 보조 코드는 executor의 실행 환경에 있어야 합니다. 입력 NULL, 자료형과 파싱 실패를 함수의 계약으로 명시하고, 바깥 조건식이 UDF 호출을 항상 단락 평가한다고 가정하지 않습니다.