对于涉及多个领域的大型代码生成方法，如何提高火花性能？

Question

我想做一个数组大小为100的数组聚合总和，这是我的蛮力。 （我知道我可以写一个标量udaf，但是我想将spark的代码生成推到极限。

Spark v2.4

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession \
    .builder \
    .master('local[10]') \
    .appName('Notebook') \
    .config('spark.sql.codegen.maxFields', '10000') \
    .config('spark.sql.codegen.methodSplitThreshold', '100000') \
    .getOrCreate()

df = spark.range(2969622).select(F.array(*[F.rand() for i in range(100)]).alias('v')).cache()
df.count()
# Trigger cache.

# The followings are equivalent to
# SELECT
#     ARRAY(SUM(v[0]), SUM(v[1]),...) as v
# FROM ...

df.agg(
    F.array(*[F.sum(F.col('v')[i]) for i in range(0, 30)]).alias('v'),
).show()
# 357 ms

df.agg(
    F.array(*[F.sum(F.col('v')[i]) for i in range(0, 40)]).alias('v'),
).show()
# 5.51 sec

似乎有一个阈值，一旦达到该阈值，性能就会大大降低。

所以，我最终这样做了

a = df.agg(
    *[F.sum(F.col("v")[i]) for i in range(0, 25)],
)
b = df.agg(
    *[F.sum(F.col("v")[i]) for i in range(25, 50)],
)
c = df.agg(
    *[F.sum(F.col("v")[i]) for i in range(50, 75)],
)
d = df.agg(
    *[F.sum(F.col("v")[i]) for i in range(75, 100)],
)
x = a.crossJoin(b).crossJoin(c).crossJoin(d).selectExpr('array(*) as v').toPandas()
# 2.4 sec

我想知道在这种情况下spark中是否有一个配置可以控制spark代码源的优化。在这个特定的示例中，也许我可以放松一点以获得更好的性能。

我尝试了spark.sql.codegen.maxFields和spark.sql.codegen.methodSplitThreshold，但没有成功。