PySpark 4 — UDF, pandas_udf & Arrow

Ranh giới Python↔JVM quyết định hiệu năng

Ở PySpark tổng quan ta đã nói: engine Spark chạy trên JVM, còn code Python của bạn chạy trong tiến trình Python riêng. Khi viết bằng DataFrame API với các hàm dựng sẵn (F.col, F.when, F.sum...), bạn chỉ mô tả phép biến đổi; toàn bộ tính toán thực chất diễn ra bên JVM. Nhưng ngay khi bạn viết một hàm Python tự định nghĩa (UDF — User Defined Function) và áp lên cột, dữ liệu buộc phải rời JVM, sang tiến trình Python, rồi quay về. Chính chuyến "vượt biên" này là nguyên nhân số một khiến một job PySpark chậm gấp hàng chục lần mà nhìn code không thấy gì sai.

Bài này trả lời ba câu hỏi thực chiến: khi nào được phép viết UDF, vì sao Python UDF chậm, và làm sao lấy lại tốc độ bằng pandas_udf (vectorized UDF) dựa trên Apache Arrow. Đây là chủ đề đặc thù của PySpark — người dùng Scala/SQL gần như không gặp — nên nắm chắc nó là điểm phân biệt kỹ sư "biết dùng" với "dùng được".

Vì sao built-in F.* luôn ưu tiên

Mọi hàm trong pyspark.sql.functions (thường import là F) không chứa logic Python thật sự. F.upper(col) chỉ tạo ra một biểu thức (Catalyst expression) mô tả "viết hoa cột này". Khi job chạy, biểu thức đó được Catalyst optimizer phân tích, gộp, sắp lại thứ tự, rồi sinh mã bytecode JVM (codegen) chạy trực tiếp trên dữ liệu nằm sẵn trong JVM. Ba hệ quả:

• Không serialize qua Python. Dữ liệu ở đâu tính ở đó, không copy sang tiến trình khác.
• Được tối ưu. Catalyst biết ngữ nghĩa của F.upper, F.filter, F.sum nên có thể đẩy lọc xuống nguồn (predicate pushdown), gộp phép chiếu cột, loại biểu thức trùng.
• Vectorize + codegen. Whole-stage codegen xử lý cả khối dữ liệu bằng vòng lặp chặt trên JVM, tận dụng CPU cache.

Vì thế, quy tắc vàng: luôn tìm hàm built-in trước. Kho pyspark.sql.functions rất giàu — chuỗi (regexp_replace, split, concat_ws), ngày (date_trunc, datediff), điều kiện (when, coalesce), mảng/JSON (explode, from_json, transform), toán/thống kê, hash, và từ Spark 3.x có cả F.expr để nhúng biểu thức SQL. Rất nhiều thứ tưởng phải viết UDF thực ra đã có sẵn.

Python UDF: mạnh nhưng chậm

Khi logic thật sự không thể diễn đạt bằng built-in (một thuật toán parse phức tạp, gọi thư viện Python đặc thù), bạn dùng UDF thường:

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

@udf(returnType=StringType())
def mask_account(acc):
    if acc is None:
        return None
    return acc[:2] + "****" + acc[-3:]

df2 = df.withColumn("acc_masked", mask_account("account_no"))

Đoạn này chạy đúng, nhưng chuyện xảy ra bên dưới rất tốn kém. Với mỗi dòng, Spark phải:

1. Lấy giá trị cột từ dạng nội bộ của JVM.
2. Serialize nó bằng pickle và gửi qua socket sang tiến trình Python worker (mở kèm mỗi executor).
3. Python chạy hàm, trả kết quả.
4. Serialize ngược kết quả, gửi về JVM, giải mã về dạng nội bộ.

Bốn bước này lặp lại cho từng dòng một (row-by-row), không vectorize. Trên tỉ dòng, chi phí pickle + chuyển tiến trình lấn át hoàn toàn chi phí tính toán thật. Tệ hơn, Catalyst coi UDF là hộp đen: nó không biết hàm làm gì nên không thể đẩy lọc qua UDF, không loại được cột, không codegen. Một UDF vô tình đặt sớm trong pipeline có thể phá luôn predicate/column pushdown của cả nhánh phía sau (xem thêm tuning).

Tóm lại Python UDF có ba cái giá: serialize từng dòng, không vectorize, và phá tối ưu Catalyst. Dùng được, nhưng phải coi là lựa chọn cuối.

pandas_udf: vectorized UDF trên Apache Arrow

pandas_udf (giới thiệu từ Spark 2.3, hoàn thiện ở 3.x) giải quyết đúng ba cái giá trên. Thay vì trao một giá trị cho Python mỗi lần, Spark gom dữ liệu thành lô (batch) và trao cho hàm dưới dạng pandas.Series (hoặc DataFrame). Bên trong hàm bạn xử lý cả Series bằng thao tác vectorized của pandas/NumPy — tức là một lời gọi Python xử lý hàng chục nghìn dòng, thay vì một lời gọi cho một dòng.

Chìa khóa tốc độ là Apache Arrow: một định dạng bộ nhớ dạng cột (columnar) dùng chung được cả JVM lẫn Python hiểu. Nhờ Arrow, dữ liệu chuyển qua ranh giới ngôn ngữ theo lô, zero-copy hoặc gần như vậy — không còn pickle từng dòng. Kết quả: pandas_udf thường nhanh hơn Python UDF nhiều lần (thực nghiệm cộng đồng thường thấy khoảng vài lần đến hơn chục lần tùy phép tính; hãy tự đo trên dữ liệu của bạn).

Ba dạng pandas_udf

1. Series → Series (scalar). Dạng phổ biến nhất, thay thế trực tiếp Python UDF trên từng cột. Nhận một hoặc nhiều Series, trả về một Series cùng độ dài:

import pandas as pd
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import DoubleType

@pandas_udf(DoubleType())
def to_usd(vnd: pd.Series) -> pd.Series:
    return vnd / 24000.0          # vectorized: chia cả Series một lần

df2 = df.withColumn("amount_usd", to_usd("amount"))

2. Grouped map — applyInPandas. Khi cần xử lý nguyên một nhóm như một DataFrame (ví dụ chuẩn hóa trong từng khách hàng, nội suy chuỗi thời gian). Gọi trên df.groupBy(...).applyInPandas(fn, schema); hàm nhận một pandas.DataFrame của cả nhóm và trả về một pandas.DataFrame:

def normalize(pdf: pd.DataFrame) -> pd.DataFrame:
    pdf["z"] = (pdf["amount"] - pdf["amount"].mean()) / pdf["amount"].std()
    return pdf

out = df.groupBy("customer_id").applyInPandas(
    normalize, schema="customer_id long, amount double, z double")

Lưu ý: mỗi nhóm phải vừa trong bộ nhớ của một executor — nhóm quá lớn (data skew) sẽ gây OOM.

3. Grouped aggregate — Series → scalar. pandas_udf trả về một giá trị cho mỗi nhóm, dùng như hàm tổng hợp tự định nghĩa trong groupBy().agg():

@pandas_udf(DoubleType())
def geo_mean(v: pd.Series) -> float:
    import numpy as np
    return float(np.exp(np.log(v).mean()))

df.groupBy("account_id").agg(geo_mean("amount").alias("gm"))

Apache Arrow trong PySpark

Arrow không chỉ tăng tốc pandas_udf. Bật cấu hình:

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

thì hai thao tác chuyển dữ liệu Python↔JVM khác cũng nhanh lên đáng kể:

• df.toPandas() — kéo kết quả về driver dưới dạng pandas DataFrame: thay vì thu từng dòng qua pickle, Spark gom thành các Arrow batch.
• spark.createDataFrame(pandas_df) — đẩy pandas DataFrame lên Spark cũng qua Arrow.

Vài lưu ý quan trọng:

• Từ Spark 3, có cờ fallback spark.sql.execution.arrow.pyspark.fallback.enabled (mặc định bật): nếu chuyển đổi Arrow lỗi (kiểu không hỗ trợ), Spark tự quay về đường không-Arrow thay vì fail. Tiện nhưng dễ khiến bạn tưởng đang nhanh mà thực ra đã fallback.
• Không phải kiểu nào cũng map sạch sang Arrow. MapType, kiểu lồng phức tạp, hoặc TimestampType với múi giờ cần chú ý; phiên bản pyarrow phải tương thích.
• toPandas() vẫn gom toàn bộ về driver — Arrow làm nó nhanh hơn chứ không làm nó "phân tán". Kết quả lớn vẫn OOM driver.
• Cần cài pyarrow trên cả driver lẫn mọi executor.

Muốn ở lại API pandas nhưng vẫn phân tán (không kéo về driver), xem PySpark 5 — pandas API on Spark.

Cây quyết định: chọn loại hàm nào

Thứ tự ưu tiên: built-in F.* > pandas_udf > Python UDF. Chỉ tụt xuống bậc dưới khi bậc trên thật sự bất khả thi.

Luồng dữ liệu: row-by-row vs batch qua Arrow

Sơ đồ dưới cho thấy khác biệt bản chất về số lần "vượt biên" và cách đóng gói dữ liệu:

Với Python UDF, số lần serialize tỉ lệ thuận với số dòng. Với pandas_udf, số lần chỉ tỉ lệ với số batch — ít hơn hàng vạn lần — và mỗi lần chuyển là một khối cột nén gọn.

So sánh cùng một logic

Giả sử cần chuẩn hóa mã tài khoản: bỏ khoảng trắng, viết hoa, thêm tiền tố. Ba cách viết cho cùng kết quả nhưng hiệu năng rất khác (ví dụ minh họa):

# (A) Built-in — nhanh nhất, chạy hẳn trên JVM
from pyspark.sql import functions as F
a = df.withColumn("acc",
    F.concat(F.lit("VN-"), F.upper(F.trim("account_no"))))

# (B) Python UDF — chậm: serialize + gọi hàm từng dòng
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType
@udf(StringType())
def norm(x):
    return None if x is None else "VN-" + x.strip().upper()
b = df.withColumn("acc", norm("account_no"))

# (C) pandas_udf — theo lô, xử lý cả Series bằng vectorized string ops
from pyspark.sql.functions import pandas_udf
import pandas as pd
@pandas_udf(StringType())
def norm_v(s: pd.Series) -> pd.Series:
    return "VN-" + s.str.strip().str.upper()
c = df.withColumn("acc", norm_v("account_no"))

Cách đo trung thực: gọi một action (ví dụ .count() hoặc .write) và xem Spark UI → SQL tab để so thời gian stage, đừng đo df.show() vì nó chỉ lấy vài dòng. Kỳ vọng định tính: (A) nhanh nhất tuyệt đối; (C) chậm hơn (A) chút nhưng nhanh hơn (B) rõ rệt khi số dòng lớn; (B) chậm nhất và càng nhiều dòng càng tụt. Với ví dụ đơn giản như trên thì (A) thắng dễ — bài học là đừng viết UDF khi built-in làm được; pandas_udf dành cho logic mà F.* không diễn đạt nổi.

Cạm bẫy thường gặp

• Phá predicate/column pushdown. UDF là hộp đen với Catalyst. Đặt df.filter(...).withColumn("x", my_udf(...)) — lọc trước, UDF sau — để Spark còn đẩy được bộ lọc xuống nguồn. Nếu lọc dựa trên kết quả UDF thì pushdown mất; cân nhắc tính bằng built-in.
• returnType sai hoặc thiếu. UDF/pandas_udf bắt buộc khai báo kiểu trả về đúng với dữ liệu thật. Khai IntegerType() mà trả float, hoặc pandas trả int64 khi Spark chờ long mà lệch, sẽ ra null hoặc lỗi Arrow khó hiểu. Với pandas_udf hãy để dtype của Series khớp kiểu Spark.
• Null handling. JVM có null; sang Python thành None (scalar UDF) hoặc NaN/NaT/None trong Series. Luôn xử lý nhánh None; đừng để hàm ném exception trên dòng null làm hỏng cả task.
• Phụ thuộc thư viện trên executor. Hàm chạy ở executor, không phải driver. import sklearn trong UDF đòi mọi executor đều cài đúng phiên bản. Đóng gói môi trường (venv/conda-pack, hoặc image) và test trên cluster, không chỉ trên local.
• Nhóm quá lớn với applyInPandas. Mỗi nhóm gom vào bộ nhớ một executor; nhóm lệch (một khách hàng chiếm phần lớn dữ liệu) gây OOM. Xử skew trước (xem tuning).
• Tưởng đang dùng Arrow. Nếu quên bật cờ, hoặc kiểu bị fallback, toPandas() âm thầm chạy đường chậm. Kiểm tra cấu hình và log.

Use case thực tế

Bối cảnh NCB. Đội rủi ro cần chạy một hàm chấm điểm giao dịch nghi ngờ (parse mã kênh, chuẩn hóa nội dung chuyển khoản bằng regex, tra một bảng tra cứu nhỏ, rồi tính điểm) trên toàn bộ lịch sử giao dịch — cỡ tỉ dòng để backfill mô hình. Logic có bước parse chuỗi tùy biến mà F.* thuần không diễn đạt gọn, nên nhóm ban đầu viết Python UDF.

Vấn đề. Job backfill chạy hàng giờ và thường xuyên timeout. Nhìn Spark UI: phần lớn thời gian nằm ở các task Python worker — chính là chi phí serialize từng dòng qua pickle. Đây là triệu chứng kinh điển của Python UDF trên khối lượng lớn.

Cách xử lý. Nhóm tách logic làm hai phần:

1. Những bước có built-in (viết hoa, regexp_replace, when, join bảng tra cứu) chuyển hẳn sang F.* và broadcast join — chạy trên JVM.
2. Phần parse còn lại không thể built-in được viết lại thành pandas_udf Series→Series, xử lý theo lô bằng thao tác str/vectorized của pandas, và bật Arrow.

Kết quả (ước lượng, minh họa — cần tự đo). Sau khi chuyển, thời gian job giảm đáng kể — bậc độ lớn thường thấy trong các bài toán tương tự là rút ngắn nhiều lần (ví dụ từ khung "vài giờ" xuống "vài chục phút"), do bỏ được serialize từng dòng và đẩy phần lớn phép tính về JVM. Con số cụ thể phụ thuộc dữ liệu, số executor và độ phức tạp hàm; các mức nêu ở đây là ước lượng định hướng, không phải cam kết đo đạc. Bài học vận hành: đừng tối ưu bằng cảm giác — đo trước/sau trên Spark UI, và ưu tiên chuyển được bước nào sang built-in thì chuyển, chỉ để lại pandas_udf cho phần thật sự cần Python.

Ghi nhớ

• Built-in F.* luôn ưu tiên: chạy trên JVM, được Catalyst tối ưu, không serialize qua Python. Tìm hàm sẵn có trước khi nghĩ đến UDF.
• Python UDF chậm vì serialize (pickle) từng dòng qua ranh giới JVM↔Python, không vectorize, và biến pipeline thành hộp đen phá tối ưu Catalyst. Là lựa chọn cuối.
• pandas_udf (vectorized UDF) xử lý theo lô dưới dạng pandas.Series/DataFrame qua Apache Arrow (định dạng cột, gần zero-copy) nên nhanh hơn Python UDF nhiều lần.
• Ba dạng: Series→Series (scalar, thay UDF từng cột), applyInPandas (grouped map, cả nhóm là một DataFrame), grouped agg (Series→scalar).
• Bật spark.sql.execution.arrow.pyspark.enabled để tăng tốc toPandas(), createDataFrame() và pandas_udf; nhớ cài pyarrow trên driver và mọi executor; coi chừng fallback âm thầm.
• Thứ tự chọn: built-in > pandas_udf > Python UDF.
• Cạm bẫy: UDF phá predicate pushdown (lọc trước, UDF sau), khai đúng returnType, xử lý null, đảm bảo thư viện có trên executor, tránh nhóm quá lớn khi applyInPandas.
• Luôn đo trước/sau bằng action thật + Spark UI, đừng đo bằng show().

Nguồn tham khảo

• Apache Spark — PySpark User Guide: Python User-defined Functions (UDF)
• Apache Spark — PySpark User Guide: pandas function APIs & pandas_udf (Apache Arrow in PySpark)
• Apache Spark — pyspark.sql.functions.pandas_udf (API reference)
• Apache Spark — pyspark.sql.GroupedData.applyInPandas (API reference)
• Apache Arrow — Documentation
• Apache Spark — SQL Performance Tuning (cấu hình spark.sql.execution.arrow.pyspark.enabled)