如何使用Python变量对Spark DataFrame执行聚合计数统计

你提供的参考SQL存在两处逻辑问题：一是将客户ID列表作为单个值做等值匹配，二是没有按customerid分组，直接执行只会返回全局单条计数，无法得到每个ID对应的统计结果。以下是两种可直接运行的实现方式：

方法1：DataFrame API 实现（推荐，性能最优） #

这种方式避免了SQL字符串拼接的风险，执行效率更高：

from pyspark.sql import functions as F

# 1. 构造过滤截止时间：取传入date的日期部分拼接21:00:00
cutoff_ts = f"{date.strftime('%Y-%m-%d')} 21:00:00"

# 2. 基础筛选+分组统计（如果不需要保留无匹配记录的ID，用这段即可）
result_df = df.filter(
    (F.col("createdate") < F.lit(cutoff_ts).cast("timestamp"))
    & (F.col("customerid").isin(customerids))
).groupBy("customerid").count()

# --- 可选：如果需要保证传入的所有customerid就算无匹配记录也返回计数0，替换成下面的逻辑 ---
id_dim_df = spark.createDataFrame(
    [(cid,) for cid in customerids],
    schema="customerid long"
)
result_df = id_dim_df.join(
    df.filter(F.col("createdate") < F.lit(cutoff_ts).cast("timestamp")),
    on="customerid",
    how="left"
).groupBy("customerid").count().fillna(0, subset=["count"])

基于你提供的样例数据，以上代码输出结果和预期完全一致：1028598965607080计数为3，其余5个客户ID计数均为1。

方法2：Spark SQL 实现 #

如果更习惯写SQL逻辑，可以用临时视图传参执行：

# 1. 注册临时视图
df.createOrReplaceTempView("customer_table")

# 2. 构造查询参数
cutoff_ts = f"{date.strftime('%Y-%m-%d')} 21:00:00"
cid_in_clause = ",".join(map(str, customerids))

# 3. 执行查询
result_df = spark.sql(f"""
SELECT customerid, count(1) as count
FROM customer_table
WHERE createdate < TO_TIMESTAMP('{cutoff_ts}')
  AND customerid IN ({cid_in_clause})
GROUP BY customerid
""")

注意事项 #

• 如果createdate字段本身已经是Timestamp类型，不需要加类型转换逻辑，直接比较即可
• 不要用for循环逐个客户ID查询统计，用isin/IN批量过滤的性能远高于循环查询
• 样例数据中所有createdate值均为2022-06-20 15:03，早于截止时间21:00:00，因此所有记录都满足时间筛选条件

内容的提问来源于stack exchange，提问作者Caroline Leite