pandas_udf()

Описание

Функция pandas_udf() создает пользовательскую функцию (UDF) на основе Pandas, которая может быть использована в PySpark. Pandas UDF обеспечивают лучшую производительность по сравнению с обычными UDF, так как работают с целыми пакетами данных.

Параметры

f: Callable - функция Pandas
returnType: DataType - тип возвращаемого значения
functionType: String - тип функции (SCALAR, GROUPED_MAP, GROUPED_AGG)

Возвращаемое значение

Column - результат применения Pandas UDF

Пример использования

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

# Создаем SparkSession
spark = SparkSession.builder.appName("pandas_udf_example").getOrCreate()

# Создаем Pandas UDF для вычисления квадрата числа
@pandas_udf(DoubleType())
def square_pandas(x: pd.Series) -> pd.Series:
    return x * x

# Создаем DataFrame
data = [
    (1,),
    (2,),
    (3,)
]
df = spark.createDataFrame(data, ["number"])

# Применяем Pandas UDF
result = df.select(
    "number",
    square_pandas(col("number")).alias("squared")
).show()

# Результат:
# +------+-------+
# |number|squared|
# +------+-------+
# |     1|    1.0|
# |     2|    4.0|
# |     3|    9.0|
# +------+-------+

Примечания

Pandas UDF работают с целыми пакетами данных, что обеспечивает лучшую производительность
Поддерживаются три типа Pandas UDF:
SCALAR: работает с отдельными значениями
GROUPED_MAP: работает с группами данных
GROUPED_AGG: выполняет агрегацию по группам
Для вызова встроенных функций используйте call_function()
Для вызова обычных UDF используйте call_udf()
Для создания обычных UDF используйте udf() или udtf()