Случайно выборочный кадр данных Pyspark с условиями столбца

Question

Я пытаюсь случайным образом выбрать фрейм данных Pyspark, где значение столбца соответствует определенному условию. Я хотел бы использовать метод sample для случайного выбора строк на основе значения столбца. Допустим, у меня есть следующий фрейм данных:

+---+----+------+-------------+------+
| id|code|   amt|flag_outliers|result|
+---+----+------+-------------+------+
|  1|   a|  10.9|            0|   0.0|
|  2|   b|  20.7|            0|   0.0|
|  3|   c|  30.4|            0|   1.0|
|  4|   d| 40.98|            0|   1.0|
|  5|   e| 50.21|            0|   2.0|
|  6|   f|  60.7|            0|   2.0|
|  7|   g|  70.8|            0|   2.0|
|  8|   h| 80.43|            0|   3.0|
|  9|   i| 90.12|            0|   3.0|
| 10|   j|100.65|            0|   3.0|
+---+----+------+-------------+------+

Я хотел бы выбрать только 1 (или любое определенное количество) каждого из 0, 1, 2, 3 на основе столбца result, поэтому я бы закончилс этим:

+---+----+------+-------------+------+
| id|code|   amt|flag_outliers|result|
+---+----+------+-------------+------+
|  1|   a|  10.9|            0|   0.0|
|  3|   c|  30.4|            0|   1.0|
|  5|   e| 50.21|            0|   2.0|
|  8|   h| 80.43|            0|   3.0|
+---+----+------+-------------+------+

Есть ли хороший программный способ для достижения этой цели, то есть взять одинаковое количество строк для каждого из значений, заданных в определенном столбце? Любая помощь очень ценится!

Answer 1

Вы можете использовать samplyBy(), который возвращает стратифицированный образец без замены на основе доли, заданной для каждого слоя.

>>> from pyspark.sql.functions import col
>>> dataset = sqlContext.range(0, 100).select((col("id") % 3).alias("result"))
>>> sampled = dataset.sampleBy("result", fractions={0: 0.1, 1: 0.2}, seed=0)
>>> sampled.groupBy("result").count().orderBy("key").show()

+------+-----+
|result|count|
+------+-----+
|     0|    5|
|     1|    9|
+------+-----+