Pyspark groupBy and consolidatng on multiple distinct column values

Viewed 201

Trying to extract records with latest date for distinct values of column A and column B (below)

enter image description here

Ideal Result:

enter image description here

Current Solution:

from pyspark.sql import functions as f
test = df.groupBy(df['A'], df['B']).agg(f.first(df['C']), f.first(df['D']), f.max(df['E']))

Any loopholes to look for or suggestiones to optimize the above ?

2 Answers

Use Window function with max on date column and use it to filter.

from pyspark.sql import functions as F
from pyspark.sql.window import Window


w=Window().partitionBy("A","B")

df1.withColumn("max", F.max(F.to_date("E","yyyyMMdd")).over(w))\
  .filter(F.to_date(F.col("E"),"yyyyMMdd")==F.col("max")).drop("max").show()

If E(date) column is not of StringType, use this:

w=Window().partitionBy("A","B")

df.withColumn("max", F.max(F.to_date(F.col("E").cast('string'),"yyyyMMdd")).over(w))\
  .filter(F.to_date(F.col("E").cast('string'),"yyyyMMdd")==F.col("max")).drop("max").show()

Output:

#+---+---+----+---+--------+
#|  A|  B|   C|  D|       E|
#+---+---+----+---+--------+
#| 12|ERP|7500|  D|20200330|
#| 12|ERF|4500|  D|20200430|
#+---+---+----+---+--------+

You can sort on E column by descending then use row_number function to extract only the latest data.

df.show()
#+---+---+----+---+--------+
#|  A|  B|   C|  D|       E|
#+---+---+----+---+--------+
#| 12|ERP|1000|  M|20200130|
#| 12|ERP|2000|  M|20200228|
#| 12|ERP|7500|  D|20200330|
#| 12|ERF|4500|  D|20200430|
#| 12|ERF|4000|  L|20200228|
#| 12|ERF|3400|  L|20200330|
#+---+---+----+---+--------+
from pyspark.sql.functions import *
from pyspark.sql import *

w=Window.partitionBy("A","B").orderBy(col("Z").desc())

df.withColumn("z",to_date(col("E"),"yyyyMMdd")).\
withColumn("rn",row_number().over(w)).\
filter(col("rn") == 1).\
drop(*['z','rn']).\
show()
#+---+---+----+---+--------+
#|  A|  B|   C|  D|       E|
#+---+---+----+---+--------+
#| 12|ERP|7500|  D|20200330|
#| 12|ERF|4500|  D|20200430|
#+---+---+----+---+--------+
Related