Find new, left and existing records in PySpark - pyspark

Question

I have a dataframe with data like this
I want to compare each record with next year's record and see if that id is there. If it's there in the year and the next year it's 'Existing', If it's there in the year and not there in the next year, it's 'Left'. If it's not there in the year but there in the next year, it's 'New'. I want output like this below. The columns 2017-18, 2018-19 etc. should be created dynamically.
How do I achieve this?
After getting the data in the above format, I need to aggregate the sales for each year band like this as below. For example for 2017-2018,
New_sales = sum of all sales of 2018 (which is the later year in 2017-2018) where it's marked as 'New' which is 25 here.
Left_sales = sum of all sales of 2017 (the earlier year in 2017-2018) where it's marked as 'Left' which is 100 here.
Existing_sales = sum of sales of 2017 where it's marked as 'Existing' subtract sum of sales of 2018 where it's marked as 'Existing'
Existing_sales = 50+75 (sales of 2017, 'Existing') - (20+50) (sales of 2018, 'Existing') = 125-70 = 55
How do I achieve this?

Answer 1

As your date is a string, I think you can:
df = spark.createDataFrame(
data=[
(1, '31/12/2017'),
(2, '31/12/2017'),
(3, '31/12/2017'),
(1, '31/12/2018'),
(3, '31/12/2018'),
(5, '31/12/2018'),
],
schema=['id', 'date']
)
First you can get the year first:
df2 = df.withColumn('year', func.split(func.col('Date'), '/').getItem(2))
df2.show(10, False)
+---+----------+----+
|id |date |year|
+---+----------+----+
|1 |31/12/2017|2017|
|2 |31/12/2017|2017|
|3 |31/12/2017|2017|
|1 |31/12/2018|2018|
|3 |31/12/2018|2018|
|5 |31/12/2018|2018|
+---+----------+----+
Then you can collect the list of year by id as a reference:
df3 = df2.groupby('id')\
.agg(func.collect_set('year').alias('year_lst'))
df3.show(3, False)
+---+------------+
|id |year_lst |
+---+------------+
|1 |[2017, 2018]|
|2 |[2017] |
|3 |[2017, 2018]|
+---+------------+
Then you can join the reference back to the data:
df4 = df2.join(df3, on='id', how='left')
df4.show(10, False)
+---+----------+----+------------+
|id |date |year|year_lst |
+---+----------+----+------------+
|1 |31/12/2017|2017|[2017, 2018]|
|2 |31/12/2017|2017|[2017] |
|3 |31/12/2017|2017|[2017, 2018]|
|1 |31/12/2018|2018|[2017, 2018]|
|3 |31/12/2018|2018|[2017, 2018]|
|5 |31/12/2018|2018|[2018] |
+---+----------+----+------------+
The last step is to create the column dynamically. I think you can use a for loop:
year_loop = ['2017', '2018', '2019', '2020', '2021']
for idx in range(len(year_loop)-1):
this_year = year_loop[idx]
next_year = year_loop[idx+1]
column_name = f"{this_year}-{next_year}"
new_condition = (~func.array_contains(func.col('year_lst'), this_year)) & (func.array_contains(func.col('year_lst'), next_year))
exist_condition = (func.array_contains(func.col('year_lst'), this_year)) & (func.array_contains(func.col('year_lst'), next_year))
left_condition = (func.array_contains(func.col('year_lst'), this_year)) & (~func.array_contains(func.col('year_lst'), next_year))
df4 = df4.withColumn(column_name, func.when(new_condition, func.lit('New'))
.when(exist_condition, func.lit('Existing'))
.when(left_condition, func.lit('Left')))
df4.show(10, False)
+---+----------+----+------------+---------+---------+---------+---------+
|id |date |year|year_lst |2017-2018|2018-2019|2019-2020|2020-2021|
+---+----------+----+------------+---------+---------+---------+---------+
|1 |31/12/2017|2017|[2017, 2018]|Existing |Left |null |null |
|2 |31/12/2017|2017|[2017] |Left |null |null |null |
|3 |31/12/2017|2017|[2017, 2018]|Existing |Left |null |null |
|1 |31/12/2018|2018|[2017, 2018]|Existing |Left |null |null |
|3 |31/12/2018|2018|[2017, 2018]|Existing |Left |null |null |
|5 |31/12/2018|2018|[2018] |New |Left |null |null |
+---+----------+----+------------+---------+---------+---------+---------+