spark-groupBy 和 sampleBy答案

【问题标题】：spark- groupBy together with sampleByspark-groupBy 和 sampleBy
【发布时间】：2021-11-06 22:08:16
【问题描述】：

我希望使用sampleBy根据一列的分布来获取样本。例如，在每个prod_name 组中，我想在colour 列sampleBy("colour", fractions ={"blue":0.5, "yellow",0.1, green: 0.3} 的基础上做一个sampleBy，如何将这两种方法结合使用？非常感谢您的帮助！

    prod_name | colour | value   | code
 -------------------------------
   A      | blue    |100     | Y  
   A      | blue    |200.    | N
   A      | blue.   |300.    | Y 
   A      | blue.   |400.    | Y 
   A      | yellow. |500.    | N 
   B      | green.  |600     | Y 
   B      | green.  |650     | Y 
   B      | blue.   |700     | N
   C      | red.    |800.    | Y
   C      | blue    |900.    | N 
   C      | green   |1000    | N

【问题讨论】：

如果prod_names 在原始数据中分布不均匀，预期的结果是什么？假设原始数据包含 10 个As（颜色不同）和 1000 个Bs。结果应该包含（大约）相同数量的As 和Bs，还是应该反映As 和Bs 的原始分布？

标签： python dataframe apache-spark pyspark apache-spark-sql

【解决方案1】：

这个方法有点棘手，但应该能正确完成它的工作。
基本上，我们将构建一个新列，它是prod_name 和colour 的串联，因此我们可以在该列上使用sampleBy。我们将为找到的颜色创建一个包含重复值的新字典。

# collect distinct values
list_prod = df.select('prod_name').distinct().rdd.map(lambda r: r[0]).collect()
list_colours = df.select('colour').distinct().rdd.map(lambda r: r[0]).collect()

# cartesian product of lists
list_combined = [a + '_' + b for a in list_prod for b in list_colours]

# original dictionary
fractions = {'blue': 0.5, 'yellow': 0.1, 'green': 0.3, 'red': 0.8}

# create new dictionary with repeated numbers
new_dict = {e: fractions[k] for e in list_combined for k in fractions.keys() if k in e}

df \
  .withColumn('combined', F.concat_ws('_', 'prod_name', 'colour')) \
  .sampleBy('combined', fractions=new_dict, seed=42) \
  .show()

+---------+------+------+----+--------+
|prod_name|colour| value|code|combined|
+---------+------+------+----+--------+
|        B| green| 600.0|   Y| B_green|
|        C|   red| 800.0|   Y|   C_red|
|        C|  blue| 900.0|   N|  C_blue|
|        C| green|1000.0|   N| C_green|
+---------+------+------+----+--------+

现在行数很少，因此结果可能很奇怪。在更大的数据框上尝试一下，它应该可以工作。

【讨论】：