python - spark- groupBy 和 sampleBy

Question

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

    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

Answer 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|
+---------+------+------+----+--------+

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