dataframe - 有没有一种惯用的方法来缓存 Spark 数据帧？

Question

我有一个大型 Parquet 数据集，我正在使用 Spark 读取它。阅读后，我筛选出在许多应用不同转换的函数中使用的行子集:

以下是与我要完成的类似但不完全相同的逻辑:

df = spark.read.parquet(file)
special_rows = df.filter(col('special') > 0)

# Thinking about adding the following line
special_rows.cache()

def f1(df):
  new_df_1 = df.withColumn('foo', lit(0))
  return new_df_1

def f2(df):
  new_df_2 = df.withColumn('foo', lit(1))
  return new_df_2

new_df_1 = f1(special_rows)
new_df_2 = f2(special_rows)
output_df = new_df_1.union(new_df_2)
output_df.write.parquet(location)

因为许多函数可能正在使用这个过滤后的行子集，所以我想缓存或保留它以潜在地加快执行速度/内存消耗。我知道在上面的示例中，在我最终写入 parquet 之前没有调用任何操作。

我的问题是，我是否需要插入对 count() 的某种调用，例如，为了触发缓存，或者如果 Spark 在最终写入 parquet 调用期间将是能够看到此数据帧正在 f1 和 f2 中使用，并将缓存数据帧本身。

如果是，这是一种惯用的方法吗？这是否意味着在依赖缓存的生产和大规模 Spark 作业中，会经常使用强制对数据帧执行先发制人操作的随机操作，例如调用 count？

Answer 1

there is no action called until my final write to parquet.

和

Spark during that final write to parquet call will be able to see that this dataframe is being used in f1 and f2 and will cache the dataframe itself.

是正确的。如果您执行 output_df.explain()，您将看到查询计划，这将表明您所说的是正确的。

因此，无需执行 special_rows.cache()。通常，cache 仅当您打算在 强制 Spark 计算某些内容后重用数据帧时才需要，例如在 write 或 show 之后。如果您发现自己有意调用 count()，则您可能做错了什么。