如何在PySpark中无需转换类型按日期列使用sampleBy采样？

问题背景 #

PySpark的sampleBy方法要求分组列的类型为float、int或str，直接对Date类型列使用时会触发类型错误；若仅将fractions字典的键转为字符串，会因与原Date列类型不匹配导致采样结果为空，而直接修改原Date列为字符串又会增加不必要的转换开销。

解决方案 #

无需修改原Date列，可通过以下两种方式实现分组采样：

方法一：利用Date的整数映射适配sampleBy

Date类型在Spark底层可转换为整数（自1970-01-01起的天数），以此作为临时分组键满足sampleBy的类型要求，采样后删除临时列即可保留原Date列：

from datetime import date
from pyspark.sql import Row
import pyspark.sql.functions as F

# 生成测试数据
df_rows = [Row(date = date(2020, 1, d), value=d*2) for d in range(1,31)]
df = spark.createDataFrame(df_rows)

# 生成以日期整数为键的fractions字典
fractions = (df
    .withColumn("date_int", F.datediff(F.col("date"), F.lit(date(1970, 1, 1))))
    .select("date_int")
    .distinct()
    .withColumn("fraction", F.lit(1.0))
    .rdd.collectAsMap()
)

# 执行采样并清理临时列
sampled_df = (df
    .withColumn("date_int", F.datediff(F.col("date"), F.lit(date(1970, 1, 1))))
    .sampleBy("date_int", fractions, seed=42)
    .drop("date_int")
)

# 验证采样结果
sampled_df.groupBy("date").count().show(5)

方法二：用窗口函数绕开sampleBy实现分组采样

若无需为不同日期设置自定义采样比例，可直接用窗口函数+随机数实现按Date分组采样，无需依赖sampleBy：

# 场景1：每个日期随机保留1条数据
sampled_df = (df
    .withColumn("rand", F.rand(seed=42))
    .withColumn("row_num", F.row_number().over(F.window.partitionBy("date").orderBy("rand")))
    .filter(F.col("row_num") == 1)
    .drop("rand", "row_num")
)

# 场景2：每个日期按固定比例（如50%）采样
sampled_df = (df
    .withColumn("rand", F.rand(seed=42))
    .filter(F.col("rand") <= 0.5)
    .groupBy("date")
    .agg(F.collect_list(F.struct("value", "rand")).alias("items"))
    .withColumn("sampled_item", F.expr("array_shuffle(items)[0]"))
    .select("date", "sampled_item.value")
)

说明 #

• 方法一完全保留原Date列结构，仅通过临时整数列适配sampleBy的类型限制，性能开销极低；
• 方法二更灵活，适合统一比例的分组采样场景，无需提前生成fractions字典。

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