Performing an Excel-like "vlookup" method in Pyspark

Viewed 1301

I'm new to Pyspark I would like some help with this. I have a Pyspark dataframe df1 like the one below:

df1 =
|---------------------|------------------|------------------|
|     ID_Machine      |  Event_Duration  |     Timestamp    |
|---------------------|------------------|------------------|
|          1          |         34       |        213       |
|---------------------|------------------|------------------|
|          1          |         97       |        572       |
|---------------------|------------------|------------------|
|          1          |         78       |        872       |
|---------------------|------------------|------------------|
|          2          |         83       |        345       |
|---------------------|------------------|------------------|
|          2          |         14       |        718       |
|---------------------|------------------|------------------|
|          2          |         115      |        884       |
|---------------------|------------------|------------------|

From it, I have to perform a groupBy with an aggregate method:

import pyspark.sql.functions as F

df2 = df1.groupBy("ID_Machine").agg(F.max("Event_duration").alias("Max_Event_Duration")

Thus obtaining:

df2 = 
|---------------------|---------------------------|
|      ID_Machine     |     Max_Event_Duration    |
|---------------------|---------------------------|
|           1         |             97            |
|---------------------|---------------------------|
|           2         |             115           |
|---------------------|---------------------------|

So far, so good. However, now I would like to perform some sort of fuction like vlookup in Excel, where I retrieve the Timestamp value in df1 corresponding to the Max_Event_Duration in df2, obtaining something like the following:

|---------------------|---------------------|------------------|
|     ID_Machine      |  Max_Event_Duration |     Timestamp    |
|---------------------|---------------------|------------------|
|          1          |          97         |        572       |
|---------------------|---------------------|------------------|
|          2          |          115        |        884       |
|---------------------|---------------------|------------------|

Does anybody know how to create this third dataframe or how to modify the code which creates df2 in order to include the respective Timestamp value?

Thanks!

2 Answers

You can make a new column of max_event_duration and filter the columns where event_duration = max_event_duration

from pyspark.sql import functions as F, Window

df2 = df1.select('*', F.max("Event_duration").over(Window.partitionBy("ID_Machine")).alias("Max_Event_Duration")) \
         .filter("Event_duration = Max_Event_Duration") \
         .drop("Max_Event_Duration")

You would have to first create a dataframe with the max event named same as in the OG dataframe

import pyspark.sql.functions as F

df_max = df1.groupBy("ID_Machine").agg(F.max("Event_Duration").alias("Event_Duration"))
df_max.show()

which would give you df:

+----------+--------------+
|ID_Machine|Event_Duration|
+----------+--------------+
|         1|            97|
|         2|           115|
+----------+--------------+

Then join the dataframe by the two similar named columns and then rename the event duration

df_combined = df_max.join(df1, ["ID_Machine", "Event_Duration"]).withColumnRenamed("Event_Duration", "Max_Event_Duration")
df_combined.show()

which would give you:

+----------+------------------+---------+
|ID_Machine|Max_Event_Duration|Timestamp|
+----------+------------------+---------+
|         2|               115|      884|
|         1|                97|      572|
+----------+------------------+---------+
Related