PySpark如何通过单次操作收集多DataFrame观测指标避免多次Action？

我有一个PySpark作业，从表a读取数据，执行若干转换与过滤操作后将结果写入表b。以下是简化后的代码：

import pyspark.sql.functions as F

spark = ...  # initialization

df = spark.table("a").where(F.col("country") == "abc")
df_unique = df.distinct()
users_without_kids = df_unique.where(F.col("kid_count") == 0)

observation = Observation()
observed_df = users_without_kids.observe(observation, F.count(F.lit(1)).alias("row_count"))

observed_df.writeTo("b")
print(observation.get["row_count"])

这段代码运行正常，我可以获取写入表b的记录数。但我还想了解：

• 第一次过滤后（即df）的记录数
• 去重后（即df_unique）的记录数

我希望避免触发额外的Action（例如不多次调用count()），理想情况是在单次Action（即writeTo操作）中收集所有指标。我尝试过添加多个observe调用或在单个Observation中添加多个指标，但在仅执行末尾单次Action时似乎无法生效。

核心问题：在PySpark中是否存在一种方法，可在单次操作中观测多个DataFrame（或多个指标），从而在不执行额外作业的前提下捕获上述三个DataFrame的记录数？

当然可以实现。核心思路是在数据流转的关键阶段添加条件标记列，然后在最终的Observation中一次性统计所有阶段的记录数，全程只触发一次Action。

修改后的代码 #

import pyspark.sql.functions as F

spark = ...  # initialization

# 1. 读取数据并添加第一次过滤的标记，再执行过滤
df = spark.table("a").withColumn(
    "is_after_first_filter",
    F.when(F.col("country") == "abc", 1).otherwise(0)
).where(F.col("country") == "abc")

# 2. 去重后添加去重阶段的标记
df_unique = df.distinct().withColumn(
    "is_after_distinct", 1
)

# 3. 过滤无子女用户，添加最终阶段标记
users_without_kids = df_unique.where(F.col("kid_count") == 0).withColumn(
    "is_final", 1
)

# 创建Observation，一次性统计三个阶段的记录数
observation = Observation()
observed_df = users_without_kids.observe(
    observation,
    F.sum("is_after_first_filter").alias("first_filter_count"),
    F.sum("is_after_distinct").alias("distinct_count"),
    F.count(F.lit(1)).alias("final_row_count")
)

# 执行单次Action：写入表b
observed_df.writeTo("b")

# 获取所有统计指标
print("第一次过滤后记录数:", observation.get["first_filter_count"])
print("去重后记录数:", observation.get["distinct_count"])
print("最终写入记录数:", observation.get["final_row_count"])

逻辑说明 #

1. 条件标记：在每个数据处理阶段添加标记列，标记当前数据是否属于该阶段的有效数据：
   • is_after_first_filter：标记通过第一次过滤的记录（值为1）
   • is_after_distinct：标记去重后的记录（去重后的每条数据都是唯一有效项，直接设为1）
2. 合并统计：在最终的observe中，通过sum函数统计标记列的总和，即可得到对应阶段的总记录数——因为每个有效记录贡献1，总和就是该阶段的行数。
3. 单次Action：所有统计逻辑都附加在最终的数据流转链路中，只有writeTo这一个Action触发作业，所有指标会在作业执行过程中被计算并收集到Observation中，完全避免额外作业开销。

内容的提问来源于stack exchange，提问作者עומר אמזלג