spark coalesce（20）覆盖repartition（1000）.groupby（xxx）.apply（func）的并行性

Question

注意：这不是问合并与重新分区之间的区别的问题，谈论这个问题有很多，我的是不同的。

我有一份pysaprk工作

df = spark.read.parquet(input_path)

@pandas_udf(df.schema, PandasUDFType.GROUPED_MAP)
def train_predict(pdf):
    ...
    return pdf

df = df.repartition(1000, 'store_id', 'product_id')
df1 = df.groupby(['store_id', 'product_id']).apply(train_predict)

df1 = df1.withColumnRenamed('y', 'yhat')

print('Partition number: %s' % df.rdd.getNumPartitions())

df1.write.parquet(output_path, mode='overwrite')

默认200分区将要求大内存，因此我将重新分区更改为1000。

spark webui的工作详细信息如下：

由于输出仅为44M，我尝试使用coalesce来避免过多的小文件使hdfs变慢。 我所做的只是在.coalesce(20)之前添加.write.parquet(output_path, mode='overwrite')：

df = spark.read.parquet(input_path)

@pandas_udf(df.schema, PandasUDFType.GROUPED_MAP)
def train_predict(pdf):
    ...
    return pdf

df = df.repartition(1000, 'store_id', 'product_id')
df1 = df.groupby(['store_id', 'product_id']).apply(train_predict)

df1 = df1.withColumnRenamed('y', 'yhat')

print('Partition number: %s' % df.rdd.getNumPartitions())  # 1000 here

df1.coalesce(20).write.parquet(output_path, mode='overwrite')

然后spark webui显示：

似乎只有20个任务正在运行。

当repartion（1000）时，并行度取决于我的vcores数，此处为36。而且我可以直观地跟踪进度（进度条大小为1000）。 合并（20）之后，先前的分区（1000）失去了功能，并行性降低到20，也失去了直觉。 并且添加coalesce(20)会导致整个工作陷入困境，并且在没有通知的情况下失败。

将coalesce(20)更改为repartition(20)是可行的，但是根据文档，coalesce(20)效率更高，并且不会引起此类问题。

我想要更高的并行度，只有结果合并为20。正确的方法是什么？

Answer 1

coalesce被Spark优化器认为是狭窄的转换，因此它将创建一个从您的groupby到输出的单个WholeStageCodegen阶段，从而将并行度限制为20。

repartition是一种广泛的转换（即强制改组），如果您使用它而不是coalesce（如果添加了新的输出级但保留了groupby-train并行性）。

repartition(20)在您的用例中是一个非常合理的选择（混洗很小，因此成本很低）。

另一种选择是明确阻止Spark优化器合并您的预测阶段和输出阶段，例如在合并之前使用cache或persist：

# Your groupby code here

from pyspark.storagelevel import StorageLevel

df1.persist(StorageLevel.MEMORY_ONLY)\
   .coalesce(20)\
   .write.parquet(output_path, mode='overwrite')

鉴于您的小输出大小，MEMORY_ONLY持久+合并应该比分区快，但是当输出大小增加时，这将不成立