对数组进行操作的 Pandas UDF答案

【问题标题】：Pandas UDF that operates on arrays对数组进行操作的 Pandas UDF
【发布时间】：2019-06-23 06:34:17
【问题描述】：

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

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

是否可以将其转换为标量 pandas_udf？ Pandas 是否提供必要的矢量化操作？

谢谢

【问题讨论】：

标签： pandas pyspark apache-spark-sql

【解决方案1】：

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

已编辑 正如 OP 在 cmets 中指出的那样，我之前使用索引的答案（pd.Series([v[0][1:]]) 是错误的，并且仅在某些条件下有效。现在已更正.

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]|
+---------+----+------+

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

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

【讨论】：

谢谢 Ganesh。 1. 您期望从这两种解决方案中的哪一种获得最大的速度收益？ 2. v[0][1:]中的“[0]”是什么意思？ 3. 我注意到您的第一个解决方案往往只适用于从头开始创建的 Spark Dataframes。在其他情况下，我们得到一半的时间：RuntimeError：来自 pandas_udf 的结果向量不是所需的长度：预期 XXX，得到 1 我认为这与分区有关，因为即使在您的示例中，如果您将其修改为：df .repartition(3).withColumn('col3',func(df.col1)).show() 你会得到同样的错误。想法？谢谢
@GoErlangen 感谢您的询问并指出我的错误。 1.pandas apply 方法应该更快。 2 & 3 实际上是相关的。 2.当将pandas udf应用于列时，它将列作为一个系列。所以我正在访问该系列的第一行。所以我的答案只返回第一行。 3. 是的，这是由于分区和根据执行程序的数量，在我的示例中，单个执行程序可能会获得 1 行或所有 2 行。这会导致错误，因为当执行程序获取所有 2 行时，我只返回方法中的第一行。我将编辑答案。