Spark SQL带WHERE子句的查询如何正确获取指定行数的采样结果

问题根因 #

你使用的TABLESAMPLE (100 rows)语法本身没有问题，返回结果仅包含首个分区数据是Spark SQL该语法的执行逻辑导致的：

Spark的TABLESAMPLE (n ROWS)采用逐分区扫描逻辑，扫描到第一个符合条件的分区后就会开始取数，只要当前分区的数据量足够覆盖你需要的n行，就会直接终止后续分区的扫描，因此你只会拿到2021-09-14分区的记录。

可行解决方案 #

方案1：全量随机采样（适合中小数据量，采样最均匀） #

直接对过滤后的结果集随机排序后取前n行，能保证采样数据覆盖所有符合条件的分区：

SELECT * FROM table 
WHERE ts IN ('2021-09-14', '2021-09-15')
ORDER BY RAND()
LIMIT 100

方案2：百分比采样+取数（适合大数据量，性能更优） #

如果过滤后的数据量极大，全量随机排序的shuffle开销太高，可以先按比例采样所有分区的数据，再取需要的行数：

SELECT * FROM (
  SELECT * FROM table 
  WHERE ts IN ('2021-09-14', '2021-09-15')
  TABLESAMPLE (0.1 PERCENT)
) t
LIMIT 100

提示：百分比数值可根据总数据量调整，只要保证采样后的总数据量大于你需要的n行即可。

方案3：分区均匀采样（需要每个分区都有采样数据时使用） #

如果要求每个符合条件的分区都有等量的采样数据，可以用窗口函数给每个分区的行打随机排序的序号，再按数量过滤：

WITH partition_ranked AS (
  SELECT *,
         ROW_NUMBER() OVER (PARTITION BY ts ORDER BY RAND()) AS rn
  FROM table
  WHERE ts IN ('2021-09-14', '2021-09-15')
)
SELECT * EXCEPT(rn) FROM partition_ranked WHERE rn <= 50

上述示例会从2021-09-14、2021-09-15两个分区各取50行，合计100行采样数据。

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