Groupby and aggregate distinct values as a string

Viewed 975

I have a table as below:

ID   start date     name        type
 1   2020/01/01   cheese,meat    A, B
 1   2020/01/01   cheese,fruit   A, C

And the desired output should be:

ID    start date    count                 type 
1     2020/01/01   cheese,meat,fruit      A,B,C

I tried with collect_list and collect_set, but both did not work.

3 Answers

You can split and explode the columns, then group by and collect_set:

import pyspark.sql.functions as F

df2 = df.withColumn(
    'name',
    F.explode(F.split('name', ','))
).withColumn(
    'type',
    F.explode(F.split('type', ','))
).groupBy(
    'ID', 'start date'
).agg(
    F.concat_ws(',', F.collect_set('name')).alias('name'),
    F.concat_ws(',', F.collect_set('type')).alias('type')
)

df2.show()
+---+----------+-----------------+-----+
| ID|start date|             name| type|
+---+----------+-----------------+-----+
|  1|2020/01/01|fruit,meat,cheese|C,B,A|
+---+----------+-----------------+-----+

You can use array_distinct to remove duplicates after collect_set :

from pyspark.sql import functions as F

df1 = df.groupBy("ID", "start date").agg(
    F.concat_ws(",", F.collect_set("name")).alias("name"),
    F.concat_ws(",", F.collect_set("type")).alias("type"),
).select(
    "ID",
    "start date",
    F.array_join(F.array_distinct(F.split("name", ",")), ",").alias("name"),
    F.array_join(F.array_distinct(F.split("type", ",")), ",").alias("type")
)

df1.show()

# +---+----------+-----------------+-------+
# | ID|start date|             name|   type|
# +---+----------+-----------------+-------+
# |  1|2020/01/01|cheese,fruit,meat|A, C, B|
# +---+----------+-----------------+-------+

Another way using regexp_replace to remove the duplicates:

df1 = df.groupBy("ID", "start date").agg(
    F.concat_ws(",", F.collect_set("name")).alias("name"),
    F.concat_ws(",", F.collect_set("type")).alias("type"),
).select(
    "ID",
    "start date",
    F.regexp_replace("name", r"\b(\w+)\b\s*,\s*(?=.*\1)", "").alias("name"),
    F.regexp_replace("type", r"\b(\w+)\b\s*,\s*(?=.*\1)", "").alias("type")
)

You can use:

df.select(
    df.ID,
    df.start_date,
    F.split(df.name, ',').alias('name'),
    F.split(df.type, ',').alias('type')
).groupby('ID', 'start_date').agg(
    F.concat_ws(',', F.array_distinct(F.flatten(F.collect_list('name')))).alias('name'),
    F.concat_ws(',', F.array_distinct(F.flatten(F.collect_list('type')))).alias('type')
)

Result:

+---+----------+-----------------+-----+
| ID|start_date|             name| type|
+---+----------+-----------------+-----+
|  1|2020/01/01|cheese,meat,fruit|A,B,C|
+---+----------+-----------------+-----+
Related