我有一个PySpark UDF,它接受一个数组并返回其后缀:
func.udf( lambda ng: ng[1:], ArrayType(IntegerType()) )
是否可以将其转换为标量pandas_udf? 熊猫是否提供必要的矢量化操作?
谢谢
答案 0 :(得分:2)
您可以使用标量熊猫udf获得相同的功能,但是请确保从udf返回带有列表列表的Series,因为该系列通常需要元素列表,并且如果您将row array展平并转换为多行直接将列表作为系列返回。
已编辑正如OP在评论中指出的那样,我以前使用index的答案(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:])