Как случайным образом выбирать строки в кадре данных Spark, в то время как условие, основанное на столбце, должно выполняться

Question

Допустим, у нас есть Spark-фрейм данных df со столбцом col, значения которого в этом столбце только 0 и 1. Как мы можем выбрать все строки, где col == 1, а также 50% строк, где col == 0?50% населения с col == 0 должны быть выбраны случайным образом.

Метод выборки допускает случайный выбор 50% строк, но никакие другие условия не могут быть наложены.

Решение, которое я в настоящее время решаюесть следующее, что кажется мне немного уродливым.Интересно, есть ли лучшее решение?

    from pyspark.sql import functions as F
    df = df.withColumn('uniform', F.rand())
    df = df.filter((df.uniform<0.5) | (df.col==1))

Answer 1

Это не гарантирует ровно 50%, но этого должно быть достаточно при достаточно большом наборе данных.

df.where($"col" == 1 or rand() > rand())

примечание: это будет возвращать различный набор случайных строк каждый раз, когда вычисляется набор данных / набор данных. Чтобы исправить это, добавьте квалификацию rand() > rand() в качестве столбца в DF, т.е. df.withColumn("lucky", rand() > rand())