python - 一列的乘法，如 f.sum

Question

我在数据框中有一列。我需要通过乘以该列中的值而不是将它们相加来聚合列。

ex = spark.createDataFrame([[1,2],[4,5]],['a','b'])
ex.show()
ex.agg(f.sum('a')).show()

我想将列“a”与类似的语法相乘，而不是求和:

ex.agg(f.mul('a')).show()

我想到的解决方法是:

ex.agg(f.exp(f.sum(f.log('a')))).show()

但是计算 exp(sum(log)) 可能不够高效

结果应该是 4。什么是最有效的方法？

Answer 1

没有内置的乘法聚合。您的解决方法对我来说似乎很有效，其他解决方案需要构建自定义聚合函数。

import pyspark.sql.functions as F
ex = spark.createDataFrame([[1,2],[4,5], [6,7], [3,2], [9,8], [4,2]],['a','b'])
ex.show()

+---+---+
|  a|  b|
+---+---+
|  1|  2|
|  4|  5|
|  6|  7|
|  3|  2|
|  9|  8|
|  4|  2|
+---+---+

# Solution 1
ex.agg(F.exp(F.sum(F.log('a')))).show()

+----------------+
|EXP(sum(LOG(a)))|
+----------------+
|          2592.0|
+----------------+

# Solution 2
from pyspark.sql.types import IntegerType

def mul_list(l):
    return reduce(lambda x,y: x*y, l)  # In Python 3, use `from functools import reduce`

udf_mul_list = F.udf(mul_list, IntegerType())
ex.agg(udf_mul_list(F.collect_list('a'))).show()

+-------------------------------+
|mul_list(collect_list(a, 0, 0))|
+-------------------------------+
|                           2592|
+-------------------------------+

# Solution 3
seqOp = (lambda local_result, row: local_result * row['a'] )
combOp = (lambda local_result1, local_result2: local_result1 * local_result2)
ex_rdd = ex.rdd
ex_rdd.aggregate( 1, seqOp, combOp)

Out[4]: 2592

现在让我们比较性能:

import random
ex = spark.createDataFrame([[random.randint(1, 10), 3] for i in range(10000)],['a','b'])

%%timeit
ex.agg(F.exp(F.sum(F.log('a')))).count()

10 loops, best of 3: 84.9 ms per loop

%%timeit
ex.agg(udf_mul_list(F.collect_list('a'))).count()

10 loops, best of 3: 78.8 ms per loop

%%timeit
ex_rdd = ex.rdd
ex_rdd.aggregate( 1, seqOp, combOp)

10 loops, best of 3: 94.3 ms per loop

性能在本地的一个分区上看起来差不多。在多个分区上尝试更大的数据框。

为了提高解决方案 2 和 3 的性能:build a custom aggregation function in Scala和 wrap it in Python