Pandas 在阵列上运行的UDF

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

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

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

---

感谢

您可以使用udf获得相同的功能，但请确保从udf返回包含列表的序列，因为序列通常需要元素列表，如果您直接以序列形式返回列表，则行数组将被展平并转换为多行

编辑正如OP在评论中指出的那样，我以前的答案（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数据帧。在其他情况下，我们得到的一半时间：RuntimeError:Result vector from pandas_udf不是所需的长度：预期XXX，得到1我认为这与分区有关，因为即使在您的示例中，如果您将其修改为：df.rep分区（3）.withColumn（'col3'，func（df.col1））.show（）你也会犯同样的错误，你有什么想法？Thanks@GoErlangen感谢您的查询并指出我的错误。1.pandas apply方法应该快得多。2和3实际上是相关的。2.当将pandas udf应用于列时，它将列作为一个系列。因此我访问系列的第一行。因此我的答案只返回第一行。3。是的，这是由于分区的原因，并且取决于执行器的数量，在我的示例中，单个执行器可能会将1行或全部2行放在一起。这会导致错误，因为当执行器获取全部2行时，我只返回方法中的第一行。我将编辑答案。