在 PySpark groupBy 中，如何按组计算执行时间？

Question

我在一个大学项目中使用 PySpark，我有大型数据框，我使用 groupBy 应用了 PandasUDF。基本上调用是这样的：

df.groupBy(col).apply(pandasUDF)

我在我的 Spark 配置中使用 10 个内核 (SparkConf().setMaster('local[10]'))。

目标是能够报告每个组运行我的代码所花费的时间。我想要每组完成的时间，以便我可以取平均值。我也对计算标准差感兴趣。

我现在正在使用我知道将分成 10 组的已清理数据进行测试，并且我让 UDF 使用 time.time() 打印运行时间。但是，如果我要使用更多组，这是不可能的（对于上下文，我的所有数据都将被分成 3000 个左右的组）。有没有办法衡量每个组的执行时间？

Answer 1

如果不想将执行时间打印到标准输出，您可以将其作为 Pandas UDF 的额外列返回，例如

@pandas_udf("my_col long, execution_time long", PandasUDFType.GROUPED_MAP)
def my_pandas_udf(pdf):
    start = datetime.now()
    # Some business logic
    return pdf.assign(execution_time=datetime.now() - start)

或者，要计算驱动程序应用程序的平均执行时间，您可以使用两个Accumulators 在 UDF 中累积执行时间和 UDF 调用次数。例如

udf_count = sc.accumulator(0)
total_udf_execution_time = sc.accumulator(0)

@pandas_udf("my_col long", PandasUDFType.GROUPED_MAP)
def my_pandas_udf(pdf):
    start = datetime.now()
    # Some business logic
    udf_count.add(1)
    total_udf_execution_time.add(datetime.now() - start)
    return pdf

# Some Spark action to run business logic

mean_udf_execution_time = total_udf_execution_time.value / udf_count.value