Shift rows dynamically based on column value - dataframe

Question

Below is my input dataframe:
+---+----------+--------+
|ID |date |shift_by|
+---+----------+--------+
|1 |2021-01-01|2 |
|1 |2021-02-05|2 |
|1 |2021-03-27|2 |
|2 |2022-02-28|1 |
|2 |2022-04-30|1 |
+---+----------+--------+
I need to groupBy "ID" and shift based on the "shift_by" column. In the end, the result should look like below:
+---+----------+----------+
|ID |date1 |date2 |
+---+----------+----------+
|1 |2021-01-01|2021-03-27|
|2 |2022-02-28|2022-04-30|
+---+----------+----------+
I have implemented the logic using UDF, but it makes my code slow. I would like to understand if this logic can be implemented without using UDF.
Below is a sample dataframe:
from datetime import datetime
from pyspark.sql.types import *
data2 = [(1, datetime.date(2021, 1, 1), datetime.date(2021, 3, 27)),
(2, datetime.date(2022, 2, 28), datetime.date(2022, 4, 30))
]
schema = StructType([
StructField("ID", IntegerType(), True),
StructField("date1", DateType(), True),
StructField("date2", DateType(), True),
])
df = spark.createDataFrame(data=data2, schema=schema)

Answer 1

based on the comments and chats, you can try to calculate first and last values of the lat/lon fields of concern.
import pyspark.sql.functions as func
from pyspark.sql.window import Window as wd
import sys
data_sdf. \
withColumn('foo_first', func.first('foo').over(wd.partitionBy('id').orderBy('date').rowsBetween(-sys.maxsize, sys.maxsize))). \
withColumn('foo_last', func.last('foo').over(wd.partitionBy('id').orderBy('date').rowsBetween(-sys.maxsize, sys.maxsize))). \
select('id', 'foo_first', 'foo_last'). \
dropDuplicates()
OR, you can create structs and take min/max
data_sdf = spark.createDataFrame(
[(1, '2021-01-01', 2, 2),
(1, '2021-02-05', 3, 2),
(1, '2021-03-27', 4, 2),
(2, '2022-02-28', 1, 5),
(2, '2022-04-30', 5, 1)],
['ID', 'date', 'lat', 'lon'])
data_sdf. \
withColumn('dt_lat_lon_struct', func.struct('date', 'lat', 'lon')). \
groupBy('id'). \
agg(func.min('dt_lat_lon_struct').alias('min_dt_lat_lon_struct'),
func.max('dt_lat_lon_struct').alias('max_dt_lat_lon_struct')
). \
selectExpr('id',
'min_dt_lat_lon_struct.lat as lat_first', 'min_dt_lat_lon_struct.lon as lon_first',
'max_dt_lat_lon_struct.lat as lat_last', 'max_dt_lat_lon_struct.lon as lon_last'
)
# +---+---------+---------+--------+--------+
# | id|lat_first|lon_first|lat_last|lon_last|
# +---+---------+---------+--------+--------+
# | 1| 2| 2| 4| 2|
# | 2| 1| 5| 5| 1|
# +---+---------+---------+--------+--------+

Answer 2

Aggregation using min and max seems could work in your case.
from pyspark.sql import functions as F
df = spark.createDataFrame(
[(1, '2021-01-01', 2),
(1, '2021-02-05', 2),
(1, '2021-03-27', 2),
(2, '2022-02-28', 1),
(2, '2022-04-30', 1)],
['ID', 'date', 'shift_by'])
df = df.groupBy('ID').agg(
F.min('date').alias('date1'),
F.max('date').alias('date2'),
)
df.show()
# +---+----------+----------+
# | ID| date1| date2|
# +---+----------+----------+
# | 1|2021-01-01|2021-03-27|
# | 2|2022-02-28|2022-04-30|
# +---+----------+----------+