Spark SQL中是否有等价于DataFrame.sampleBy的分层抽样实现方式？

当然可以用原生Spark SQL实现和sampleBy等价的分层抽样逻辑，核心思路是通过关联权重配置+随机数过滤来实现按指定列、指定比例的分层抽取，具体步骤如下：

• 第一步：定义分层权重表
  把sampleBy中用到的权重字典转换成Spark可查询的临时视图或CTE，明确每个分层值对应的抽样比例。

• 第二步：关联原表与权重表
  将目标数据表和权重表按分层列进行关联，让每一行数据都带上对应的抽样比例。

• 第三步：随机数过滤实现抽样
  使用rand()函数生成0到1之间的随机数，只保留随机数小于对应分层权重的行，完成抽样。

示例代码 #

假设我们有一张名为product_data的表，分层列是category，需要按{'electronics': 0.4, 'clothing': 0.6, 'home': 0.2}的比例抽样：

1. 创建权重临时视图：

CREATE OR REPLACE TEMP VIEW sample_weights AS
SELECT 'electronics' AS category, 0.4 AS weight UNION ALL
SELECT 'clothing' AS category, 0.6 AS weight UNION ALL
SELECT 'home' AS category, 0.2 AS weight;

2. 执行分层抽样：

SELECT p.*
FROM product_data p
JOIN sample_weights w ON p.category = w.category
WHERE rand() < w.weight;

补充说明 #

• 如果需要可复现的抽样结果，替换rand()为rand(seed)，比如rand(123)，对应sampleBy方法中的seed参数。
• 该逻辑和sampleBy完全等价：属于无放回分层抽样，每个分层独立按照指定比例抽取数据。
• 对于大表，Spark会自动基于分层列做分区优化，保证抽样效率。

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