PySpark UDF가 느린 이유와 Pandas UDF — 10배 빠르게 만들기

번역가를 한 명 고용했는데, 그 번역가가 단어 하나를 옮길 때마다 두 방 사이를 직접 뛰어다녀야 한다고 상상해 보세요. 단어 한 개, 왕복 한 번. 백만 단어라면 백만 번의 왕복입니다. 일반 Python UDF 가 JVM 과 Python 사이를 오가며 하는 일이 정확히 이렇습니다. PySpark 잡을 프로파일링하면, 전체 시간의 80%를 UDF 하나가 잡아먹는 경우가 흔합니다. "파이썬 함수 하나 적용했을 뿐인데" 잡이 10배 느려지는 이유가 여기 있습니다 — 일반 Python UDF 는 행 단위로, JVM 과 Python 사이를 직렬화하며 동작합니다.

이 글은 Python UDF 가 왜 느린지 구조적으로 설명하고, Arrow 기반 벡터화된 Pandas UDF 로 성능을 끌어올리는 법, 그리고 가장 좋은 해법인 "UDF 자체를 피하기"를 정리합니다.

이 글에서 배우는 것

• 일반 Python UDF 가 느린 근본 원인(직렬화·행 단위 왕복)
• 내장 함수로 UDF 를 아예 없애는 1순위 전략
• Apache Arrow 기반 Pandas UDF 로 배치 벡터화하는 방법
• 그룹별 처리(applyInPandas)와 파티션 스트리밍(mapInPandas) 패턴
• UDF 를 쓸 수밖에 없을 때 주의해야 할 함정들

1. 왜 Python UDF 는 느린가

Spark 는 JVM 위에서 동작합니다. Python UDF 를 만나는 순간, Spark 는 하던 일을 멈추고 데이터를 JVM → Python → JVM 으로 왕복시켜야 합니다. 이 왕복이 행 하나마다 반복된다는 것이 핵심 문제입니다.

세 가지 비용이 겹칩니다.

내장 함수(F.col, F.when, …)는 JVM 안에서 컬럼 단위로 실행되어 이 비용이 전혀 없습니다. 그래서 첫 번째 원칙은 "내장 함수로 되면 UDF 쓰지 말 것" 입니다.

2. 1순위 해법 — 내장 함수로 대체

UDF 를 쓰기 전에 먼저 자문해 보세요: "이거 내장 함수로 안 될까?" 여러분이 생각하는 것보다 훨씬 많은 UDF 가 사실 내장 함수 조합으로 깔끔하게 표현됩니다.

from pyspark.sql import functions as F
 
# BAD: Python UDF
@F.udf("string")
def grade(score):
    if score >= 90: return "A"
    elif score >= 80: return "B"
    else: return "C"
df = df.withColumn("g", grade("score"))
 
# GOOD: 내장 when/otherwise (JVM 내 실행, 수배~수십배 빠름)
df = df.withColumn("g",
    F.when(F.col("score") >= 90, "A")
     .when(F.col("score") >= 80, "B")
     .otherwise("C"))

문자열·날짜·정규식·JSON 가공은 거의 다 내장 함수가 있습니다(regexp_replace, split, from_json, date_format, transform, filter 등). UDF 를 쓰기 전에 내장 함수를 먼저 찾아보세요 — 대부분 이미 있거든요.

한 문장으로: 내장 함수는 JVM 안에서 컬럼 단위로 실행되어, Python 왕복 비용이 아예 없습니다.

3. 2순위 해법 — Pandas UDF (벡터화)

내장 함수로는 도저히 안 되는 진짜 커스텀 로직 — 복잡한 파싱, ML 추론, 특수 통계 계산 — 이 있을 때 꺼내는 카드가 Pandas UDF 입니다. Apache Arrow 를 이용해 데이터를 배치(컬럼) 단위로 한 번에 넘기기 때문에, 행마다 왕복하던 직렬화·인터프리터 비용이 사라집니다.

Pandas UDF 를 사용하기 전에 먼저 Arrow 를 켜 두세요. 이 설정 한 줄이 없으면 속도 향상이 절반으로 줄어듭니다.

spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")

Series → Series (스칼라 Pandas UDF)

Pandas UDF 중 가장 자주 만나는 형태입니다. 입력 컬럼 하나를 pd.Series 로 받아, pandas/numpy 벡터 연산으로 결과를 돌려줍니다. 루프가 없는 덕분에 일반 UDF 보다 훨씬 빠릅니다.

import pandas as pd
from pyspark.sql.functions import pandas_udf
 
@pandas_udf("double")
def celsius_to_f(c: pd.Series) -> pd.Series:
    return c * 9 / 5 + 32          # pandas 벡터 연산, 루프 없음
 
df = df.withColumn("temp_f", celsius_to_f("temp_c"))

여러 컬럼 입력

@pandas_udf("double")
def weighted(a: pd.Series, b: pd.Series, w: pd.Series) -> pd.Series:
    return (a * w + b * (1 - w))
 
df = df.withColumn("score", weighted("x", "y", "w"))

4. applyInPandas — 그룹별 pandas 처리

"그룹 단위로 pandas 를 자유자재로 쓰고 싶다"는 요구에 딱 맞는 도구입니다. 그룹별 정규화, 시계열 보간, 그룹별 모델 적용처럼 행 단위나 컬럼 단위로는 표현하기 어려운 로직을 깔끔하게 처리할 수 있습니다.

from pyspark.sql.types import StructType, StructField, DoubleType, StringType
 
schema = StructType([
    StructField("user_id", StringType()),
    StructField("z_score", DoubleType()),
])
 
def normalize(pdf: pd.DataFrame) -> pd.DataFrame:
    # pdf 는 한 그룹 전체 (pandas DataFrame)
    pdf["z_score"] = (pdf["amount"] - pdf["amount"].mean()) / pdf["amount"].std()
    return pdf[["user_id", "z_score"]]
 
result = df.groupBy("user_id").applyInPandas(normalize, schema=schema)

주의: applyInPandas 는 그룹 하나가 한 익스큐터 메모리에 다 들어가야 합니다. 그룹이 거대하면(스큐) OOM 이 납니다. 그룹 크기를 확인하고, 큰 그룹은 사전 분할하세요.

5. mapInPandas — 파티션 단위 스트리밍 처리

applyInPandas 가 그룹 단위라면, mapInPandas 는 파티션 단위입니다. 파티션을 pandas DataFrame 배치들의 이터레이터로 받기 때문에 전체를 한꺼번에 메모리에 올리지 않아도 됩니다. 대용량 데이터에 ML 추론을 태울 때 특히 유용한 패턴이죠.

from typing import Iterator
 
def predict(batches: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]:
    model = load_model()                 # 파티션당 1회 로드 (행마다 아님!)
    for pdf in batches:
        pdf["pred"] = model.predict(pdf[features])
        yield pdf
 
result = df.mapInPandas(predict, schema="... pred double")

mapInPandas 의 큰 장점은 모델·커넥션 같은 무거운 객체를 배치 루프 밖에서 한 번만 초기화할 수 있다는 것입니다. 일반 UDF 로 추론하면 행마다 모델을 들고 다녀 비효율적입니다. (대규모 배치 추론 패턴은 별도 글 "Spark + LLM 연동 가이드"에서도 다룹니다.)

6. 어떤 걸 언제 쓰나

지금까지 소개한 방법들을 언제 꺼내야 할지 한눈에 정리해 보겠습니다. 결정 기준은 간단합니다 — 내장 함수부터 따져보고, 안 될 때만 아래로 내려가면 됩니다.

7. UDF 사용 시 추가 주의

UDF 를 선택했다면, 다음 함정들을 미리 알고 가세요. 이 중 하나라도 놓치면 성능이 다시 바닥으로 돌아오거나 조용히 틀린 결과가 나올 수 있습니다.

• WHERE 절에 UDF 금지: UDF 로 필터하면 Catalyst 가 못 들여다봐 pushdown·파티션 프루닝이 깨집니다. 필터는 내장 표현식으로, UDF 는 select(투영) 단계에서.
• Arrow 활성화 확인: spark.sql.execution.arrow.pyspark.enabled=true 가 꺼져 있으면 Pandas UDF 도 느려집니다.
• 타입 일치: UDF 반환 타입과 선언 스키마가 어긋나면 조용히 NULL 이 되거나 에러. 명시적으로.
• Python 메모리: Pandas UDF 는 Python 워커 메모리를 씁니다. 배치가 크면 overhead OOM 위험(별도 글 "PySpark Executor OOM 정복" 참고).
• null 처리: pandas 연산에서 NaN/None 처리를 명시하지 않으면 결과가 틀어질 수 있습니다.

마치며 — 핵심 요약

• 일반 Python UDF 는 행마다 JVM↔Python 을 왕복합니다. 1억 행이면 1억 번의 직렬화 비용이 발생합니다.
• 내장 함수가 언제나 1순위입니다. JVM 안에서 컬럼 단위로 실행되므로 Python 왕복이 전혀 없습니다.
• Pandas UDF 는 Arrow 배치로 한 번에 전달합니다. 커스텀 로직이 필요하다면 일반 UDF 대신 Pandas UDF 를 선택하세요.
• 그룹별 복잡 로직은 applyInPandas, 대용량 ML 추론처럼 무거운 초기화가 필요하면 mapInPandas 를 쓰세요.
• WHERE 절에 UDF 를 쓰면 Catalyst 최적화가 깨집니다. 필터는 반드시 내장 표현식으로 처리하세요.

"내장으로 되나? 안 되면 Pandas UDF" — 이 한 줄 습관만 들여도, 여러분의 Spark 잡은 오늘보다 훨씬 빨라질 수 있습니다.

8. 정리

PySpark 성능 문제의 단골 원인은 "당연하게 쓴 Python UDF" 입니다. 핵심 통찰은 두 가지 — Spark 는 JVM 엔진이므로 행 단위 Python 왕복이 가장 비싸고, 이를 없애는 길은 내장 함수(JVM 내) 아니면 Arrow 벡터화(배치 단위) 라는 것입니다. "내장으로 되나? 안 되면 Pandas UDF" 한 줄만 습관이 되어도, 잡의 절반이 빨라지는 경험을 하게 됩니다.

이 글은 Spark 3.5 + PyArrow 기준으로 작성되었습니다. PySpark 파이프라인의 UDF 성능 개선이나 대규모 추론 최적화가 필요하시면 언제든 문의해 주세요.

— Data Dynamics 엔지니어링 팀