Spark SQL group, map reduce

Viewed 371

I have the following dataset with the name 'data':

+---------+-------------+------+
|   name  |      subject| mark |
+---------+-------------+------+
|     Anna|         math|    80|
|     Vlad|      history|    67|
|     Jack|          art|    78|
|    David|         math|    71|
|   Monica|          art|    65|
|     Alex|          lit|    59|
|     Mark|         math|    82|
+---------+-------------+------+

I would like to do a map-reduce job.

The result show look like this or similar:

Anna, David : 1
Anna, Mark : 1
David, mark: 1
Vlad, None : 1
Jack, Monica: 1
Alex, None : 1

I have tried to do the following:

new_data = data.select(['name', 'subject']).show()

+---------+-------------+
|   name  |      subject| 
+---------+-------------+
|     Anna|         math|  
|     Vlad|      history|  
|     Jack|          art|   
|    David|         math|   
|   Monica|          art|    
|     Alex|          lit|    
|     Mark|         math|    
+---------+-------------+


data_new.groupBy('name','subject').count().show(10)

However, this command does not give what I need.

2 Answers

You can do a self left join using the subject, get the distinct pairs, and add a column of 1.

import pyspark.sql.functions as F

result = df.alias('t1').join(df.alias('t2'),
    F.expr("t1.subject = t2.subject and t1.name != t2.name"), 
    'left'
).select(
    F.concat_ws(
        ', ',
        F.greatest('t1.name', F.coalesce('t2.name', F.lit('None'))),
        F.least('t1.name', F.coalesce('t2.name', F.lit('None')))
    ).alias('pair')
).distinct().withColumn('val', F.lit(1))

result.show()
+------------+---+
|        pair|val|
+------------+---+
|  Alex, None|  1|
| Anna, David|  1|
|  Anna, Mark|  1|
|  None, Vlad|  1|
| David, Mark|  1|
|Jack, Monica|  1|
+------------+---+

The process could be:

  1. Grouping student with the same subject in an array
  2. Call a udf function to create the array items permutation
  3. Add a column that shows numbers for each subject
  4. Call explode function to create separate3 rows for each item in the array

Let's do the steps one by one: Step 1: Grouping

import pyspark.sql.functions as F

grouped_df = data_new.groupBy('subject').agg(F.collect_set('name').alias('students_array'))

Step 2: udf function

from itertools import permutations
def permutatoin(df_col):
    result = sorted([e for e in set(permutations(df_col))])
    return result 
spark.udf.register("perWithPython", permutatoin)
grouped_df  = grouped_df.select('*', permutatoin('students_array'))

Step 3: Create a new digit value column for each subject

grouped_df = grouped_df .withColumn('subject_no', F.rowNumber().over(Window.partitionBy('subject'))

Step 4: create separate rows

grouped_df.select(grouped_df.subject_no, explode(grouped_df.students_array)).show(truncate=False)
Related