寻找一种解决方案来加速大型数据集上的 pyspark.sql.GroupedData.applyInPandas 处理

Question

我正在使用存储在 S3 存储桶（镶木地板文件）中的数据集，该数据集总共包含 ~165 million 记录（约 30 列）。现在，要求是首先groupby 某个ID 列，然后根据数据为这些分组记录中的每一个生成250+ features。使用多个 Pandas 功能以及 10 多个支持功能来构建这些功能非常复杂。 groupby 函数应该生成~5-6 million records，因此最终输出应该是6M x 250 形状的数据帧。

现在，我在一个较小的样本上测试了代码，它运行良好。问题是，当我在整个数据集上实现它时，它需要很长时间 - Spark 显示中的进度条即使在运行 4 个多小时后也不会改变。我在连接到集群（1 个 m5.xlarge 主节点和 2 个 m5.xlarge 核心节点）的 AWS EMR Notebook 中运行它。 我试过1 m5.4xlarge Master & 2 m5.4xlarge Core Nodes、1 m5.xlarge Master & 8 m5.xlarge Core Nodes 等组合。他们都没有表现出任何进展。 我已经尝试在我的本地机器的内存中的 Pandas 中运行它以获得大约 65 万条记录，进度是大约 3.5 次迭代/秒，达到 ~647 hours 的 ETA。

所以，问题是 - 任何人都可以分享任何更好的解决方案来减少时间消耗并加快处理速度吗？是否应该为此用例使用另一种集群类型？是否应该对其进行重构，或者是否应该删除 Pandas 数据框的使用，或者任何其他指针都会非常有帮助。

提前非常感谢！

Answer 1

首先要做的事情是：您的数据分区是否足以利用所有员工？如果您的流程的某些部分导致它合并为例如一个分区，那么你基本上是在运行单线程。

除此之外，如果没有看到代码，我不确定，但这里有一个微妙的行为可能导致运行时变得庞大：

source_df = # some pandas dataframe with a lot of features in columns

flattened_df = your_df.stack().reset_index().unstack() # Turn the features into rows

spark_df = spark.createDataFrame(flattened_df) # 'index' is the column that contains the feature name

# a function to do a linear regression and calculate residual
def your_good_pandas_function(key, slice):
  clf = LinearRegression()
  X = slice[subset,of,columns]
  y = slice[key]
  clf.train(X,y)
  predicted = clf.predict(X)
  return y-predicted

def your_bad_pandas_function(key, slice):
  clf = LinearRegression()
  X = slice[subset,of,columns]
  y = slice[key]
  clf.train(X,y)
  predicted = clf.predict(X)
  return source_df[key]-predicted

spark_df.groupBy('index').applyInPandas(your_good_pandas_function,schema=some_schema) #fast
spark_df.groupBy('index').applyInPandas(your_bad_pandas_function,schema=some_schema) #slow

这两个 ApplyInPandas 函数做同样的事情——它们对某些特征进行线性回归并计算残差。第一个使用 pandas UDF 范围内的变量。第二个使用了一个超出 pandas UDF 范围的变量。在第二种情况下，Spark 将通过广播source_df 来帮助您对 pandas UDF 的每一次调用。这将导致大量内存使用并肯定会杀死您的工作。

您的数据似乎不够大，无法花费那么长时间，所以我的猜测是它适用于小子集而不是较大集的原因可能是因为您无意中将较大的集合广播到您的 applyInPandas 函数调用。