UDF Pandas, работающая с массивами

Question

У меня есть PySpark UDF, который принимает массив и возвращает его суффикс:

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

Можно ли превратить его в скалярный pandas_udf?Предлагает ли Pandas необходимые векторизованные операции?

Спасибо

Answer 1

Вы можете получить ту же функциональность со скалярными пандами 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:])