在熊猫udf pyspark内部使用numpy

Question

我正在尝试定义一个熊猫udf，以计算每个周期的对数正态分布的偏度。

我目前已执行以下操作：

@pandas_udf("double", PandasUDFType.GROUPED_AGG)  
def lognormal_skew(v):
  return (np.exp(v.std()) + 2) * np.sqrt(np.exp(v.std()) - 1)

my_df.groupBy('period').agg(lognormal_skew(my_df['my_columns'])).show()

但是我得到一个错误：

rg.apache.spark.SparkException: Job aborted due to stage failure: Task 0 in stage 3047.0 failed 4 times, most recent failure: Lost task 0.3 in stage 3047.0 (TID 208812, 10.139.64.8, executor 82): org.apache.spark.api.python.PythonException: Traceback (most recent call last):

我的猜测是，这与numpy有一定关系，因为如果我尝试按以下方式定义偏斜：

@pandas_udf("double", PandasUDFType.GROUPED_AGG)  
def skew(v):
  return v.skew()

my_df.groupBy('period').agg(skew(my_df['my_columns'])).show()

它输出一个DataFrame并且没有错误。

Answer 1

序言

根据我的经验，我认为只要可以使用pyspark内置函数来实现某些东西，它就比用户定义的函数更可取。

udf的问题之一是错误消息很难解密。例如，就您而言，我不知道您为什么遇到此错误。

pyspark.sql.functions允许您做更多的事情，如果您愿意接受更多操作。但是，就性能而言，这是很难克服的，因为这些功能是由专家优化的。如果您想用pyspark.sql.functions做不到的事情（发生），我更喜欢使用rdd而不是udf。为了应用rdd功能，Python更自然。相对于内置的DataFrame方法，您的性能降低了，但是您获得了一些灵活性。

也许有关您的问题的示例可能具有启发性。

Python

让我们以基于numpy的示例为例。您提供了python实现：

import numpy as np
def lognormal_skew_numpy(v):
    return (np.exp(v.std()) + 2) * np.sqrt(np.exp(v.std()) - 1)

它可以用来控制其他实现是否一致：

print(lognormal_skew_numpy(np.array([1,3,5])))
print(lognormal_skew_numpy(np.array([5,6])))
# 14.448897615797454
# 2.938798148174726

DataFrame API逻辑

现在，让我们来看看Spark。我将使用以下DataFrame：

df = spark.createDataFrame([(1, 'a'), (3, 'a'), (5, 'a'), (5,'b'), (6,'b')], ['x','period'])
df.show(2)

+---+------+
|  x|period|
+---+------+
|  1|     a|
|  3|     a|
+---+------+
only showing top 2 rows

偏度函数仅执行基本的数学运算。它们都是在pyspark.sql.functions中实现的，因此在这种情况下创建一个执行该功能的函数并不难

import pyspark.sql.functions as psf

def lognormal_skew(df, xvar = 'x'):
    df_agg = (df
              .groupBy('period')
              .agg(psf.stddev_pop(xvar).alias('sd'))
             )
    df_agg = df_agg.withColumn('skew', (psf.exp(psf.col('sd')) + 2)*psf.sqrt(psf.exp('sd') - 1))
    return df_agg

请注意，psf中存在不同的函数来计算标准差：我使用的stddev_pop效率较低，但报告的是总体水平方差，而不是估算值（精度为3或2点，一个估计量将是很差的。

我们可以控制产生期望的输出：

lognormal_skew(df).show(2)
+------+-----------------+------------------+
|period|               sd|              skew|
+------+-----------------+------------------+
|     b|              0.5| 2.938798148174726|
|     a|1.632993161855452|14.448897615797454|
+------+-----------------+------------------+

我们设法通过纯DataFrame逻辑获得了预期的结果。

rdd

让我们将数据安排为一个rdd，看起来像是并行的numpy数组：

rdd = df.rdd
rdd = rdd.mapValues(lambda l: l).map(lambda l: (l[1], [l[0]] )).reduceByKey(lambda x,y: x + y)
rdd.take(2)
[('b', [5, 6]), ('a', [1, 3, 5])]

在这里，我们使用reduceByKey将值分组到一个列表中。在这一步中，大量数据可能会使RAM爆炸。

最后，您可以轻松地将函数与该结构并行进行：

rdd = rdd.map(lambda l: (l[0], np.array(l[1]))).map(lambda l: (l[0], lognormal_skew_numpy(l[1])))
rdd.take(2)
[('b', 2.938798148174726), ('a', 14.448897615797454)]

我们再次得到相同的结果。我看到这种方法有两个缺陷：

• 它的可读性和便携性较差。如果要在其他数据集上重复使用代码，则必须做更多的工作
• 效率较低（速度和内存）。这里的reduceByKey操作是主要瓶颈。

但是，您获得了一些灵活性。这是一个权衡。