In PySpark, how can I populate a new column based on a lookup in another DataFrame?

user_11661982 0
Viewed 1884

New to Spark and PySpark, I am trying to add a field / column in a DataFrame by looking up information in another DataFrame. I have spent the past several hours trying to read up on RDDs, DataFrames, DataSets, maps, joins, etc. but the concepts are all still new to me and I am still having a hard time making heads or tails of it all.

So for example, given the following input DataFrame with IDs and roles:

+----+-------------+
| ID | Role        |
+----+-------------+
|  1 | Author      |
|  1 | Editor      |
|  2 | Author      |
|  2 | Publisher   |
|  3 | Editor      |
|  3 | Assistant   |
+----+-------------+

And another input DataFrame with a unique mapping of IDs and names:

+----+-------------+
| ID | Name        |
+----+-------------+
|  1 | John Smith  |
|  2 | John Doe    |
|  3 | Bob Jim Bob |
+----+-------------+

I need to somehow merge this data into a single output DataFrame with the Name column added like so:

+----+-------------+-------------+
| ID | Name        | Role        |
+----+-------------+-------------+
|  1 | John Smith  | Author      |
|  1 | John Smith  | Editor      |
|  2 | John Doe    | Author      |
|  2 | John Doe    | Publisher   |
|  3 | Bob Jim Bob | Editor      |
|  3 | Bob Jim Bob | Assistant   |
+----+-------------+-------------+

The actual volume of the data is on the order of 200k+ rows for both input tables.

My initial thoughts were to add a new Name column to the ID + Role table, then iterate through each row and manually look up the Name from the ID + Name table but that is really slow and I am sure fails to capitalize on the strengths of Spark.

Again, still really new to all this so please forgive me if this is a really obvious question.

Thank you!

pyspark apache-spark-sql apache-spark python
1 Answers

You can merge both dataframeS as below in pySpark:

>>> df1.show()
+---+---------+
| ID|     Role|
+---+---------+
|  1|   Author|
|  1|   Editor|
|  2|   Author|
|  2|Publisher|
|  3|   Editor|
|  3|Assistant|
+---+---------+

>>> df2.show()
+---+-----------+
| ID|       Name|
+---+-----------+
|  1| John Smith|
|  2|   John Doe|
|  3|Bob Jim Bob|
+---+-----------+

>>> df3 = df2.join(df1,"ID")
>>> df3.show()
+---+-----------+---------+
| ID|       Name|     Role|
+---+-----------+---------+
|  1| John Smith|   Author|
|  1| John Smith|   Editor|
|  2|   John Doe|   Author|
|  2|   John Doe|Publisher|
|  3|Bob Jim Bob|   Editor|
|  3|Bob Jim Bob|Assistant|
+---+-----------+---------+

Note: I am assuming "ID" as foreign Key, Please let comment out in case any concerns.

user_115923860
Related