Q:

Populating new columns based on conditions in pyspark

I have a data frame in pyspark like below. 
df = sqlContext.createDataFrame(
[
("101", "99.10", "2019-06-04"), 
("102", "89.27", "2019-06-04"), 
("102", "89.10", "2019-03-04"), 
("103", "73.11", "2019-09-10"), 
("101", "-69.81", "2019-09-11"), 
("101", "12.51", "2018-12-14"), 
("101", "43.23", "2018-09-11")
], 
("user_id", "amount", "trans_date"))

I want to find out the following things
1) First find the max of trans_date in the data frame and store as variable.  I got it while doing like below
from pyspark.sql import functions as f
from pyspark.sql import Window
max_date=df.groupby().agg(f.max('trans_date')).collect()[0].asDict()['max(trans_date)']

2) Sum of amount for each user_id
df1=df.withColumn('Balance', f.sum('amount').over(Window.partitionBy('user_id')))

3) Now using this max_date I want to create some columns and populate values for example
`days_0_30` if `trans_date` is between `max_date` and `30 days before
`days_31_60` if `trans_date` is between `max_date - 31 days` and `60 days before
and so on.  I am able to do it using below

df2 = df1.withColumn(days_0_30, f.when((df1.trans_date <= '2019-09-11') & (df1.trans_date >= '2019-06-11'), df1.Balance).otherwise('null')).withColumn(days_31_60, f.when((df1.trans_date <= '2019-06-10') & (df1.trans_date >= '2019-03-11'), df1.Balance).otherwise('null')).withColumn(days_61_90, f.when((df1.trans_date <= '2019-03-10') & (df1.trans_date >= '2018-12-11'), df1.Balance).otherwise('null')).withColumn(days_91_120, f.when((df1.trans_date <= '2018-12-10') & (df1.trans_date >= '2018-09-11'), df1.Balance).otherwise('null')).withColumn(days_121_150, f.when((df1.trans_date <= '2018-09-10') & (df1.trans_date >= '2018-06-11'), df1.Balance).otherwise('null'))

If date is not in range defined values should be null
But I have hard coded dates in the above.  I would like to have this dynamically achieve my result. 
What should I do in order to achieve my result in a more elegant way

A:

Just use list comprehensions:
Follow your steps to get df1 (make sure the column trans_date is DateType())
from pyspark.sql import functions as f, Window
from datetime import timedelta

df = df.withColumn('trans_date', f.to_date('trans_date'))

max_date = df.select(f.max('trans_date').alias('max_date')).first().max_date
# datetime.date(2019, 9, 11)

df1 = df.withColumn('balance', f.round(f.sum('amount').over(Window.partitionBy('user_id')),2))

Then use list comprehension to create a list of tuple with three elements (range_name, range_start_date, range_end_date)
dranges = [
    ('days_{}_{}'.format(i*30+1,(i+1)*30), max_date-timedelta(days=(i+1)*30), max_date-timedelta(days=i*30+1))
         for i in range(5)
]
dranges 
#[('days_1_30', datetime.date(2019, 8, 12), datetime.date(2019, 9, 10)),
# ('days_31_60', datetime.date(2019, 7, 13), datetime.date(2019, 8, 11)),
# ('days_61_90', datetime.date(2019, 6, 13), datetime.date(2019, 7, 12)),
# ('days_91_120', datetime.date(2019, 5, 14), datetime.date(2019, 6, 12)),
# ('days_121_150', datetime.date(2019, 4, 14), datetime.date(2019, 5, 13))]

"""will need to adjust the first element since it does not follow the same rules as other ranges:"""
dranges[0] = ('days_0_30', dranges[0][1], max_date)

Now use list comprehension to generate these new columns:
df2 = df1.select('*', *[
        f.when((df1.trans_date >= d[1]) & (df1.trans_date <= d[2]), df1.balance).otherwise(None).alias(d[0])
            for d in dranges
       ])

df2.show()
+-------+------+----------+-------+---------+----------+----------+-----------+------------+
|user_id|amount|trans_date|balance|days_0_30|days_31_60|days_61_90|days_91_120|days_121_150|
+-------+------+----------+-------+---------+----------+----------+-----------+------------+
|    101| 99.10|2019-06-04|  85.03|     null|      null|      null|      85.03|        null|
|    101|-69.81|2019-09-11|  85.03|    85.03|      null|      null|       null|        null|
|    101| 12.51|2018-12-14|  85.03|     null|      null|      null|       null|        null|
|    101| 43.23|2018-09-11|  85.03|     null|      null|      null|       null|        null|
|    102| 89.27|2019-06-04| 178.37|     null|      null|      null|     178.37|        null|
|    102| 89.10|2019-03-04| 178.37|     null|      null|      null|       null|        null|
|    103| 73.11|2019-09-10|  73.11|    73.11|      null|      null|       null|        null|
+-------+------+----------+-------+---------+----------+----------+-----------+------------+

Note: it looks you are using 90-day intervals instead of 30-days in your code.  but I believe it's simple for your to adjust the above code to your needs.