Pyspark flatten rows with multiple values per ID into multiple columns

I am using PySpark to transform join several dataframes, but one of the dataframe columns contains multiple values per ID, resulting in multiple rows for those IDs. I need to flatten those rows into one one, with a unique column for each value. My data looks like this:

...and the final transformation should look like this:

| ID | Text_1 | Text_2 | Text_3 |
| -------- | -------- |
| 1 | some text | more text | still more text |
| 2 | some text | still more text | null |
| 3 | null | null | null |

How can I do this without using up too much memory?

I was able to group the data by ID to count how many columns it needs to be expanded to, but I'm unsure how to move forward.

The way you started works:

1. Create the input dataframe:

spark = ...
df = spark.createDataFrame([
    [1, 	'some text'],
    [1, 	'more text'],
    [1, 	'still more text'],
    [2, 	'some text'],
    [2, 	'still more text'],
    [3, 	None]
], ["ID", "Text"])

2. Group the dataframe by ID and count the number of strings for each ID:

from pyspark.sql import functions as F

df2 = df.groupBy('ID').agg(F.collect_list('Text').alias('Text')) \
    .withColumn('cnt', F.size('Text')).cache()

The cache is important here as the df2 dataframe will be used twice.

3. Get the maximum number of strings for a single ID

max = df2.agg(F.max('cnt')).first()[0]

4. Select the correct element of the collected array for each column

textcols = [F.col('Text')[i].alias(f'Text_{i+1}') for i in range(0, max)]
df2.select([F.col('ID')] + textcols).orderBy('ID').show()

Result:

+---+---------+---------------+---------------+
| ID|   Text_1|         Text_2|         Text_3|
+---+---------+---------------+---------------+
|  1|some text|      more text|still more text|
|  2|some text|still more text|           null|
|  3|     null|           null|           null|
+---+---------+---------------+---------------+