pyspark寄存器内置函数及其在spark.sql查询中的使用

Question

在spark.sql查询中注册和使用pyspark version 3.1.2内置函数的正确方式是什么？

下面是一个创建pyspark DataFrame对象并在纯SQL中运行简单查询的最小示例。

尝试使用...TypeError: Invalid argument, not a string or column: -5 of type <class 'int'>. For column literals, use 'lit', 'array', 'struct' or 'create_map' function...的pyspark内置函数运行相同查询的代码出现错误

import pandas as pd
import pyspark.sql.functions as F
import pyspark.sql.types as T
from pyspark.sql import SparkSession

# spark session and create spark DataFrame from pandas DataFrame
spark = SparkSession.builder.getOrCreate()
spark.conf.set('spark.sql.execution.arrow.pyspark.enabled', 'true')
pdf = pd.DataFrame({'x': [0, 0, 1, 1], 'y': [-5, -6, -7, -8]})
df = spark.createDataFrame(pdf)
df.createOrReplaceTempView('df')

# pure SQL
rv1 = spark.sql(f'SELECT ABS(x + y) AS z FROM df').toPandas()

# pyspark built-in function
spark.udf.register('abs_builtin', F.abs, T.LongType())
rv2 = spark.sql(f'SELECT abs_builtin(x + y) AS z FROM df').toPandas()

Answer 1

没有内置函数的解决方案是使用pandas_udf

import numpy as np
@F.pandas_udf(returnType=T.LongType())
def abs_pandas_udf(x: pd.Series) -> pd.Series:
    return np.abs(x)
spark.udf.register('abs_pandas_udf', abs_pandas_udf)
rv3 = spark.sql(f'SELECT abs_pandas_udf(x + y) AS z FROM df').toPandas()

Answer 2

pyspark内置标准函数期望列(pyspark.sql.Column)类型作为参数，而用户定义函数期望原始类型和复杂类型(基于函数定义)。这就是你的实现失败的原因。

仅当没有选项时，才建议使用UDF。