从Spark GroupedData对象中选择随机项

时间:2015-11-17 05:44:02

标签: python python-2.7 apache-spark apache-spark-sql

我刚接触在Python中使用Spark并且无法解决此问题:在groupBy上运行pyspark.sql.dataframe.DataFrame

df = sqlsc.read.json("data.json")
df.groupBy('teamId')

如何在不更换的情况下从每个结果组中选择N个随机样本(按teamId分组)?

我基本上是在尝试从每个团队中选择N随机用户,可能使用groupBy开头是错误的吗?

2 个答案:

答案 0 :(得分:15)

嗯,这有点不对劲。 GroupedData并非真正设计用于数据访问。它只描述了分组标准并提供了聚合方法。有关详细信息,请参阅我对Using groupBy in Spark and getting back to a DataFrame的回答。

这个想法的另一个问题是选择N random samples。这是一项很难在没有数据心理分组的情况下并行实现的任务,当call分组在DataFrame上时,不会发生这种情况:

至少有两种方法可以解决这个问题:

  • 转换为RDD,groupBy并执行本地采样

    import random
    
    n = 3
    
    def sample(iter, n): 
        rs = random.Random()  # We should probably use os.urandom as a seed
        return rs.sample(list(iter), n)    
    
    df = sqlContext.createDataFrame(
        [(x, y, random.random()) for x in (1, 2, 3) for y in "abcdefghi"], 
        ("teamId", "x1", "x2"))
    
    grouped = df.rdd.map(lambda row: (row.teamId, row)).groupByKey()
    
    sampled = sqlContext.createDataFrame(
        grouped.flatMap(lambda kv: sample(kv[1], n)))
    
    sampled.show()
    
    ## +------+---+-------------------+
    ## |teamId| x1|                 x2|
    ## +------+---+-------------------+
    ## |     1|  g|   0.81921738561455|
    ## |     1|  f| 0.8563875814036598|
    ## |     1|  a| 0.9010425238735935|
    ## |     2|  c| 0.3864428179837973|
    ## |     2|  g|0.06233470405822805|
    ## |     2|  d|0.37620872770129155|
    ## |     3|  f| 0.7518901502732027|
    ## |     3|  e| 0.5142305439671874|
    ## |     3|  d| 0.6250620479303716|
    ## +------+---+-------------------+
    
  • 使用窗口函数

    from pyspark.sql import Window
    from pyspark.sql.functions import col, rand, rowNumber
    
    w = Window.partitionBy(col("teamId")).orderBy(col("rnd_"))
    
    sampled = (df
        .withColumn("rnd_", rand())  # Add random numbers column
        .withColumn("rn_", rowNumber().over(w))  # Add rowNumber over windw
        .where(col("rn_") <= n)  # Take n observations
        .drop("rn_")  # drop helper columns
        .drop("rnd_"))
    
    sampled.show()
    
    ## +------+---+--------------------+
    ## |teamId| x1|                  x2|
    ## +------+---+--------------------+
    ## |     1|  f|  0.8563875814036598|
    ## |     1|  g|    0.81921738561455|
    ## |     1|  i|  0.8173912535268248|
    ## |     2|  h| 0.10862995810038856|
    ## |     2|  c|  0.3864428179837973|
    ## |     2|  a|  0.6695356657072442|
    ## |     3|  b|0.012329360826023095|
    ## |     3|  a|  0.6450777858109182|
    ## |     3|  e|  0.5142305439671874|
    ## +------+---+--------------------+
    

但我担心两者都会相当昂贵。如果各个组的大小平衡且相对较大,我只需使用DataFrame.randomSplit

如果群组数量相对较少,则可以尝试其他方式:

from pyspark.sql.functions import count, udf
from pyspark.sql.types import BooleanType
from operator import truediv

counts = (df
    .groupBy(col("teamId"))
    .agg(count("*").alias("n"))
    .rdd.map(lambda r: (r.teamId, r.n))
    .collectAsMap()) 

# This defines fraction of observations from a group which should
# be taken to get n values 
counts_bd = sc.broadcast({k: truediv(n, v) for (k, v) in counts.items()})

to_take = udf(lambda k, rnd: rnd <= counts_bd.value.get(k), BooleanType())

sampled = (df
    .withColumn("rnd_", rand())
    .where(to_take(col("teamId"), col("rnd_")))
    .drop("rnd_"))

sampled.show()

## +------+---+--------------------+
## |teamId| x1|                  x2|
## +------+---+--------------------+
## |     1|  d| 0.14815204548854788|
## |     1|  f|  0.8563875814036598|
## |     1|  g|    0.81921738561455|
## |     2|  a|  0.6695356657072442|
## |     2|  d| 0.37620872770129155|
## |     2|  g| 0.06233470405822805|
## |     3|  b|0.012329360826023095|
## |     3|  h|  0.9022527556458557|
## +------+---+--------------------+

在Spark 1.5+中,您可以通过调用udf方法替换sampleBy

df.sampleBy("teamId", counts_bd.value)

它不能给出确切的观察数量,但是大多数时候应该足够好,只要每组的观察数量足够大以获得适当的样本。您也可以以类似的方式在RDD上使用sampleByKey

答案 1 :(得分:0)

我发现这又是一个数据框架,而不是进入rdd方式。

您可以使用window函数在组内创建排名,排名可以是随机的,以适合您的情况。然后,您可以根据每个组所需的样本(N)的数量进行过滤

window_1 = Window.partitionBy(data['teamId']).orderBy(F.rand())
data_1 = data.select('*', F.rank().over(window_1).alias('rank')).filter(F.col('rank') <= N).drop('rank')