spark groupBy与sampleBy一起使用

Question

我希望使用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

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

现在行数很少，所以结果可能会很奇怪。在更大的数据帧上尝试一下，应该可以