Spark - How to group data by multiple keys with or relationship

user_2197820 0
Viewed 195

I have a RDD like the following:

[
    Row(device='1', token='aaa', other_value=1),  
    Row(device='1', token='bbb', other_value=1), 
    Row(device='2', token='bbb', other_value=1), 
    Row(device='4', token='ddd', other_value=1), 
    ...
]

In order to do some processing (not for aggregation), I need to group rows with the same device or token together. For example, first 3 rows above should be grouped together for further processing. (row 1 and 2 have the same device 1, row 2 and 3 have the same token bbb).

Essentially, I would like to do something like

rdd.groupBy(same device or same token)

Is it possible in spark?

Update:

The expected output would be a RDD like this for the above example:

[
    [
        Row(device='1', token='aaa', other_value=1),  
        Row(device='1', token='bbb', other_value=1), 
        Row(device='2', token='bbb', other_value=1)
    ],
    [
        Row(device='4', token='ddd', other_value=1)
    ]
]
pyspark apache-spark
1 Answers

The property to which group a row belongs is transitive. Therefore a simple groupBy won't work. One option would be to interpret the data as graph and use GraphFrames to find the connected components in the graph. Each row of the input data would be an edge of the graph with source = device and destination = token.

#setup GraphFrames
import os
os.environ["PYSPARK_SUBMIT_ARGS"] = (
    "--packages graphframes:graphframes:0.8.1-spark3.0-s_2.12 pyspark-shell"
)

spark = ...

#the input data
rdd = spark.sparkContext.parallelize([   Row(device='1', token='aaa', other_value=1),  
    Row(device='1', token='bbb', other_value=1), 
    Row(device='2', token='bbb', other_value=1),     
    Row(device='4', token='ddd', other_value=1)])

df = spark.createDataFrame(rdd)

edges = df \
    .withColumnRenamed("device", "src") \
    .withColumnRenamed("token", "dst")

vertices = edges.select("src").distinct() \
    .union(edges.select("dst").distinct()) \
    .withColumnRenamed("src", "id")

#create a graph and find all connected components
g = GraphFrame(vertices, edges)
cc = g.connectedComponents()

df.join(cc.distinct(), df.device == cc.id) \
    .orderBy("component", "device", "token") \
    .show()

Output:

+------+-----+-----------+---+------------+
|device|token|other_value| id|   component|
+------+-----+-----------+---+------------+
|     4|  ddd|          1|  4| 60129542144|
|     1|  aaa|          1|  1|335007449088|
|     1|  bbb|          1|  1|335007449088|
|     2|  bbb|          1|  2|335007449088|
+------+-----+-----------+---+------------+

All rows with the same value for component belong to the same group.

user_21298010
Related