spark- groupBy together with sampleBy

user_3735871 0
Viewed 117

I hope to use sampleBy to get samples based on the distribution of a column. For example, within each prod_name group, I'd like to do a sampleBy based on colour column sampleBy("colour", fractions ={"blue":0.5, "yellow",0.1, green: 0.3} How can I combine and use these two methods together? Many thanks for your help!

    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
pyspark apache-spark-sql apache-spark dataframe python
1 Answers

This method is a bit tricky, but should do its job correctly.
Basically, we are going to build a new column which is the concatenation of prod_name and colour, so we can use sampleBy on that column. We will create a new dictionary with repeated values for the found colours.

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

Now rows are in low numbers, so results could be weird. Try it on your larger dataframe, it should work.

user_74654620
Related