apache-spark - 使用 Spark pandas_udf 创建具有动态输入列数的列

Question

我有这个 df:

df = spark.createDataFrame(
    [('row_a', 5.0, 0.0, 11.0),
     ('row_b', 3394.0, 0.0, 4543.0),
     ('row_c', 136111.0, 0.0, 219255.0),
     ('row_d', 0.0, 0.0, 0.0),
     ('row_e', 0.0, 0.0, 0.0),
     ('row_f', 42.0, 0.0, 54.0)],
    ['value', 'col_a', 'col_b', 'col_c']
)

我想使用 Pandas 中的 .quantile(0.25, axis=1) 来添加一列:

import pandas as pd
pdf = df.toPandas()
pdf['25%'] = pdf.quantile(0.25, axis=1)
print(pdf)
#    value     col_a  col_b     col_c      25%
# 0  row_a       5.0    0.0      11.0      2.5
# 1  row_b    3394.0    0.0    4543.0   1697.0
# 2  row_c  136111.0    0.0  219255.0  68055.5
# 3  row_d       0.0    0.0       0.0      0.0
# 4  row_e       0.0    0.0       0.0      0.0
# 5  row_f      42.0    0.0      54.0     21.0

性能对我来说很重要，所以我假设 pyspark.sql.functions 中的 pandas_udf 可以以更优化的方式完成。但我很难做出一个高性能和有用的功能。这是我最好的尝试:

from pyspark.sql import functions as F
import pandas as pd
@F.pandas_udf('double')
def quartile1_on_axis1(a: pd.Series, b: pd.Series, c: pd.Series) -> pd.Series:
    pdf = pd.DataFrame({'a':a, 'b':b, 'c':c})
    return pdf.quantile(0.25, axis=1)

df = df.withColumn('25%', quartile1_on_axis1('col_a', 'col_b', 'col_c'))

1. 我不喜欢我需要为每一列提供一个参数，然后在函数中分别处理这些参数以创建一个 df。所有这些列都有相同的目的，所以恕我直言，应该有一种方法可以将它们一起处理，就像在这个伪代码中一样:

   def quartile1_on_axis1(*cols) -> pd.Series:
       pdf = pd.DataFrame(cols)
   

   这样我就可以将此函数用于任意数量的列。

2. 是否需要在UDF中创建一个pd.Dataframe？对我来说，这似乎与没有 UDF(Spark df -> Pandas df -> Spark df)一样，如上所示。如果没有 UDF，它甚至更短。我真的应该尝试让它在性能方面与 pandas_udf 一起工作吗？我认为 pandas_udf 是专门为这种目的而设计的......

Answer 1

udf 方法将为您提供所需的结果，而且绝对是最直接的方法。但是，如果性能确实是重中之重，您可以为 quantile 创建自己的原生 Spark 实现。基础知识可以很容易地编写代码，如果您想使用任何其他 pandas 参数，您需要自己进行调整。

注意:这取自 pandas API docs对于 interpolation='linear'。如果您打算使用它，请测试性能并自行在大型数据集上验证结果。

import math
from pyspark.sql import functions as f

def quantile(q, cols):
    if q < 0 or q > 1:
        raise ValueError("Parameter q should be 0 <= q <= 1")

    if not cols:
        raise ValueError("List of columns should be provided")

    idx = (len(cols) - 1) * q
    i = math.floor(idx)
    j = math.ceil(idx)
    fraction = idx - i

    arr = f.array_sort(f.array(*cols))

    return arr.getItem(i) + (arr.getItem(j) - arr.getItem(i)) * fraction


columns = ['col_a', 'col_b', 'col_c']

df.withColumn('0.25%', quantile(0.25, columns)).show()

+-----+--------+-----+--------+-----+-------+
|value|   col_a|col_b|   col_c|col_d|  0.25%|
+-----+--------+-----+--------+-----+-------+
|row_a|     5.0|  0.0|    11.0|    1|    2.5|
|row_b|  3394.0|  0.0|  4543.0|    1| 1697.0|
|row_c|136111.0|  0.0|219255.0|    1|68055.5|
|row_d|     0.0|  0.0|     0.0|    1|    0.0|
|row_e|     0.0|  0.0|     0.0|    1|    0.0|
|row_f|    42.0|  0.0|    54.0|    1|   21.0|
+-----+--------+-----+--------+-----+-------+

作为旁注，还有 pandas API on spark ，但是 axis=1 尚未(尚未)实现。这可能会在未来添加。

df.to_pandas_on_spark().quantile(0.25, axis=1)

NotImplementedError: axis should be either 0 or "index" currently.