SPARK - Joining 2 dataframes on values in an array

Question

I can't find an easy and elegant solution to this one.

I have a df1 with this column :

|-- guitars: array (nullable = true)
 |    |-- element: long (containsNull = true)

I have a df2 made of guitars, and an id matching with the Long in my df 1.

root
 |-- guitarId: long (nullable = true)
 |-- make: string (nullable = true)
 |-- model: string (nullable = true)
 |-- type: string (nullable = true)

I want to join my two dfs, obviously, and instead of having an array of long, I want an array of struct guitars from df2.

I'm using array_contains() to join the two dfs, but spark is exploding the array of n Long in the df1 in n rows in the result df.

before

|   2|Eric Clapton| [1, 5]|               [,,,]|

after

|   2|Eric Clapton| [1, 5]|               [,,,]|       5|Fender|Stratocaster|            Electric|
|   2|Eric Clapton| [1, 5]|               [,,,]|       1|Gibson|          SG|            Electric|

What would be the most elegant solution to transform this array column of Long into an array column of struct from an other dataframe ?

ideal

|   2|Eric Clapton|[[Fender, Stratocaster, Electric],[Gibson, SG, Electric]]|               [,,,]|

Thanks in advance

(first question btw, be humble :P)

Is the data model correct? df1 guitars contains many ids and one/many of them should match? You might be having a cartesian join. — Salim

werner werner · Accepted Answer · 2019-12-30T19:09:35

array_contains() works, and you only have to group the result by the player afterwards.

Lets start with two datasets, one for the players and one for the guitars:

val player = Seq(("Eric Clapton", Array(1,5)), ("Paco de Lucia", Array(1,2)), ("Jimi Hendrix", Array(3))).toDF("player", "guitars")
val guitar = Seq((1, "Gibson", "SG", "Electric"), (2, "Faustino Conde", "Media Luna", "Acoustic"), (3, "Pulsebeatguitars", "Spider", "Electric"), (4, "Yamaha", "FG800", "Acoustic"), (5, "Fender", "Stratocaster", "Electric")).toDF("guitarId", "make", "model", "type")

+-------------+-------+
|       player|guitars|
+-------------+-------+
| Eric Clapton| [1, 5]|
|Paco de Lucia| [1, 2]|
| Jimi Hendrix|    [3]|
+-------------+-------+
+--------+----------------+------------+--------+
|guitarId|            make|       model|    type|
+--------+----------------+------------+--------+
|       1|          Gibson|          SG|Electric|
|       2|  Faustino Conde|  Media Luna|Acoustic|
|       3|Pulsebeatguitars|      Spider|Electric|
|       4|          Yamaha|       FG800|Acoustic|
|       5|          Fender|Stratocaster|Electric|
+--------+----------------+------------+--------+

To make the grouping operation a bit easier, the idea is to combine the three columns of the guitar dataset into a struct before the join:

val guitar2 = guitar.withColumn("guitar", struct('make, 'model, 'type))

After the join, we group the result by the player and get the correct result:

player.join(guitar2, expr("array_contains(guitars, guitarId)"))
  .groupBy("player")
  .agg(collect_list('guitar))
  .show(false)

prints

+-------------+----------------------------------------------------------------+
|player       |collect_list(guitar)                                            |
+-------------+----------------------------------------------------------------+
|Jimi Hendrix |[[Pulsebeatguitars, Spider, Electric]]                          |
|Eric Clapton |[[Gibson, SG, Electric], [Fender, Stratocaster, Electric]]      |
|Paco de Lucia|[[Gibson, SG, Electric], [Faustino Conde, Media Luna, Acoustic]]|
+-------------+----------------------------------------------------------------+

SPARK - Joining 2 dataframes on values in an array

1 Answers