spark-如何根据列值的百分比选择随机行

Question

有没有办法使用 spark sql 根据列的分布来选择随机样本？例如对于下面的数据框，我想选择总共 6 行但大约 2 行 prod_name = A 和 2 行 prod_name = B 和 2 行 prod_name = C ，因为它们各占 1/3数据？请注意，每种产品并不总是占 1/3%。这只是一个例子。非常感谢您的帮助。

prod_name | value
----------------------
   A      | 100    
   A      | 200.    
   A      | 300
   A      | 400
   B      | 500
   B      | 600     
   B      | 650   
   B      | 700   
   C      | 500
   C      | 600     
   C      | 650   
   C      | 700  

Answer 1

使用sampleBy 应该可以做到。每个prod_name 所需的分数可以通过将预期行数除以实际行数来计算：

df=...
prods=df.select('prod_name').distinct().collect()
cnt=df.count()
expected_rows=6 # change this number to adjust of rows in  the result
fractions={r['prod_name']:expected_rows/cnt for r in prods}
df.stat.sampleBy('prod_name', fractions).show()

输出：

+---------+-----+
|prod_name|value|
+---------+-----+
|        A|  200|
|        A|  400|
|        B|  500|
|        B|  700|
|        C|  500|
|        C|  700|
+---------+-----+

结果的大小可能与expected_rows 的数量不完全匹配，因为抽样涉及随机操作。结果会有所不同。