Range Array Column

I want to add an array column with the range/sequence of number of days between the start date and end date. An example of the data would be

from datetime import datetime
from pyspark.sql.types import *
from pyspark.sql import functions as F

data = [
    (
        datetime(2020,1,1),
        datetime(2020,12,31),
    ),
    
]
rdd = sc.parallelize(data)

schema = StructType(
    [
       StructField("dateStart", DateType(), True),
       StructField("dateEnd", DateType(), True)
    ]
)

df = spark.createDataFrame(rdd, schema)

df = df.withColumn(
    "dateDiff",
    F.datediff(
        F.col("dateEnd"),
        F.col("dateStart"),
    )
)

df.show()
+----------+----------+--------+
| dateStart|   dateEnd|dateDiff|
+----------+----------+--------+
|2020-01-01|2020-12-31|     365|
+----------+----------+--------+

and what I need is the following so that I can explode the range column.

+----------+----------+--------+------------+
| dateStart|   dateEnd|dateDiff|       range|
+----------+----------+--------+------------+
|2020-01-01|2020-12-31|     365| [1,2,..365]|
+----------+----------+--------+------------+

I wrote a udf to do this, but I want to know if there is a more efficient way to do this?

def range_func(x):
    return list(range(1, x+1))

range_udf = F.udf(range_func, ArrayType(IntegerType()))

df = df.withColumn("Range", range_udf(F.col("dateDiff")))
df.show()
+----------+----------+--------+--------------------+
| dateStart|   dateEnd|dateDiff|               Range|
+----------+----------+--------+--------------------+
|2020-01-01|2020-12-31|     365|[1, 2, 3, 4, 5, 6...|
+----------+----------+--------+--------------------+