如何为具有True，False和Null值的布尔字段创建分层样本？

问题描述

我有一个带有布尔字段的DataFrame。

df = spark.createDataFrame([
  [True],[False],[None],[True],[None]
]).toDF("match")

我想创建一个具有相等的True，False和Null值的分层样本（PySpark）。

我还如何获取示例中的Null值（不接受None: 0.3）

sampled = df.sampleBy("match",fractions={True: 0.3,False: 0.3}

解决方法

基于sampleBy方法的source code，参数fractions是Map[T,Double]，对于包含Spark的MapType列，不允许空键（请参阅doc）

def sampleBy[T](col: String,fractions: Map[T,Double],seed: Long): DataFrame = {
  sampleBy(Column(col),fractions,seed)
}

一种可能的解决方案是添加一个标志，以将False，True和NULL转换为0，1，2，然后根据此标志执行 sampleBy ，例如：

from pyspark.sql.functions import expr

df_sample = df.withColumn('flag',expr("coalesce(int(match),2)")) \
    .sampleBy("flag",{0:0.3,1:0.3,2:0.3}) \
    .drop("flag")