Question

我写了一个UDF。非常慢。我想用pandas_udf替换它，以利用矢量化的优势。

实际的udf有点复杂，但是我创建了一个简化的玩具版本。

我的问题：是否可以用pandas_udf替换玩具示例中的UDF，而该熊猫可以利用矢量化功能？如果没有，为什么不呢？

P.S：我知道没有UDF就能达到相同的效果。那是因为我简化了示例，但这不是我的目标。

from pyspark.sql import functions as f
from pyspark.sql.types import ArrayType, StringType
import pandas as pd

#Example data
df = spark.createDataFrame(pd.DataFrame({ 'Letter': [['A', 'A', 'C'], ['A', 'C', 'A', 'D']],
                                          'Number': [[2, 1, 1], [3, 1, 1, 2]],
                                        })
                          )

# The UDF I hope to replace with a pandas_udf
@f.udf(ArrayType(StringType()))
def array_func(le, nr):
    res=[]
    for i in range(len(nr)):
        if nr[i]==1:
            res.append(le[i])
        else:
            res.append('Nope')
    return res

# Applying the udf
df = df.withColumn('udf', array_func('Letter','Number'))
df.show()

Answer 1

这个怎么样？

from pyspark.sql import functions as F
from pyspark.sql.types import ArrayType, StringType
import pandas as pd

#Example data
df = spark.createDataFrame(pd.DataFrame({ 'Letter': [['A', 'A', 'C'], ['A', 'C', 'A', 'D']],
                                          'Number': [[2, 1, 1], [3, 1, 1, 2]],
                                        })
                          )
df.show()

# Add a dummy column so you can use groupby
df = df.withColumn('id', F.lit(1))
schm = StructType(df.schema.fields + [StructField('udf', ArrayType(StringType()), True)])
@pandas_udf(schm, PandasUDFType.GROUPED_MAP)
def array_udf(pdf):
    res=[]
    for ls, ns in zip(pdf['Letter'], pdf['Number']):
        r = [l if n == 1 else 'Nope' for l, n in zip(ls, ns)]
        res.append(r)
    pdf['udf'] = res
    return pdf

df = df.groupby('id').apply(array_udf).drop('id')
df.show()

输出：

+------------+------------+------------------+
|      Letter|      Number|               udf|
+------------+------------+------------------+
|   [A, A, C]|   [2, 1, 1]|      [Nope, A, C]|
|[A, C, A, D]|[3, 1, 1, 2]|[Nope, C, A, Nope]|
+------------+------------+------------------+

Answer 2

我已经使用pandas_udf创建了一个名为array_func_pd的新函数，只是为了区分原始的array_func，以便您可以比较和玩这两个函数。

from pyspark.sql import functions as f
from pyspark.sql.types import ArrayType, StringType
import pandas as pd

@f.pandas_udf(ArrayType(StringType()))
def array_func_pd(le, nr):
"""
   le:  pandas.Series< numpy.ndarray<string> >
   nr:  pandas.Series< numpy.ndarray<int> >

   return: pd.Series< list<string> >
"""
    res=[]
    for i, (l_lst, n_lst) in enumerate(zip(le, nr)):
        ret_lst = []
        res.append(ret_lst)
        l_lst2 = l_lst.tolist()
        n_lst2 = n_lst.tolist()
        for j,(l, n) in enumerate(zip(l_lst2, n_lst2)):
            if n == 1:
                ret_lst.append(l)
            else:
                ret_lst.append('Nope')
    return pd.Series(res)

# Applying the udf
df = df.withColumn('udf', array_func_pd('Letter','Number'))
df.show()

这是输出：

+------------+------------+------------------+
|      Letter|      Number|               udf|
+------------+------------+------------------+
|   [A, A, C]|   [2, 1, 1]|      [Nope, A, C]|
|[A, C, A, D]|[3, 1, 1, 2]|[Nope, C, A, Nope]|
+------------+------------+------------------+

有两种类型的熊猫UDF（又名矢量化UDF）。对于您的情况，我认为最好保持简单并使用Scalar Pandas UDF。

以下是官方文档中Scalar Pandas UDF的注释：

Python函数应将pandas.Series作为输入并返回相同长度的pandas.Series。在内部，Spark将通过将列拆分为批处理并为每个批处理调用该函数作为数据的子集，然后将结果串联在一起来执行Pandas UDF。

所以在我的代码中：

udf的输出应为pd.Series，并应与le或nr共享相同的计数。

在两个ArrayType（StringType（））字段上运行的pandas_udf

2 个答案: