如何从 DataFrame 中选择确切数量的随机行答案

【问题标题】：How to select an exact number of random rows from DataFrame如何从 DataFrame 中选择确切数量的随机行
【发布时间】：2017-03-20 03:45:11
【问题描述】：

如何有效地从 DataFrame 中选择 exact 数量的随机行？数据包含一个可以使用的索引列。如果我必须使用最大大小，索引列上的 count() 或 max() 哪个更有效？

【问题讨论】：

你不能只使用df.sample() 吗？
@mtoto sample() 返回一个近似数字，但算法在某些情况下要求一个准确数字。

标签： apache-spark random spark-dataframe

【解决方案1】：

这是使用 Pandas DataFrame.Sample 方法的替代方法。这使用 spark applyInPandas 方法分发组，可从 Spark 3.0.0 获得。这允许您选择每组的确切行数。

我已将args 和kwargs 添加到函数中，以便您可以访问DataFrame.Sample 的其他参数。

def sample_n_per_group(n, *args, **kwargs):
    def sample_per_group(pdf):
        return pdf.sample(n, *args, **kwargs)
    return sample_per_group

df = spark.createDataFrame(
    [
        (1, 1.0), 
        (1, 2.0), 
        (2, 3.0), 
        (2, 5.0), 
        (2, 10.0)
    ],
    ("id", "v")
)

(df.groupBy("id")
   .applyInPandas(
        sample_n_per_group(2, random_state=2), 
        schema=df.schema
   )
)

要了解大型组的限制，来自documentation：

此功能需要完全随机播放。一个组的所有数据将是加载到内存中，因此用户应注意潜在的 OOM 如果数据有偏差并且某些组太大而无法容纳，则存在风险记忆。

另见此处： How take a random row from a PySpark DataFrame?

【讨论】：

【解决方案2】：

一种可能的方法是使用.count() 计算行数，然后使用python 的random library 中的sample() 从该范围内生成任意长度的随机序列。最后使用生成的数字列表vals 对您的索引列进行子集化。

import random 
def sampler(df, col, records):

  # Calculate number of rows
  colmax = df.count()

  # Create random sample from range
  vals = random.sample(range(1, colmax), records)

  # Use 'vals' to filter DataFrame using 'isin'
  return df.filter(df[col].isin(vals))

例子：

df = sc.parallelize([(1,1),(2,1),
                     (3,1),(4,0),
                     (5,0),(6,1),
                     (7,1),(8,0),
                     (9,0),(10,1)]).toDF(["a","b"])

sampler(df,"a",3).show()
+---+---+
|  a|  b|
+---+---+
|  3|  1|
|  4|  0|
|  6|  1|
+---+---+

【讨论】：

感谢您的建议。这也是我遇到的。我不想使用这个解决方案的原因是使用 count() 方法，非常昂贵。
你也可以缓存你的df然后在函数外计算count()，或者使用agg(max)。
谢谢，在 Java 中使用了您的解决方案。