columns value in order - pyspark

Question

I have this DataFrame bellow:
Ref ° | indice_1 | Indice_2 | rank_1 | rank_2 | echelon_from | section_from | echelon_to | section_to
--------------------------------------------------------------------------------------------------------------------------------------------
70574931 | 19 | 37.1 | 32 | 62 | ["10032,20032"] | ["11/12","13"] | ["40062"] | ["14A"]
---------------------------------------------------------------------------------------------------------------------------------------------
70574931 | 18 | 36 | 32 | 62 | ["20032"] | ["13"] | ["30062,40062"] | ["14,14A"]
I want concatenate the lines that have the same Ref° number, to concatenate echelon_from values, section_from values, echelon_to values and section_to values with duplicates there values, like in example bellow, and without touch the rest of the columns.
Ref ° | Indice_1 | Indice_2 | rank_1 | rank_2 | echelon_from | section_from | echelon_to | section_to
---------------------------------------------------------------------------------------------------------------------------------------------
70574931 | 19 | 37.1 | 32 | 62 | ["10032,20032"] | ["11/12","13"] | ["30062,40062"] | ["14,14A"]
----------------------------------------------------------------------------------------------------------------------------------------------
70574931 | 18 | 36 | 32 | 62 | ["10032,20032"] | ["11/12","13"] | ["30062,40062"] | ["14,14A"]
Some columns values in my original Dataframe are duplicates, I shouldn't touch them, I should keep there values to keep the same line numer of my DataFrame.
Someone can help me please how can I do it ?
Thank you!

Answer 1

There is multiple ways of doing this. One way is to explode all the given lists and collect them back again as a set.
from pyspark.sql import functions as F
lists_to_concat = ['echelon_from', 'section_from', 'echelon_to', 'section_to']
columns_not_to_concat = [c for c in df.columns if c not in lists_to_concat]
for c in lists_to_concat:
df = df.withColumn(c, F.explode(c))
df = (
df
.groupBy(*columns_not_to_concat)
.agg(
*[F.collect_set(c).alias(c) for c in lists_to_concat]
)
)
Another more elegant way is to use flatten().
from pyspark.sql import functions as F
lists_to_concat = ['echelon_from', 'section_from', 'echelon_to', 'section_to']
for c in lists_to_concat:
df = df.withColumn(c, F.flatten(c))
References:
https://spark.apache.org/docs/latest/api/python/pyspark.sql.html#pyspark.sql.functions.flatten