对数组进行操作的 Pandas UDF

Question

我有一个 PySpark UDF，它接受一个数组并返回其后缀：

func.udf( lambda ng: ng[1:], ArrayType(IntegerType()) )

有没有可能把它变成一个标量pandas_udf？Pandas 是否提供必要的矢量化操作？

谢谢

Answer 1

您可以使用标量 Pandas udf 获得相同的功能，但请确保您从 udf 返回一个带有列表列表的系列，因为该系列通常需要一个元素列表，并且如果您直接返回，您的行数组将被展平并转换为多行列为系列。

编辑正如 OP 在评论中指出的那样，我之前的答案（pd.Series([v[0][1:]]) using index 是错误的，仅在某些条件下有效。现在更正。

df = spark.createDataFrame([([1,2,3],'val1'),([4,5,6],'val2')],['col1','col2'])
df.show()
+---------+----+
| col1|col2|
+---------+----+
|[1, 2, 3]|val1|
|[4, 5, 6]|val2|
+---------+----+

from pyspark.sql.functions import pandas_udf,PandasUDFType
from pyspark.sql.types import *
import pandas as pd

@pandas_udf(ArrayType(LongType()))
def func(v):
    res=[]
    for row in v:
        res.append(row[1:])
    return pd.Series(res)

df.withColumn('col3',func(df.col1)).show()
+---------+----+------+
|col1     |col2|col3  |
+---------+----+------+
|[1, 2, 3]|val1|[2, 3]|
|[4, 5, 6]|val2|[5, 6]|
+---------+----+------+

另一种方法是使用 apply 类似于您所做的：

@pandas_udf(ArrayType(LongType()))
def func(v):
    return v.apply(lambda x:x[1:])