Pandas UDF that operates on arrays - pandas

Question

I have a PySpark UDF that takes an array and returns its suffix:
func.udf( lambda ng: ng[1:], ArrayType(IntegerType()) )
Is it possible to turn it into a scalar pandas_udf?
Does Pandas offer the necessary vectorized operations?
Thanks

Answer 1

You can get the same functionality with scalar pandas udf but make sure that you return a Series with list of lists from the udf as the series normally expects a list of elements and your row array is flattened and converted to multiple rows if you return directly the list as series.
Edited As pointed out by OP in comments my previous answer(pd.Series([v[0][1:]]) using index is wrong and works only in certain conditions. Now corrected.
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]|
+---------+----+------+
The other way is to use apply similar to what you have done:
#pandas_udf(ArrayType(LongType()))
def func(v):
return v.apply(lambda x:x[1:])