Pyspark - Preserve order of collect list and collect set over multiple columns

user_8832641 0
Viewed 2340

I have the below pyspark dataframe.

Column_1 Column_2 Column_3 Column_4
1        A        U1       12345
1        A        A1       549BZ4G

Expected output:

Group by on column 1 and column 2. Collect set column 3 and 4 while preserving the order in input dataframe. It should be in the same order as input. There is no dependency in ordering between column 3 and 4. Both has to retain input dataframe ordering

Column_1 Column_2 Column_3 Column_4
1        A        U1,A1    12345,549BZ4G

What I tried so far:

I first tried using window method. Where I partitioned by column 1 and 2 and order by column 1 and 2. I then grouped by column 1 and 2 and did a collect set on column 3 and 4.

I didn't get the expected output. My result was as below.

Column_1 Column_2 Column_3 Column_4
1        A        U1,A1    549BZ4G,12345

I also tried using monotonically increasing id to create an index and then order by the index and then did a group by and collect set to get the output. But still no luck.

Is it due to alphanumeric and numeric values ? How to retain the order of column 3 and column 4 as it is there in input with no change of ordering.

pyspark apache-spark dataframe group-by python
4 Answers

Use monotically_increasing_id function from spark to maintain the order.you can find more info about it here

    #InputDF
    # +----+----+----+-------+
    # |col1|col2|col3|   col4|
    # +----+----+----+-------+
    # |   1|   A|  U1|  12345|
    # |   1|   A|  A1|549BZ4G|
    # +----+----+----+-------+

    df1 = df.withColumn("id", F.monotonically_increasing_id()).groupby("Col1", "col2").agg(F.collect_list("col4").alias("Col4"),F.collect_list("col3").alias("Col3"))

    df1.select("col1", "col2",F.array_join("col3", ",").alias("col3"),F.array_join("col4", ",").alias("col4")).show()

    # OutputDF
    # +----+----+-----+-------------+
    # |col1|col2| col3|         col4|
    # +----+----+-----+-------------+
    # |   1|   A|U1,A1|12345,549BZ4G|
    # +----+----+-----+-------------+

Use array_distinct on top of collect_list to have distinct values and maintain order.

    #InputDF
    # +----+----+----+-------+
    # |col1|col2|col3|   col4|
    # +----+----+----+-------+
    # |   1|   A|  U1|  12345|
    # |   1|   A|  A1|549BZ4G|
    # |   1|   A|  U1|123456 |
    # +----+----+----+-------+



    df1 = df.withColumn("id", F.monotonically_increasing_id()).groupby("Col1", "col2").agg(
        F.array_distinct(F.collect_list("col4")).alias("Col4"),F.array_distinct(F.collect_list("col3")).alias("Col3"))

    df1.select("col1", "col2", F.array_join("col3", ",").alias("col3"), F.array_join("col4", ",").alias("col4")).show(truncate=False)

    # +----+----+-----+---------------------+
    # |col1|col2|col3 |col4                 |
    # +----+----+-----+---------------------+
    # |1   |A   |U1,A1|12345,549BZ4G,123456 |
    # +----+----+-----+---------------------+
user_86658430

How about using struct ?

val result = df.groupBy(
  "Column_1", "Column_2"
).agg(
  collect_list(
    struct(col("Column_3"), col("Column_4"))
  ).alias("coll")
).select(
  col("Column_1"), col("Column_2"), col("coll.Column_3"), col("coll.Column_4")
)

This should produce the expected result. The trick is that struct preserves named elements so you can reference them via . accessor. It works in ArrayType(StructType) as well.

Grouping collocated concepts via struct feels natural since it is a structure relation that you are trying to preserve here.

In theory you might even not want to unpack structs since these values appear to have a dependance.

user_138781830

All the collect functions(collect_set, collect_list) within spark are non-deterministic since the order of collected result depends on the order of rows in the underlying dataframe which is again non-deterministic.

So to conclude you can't really preserve the order with spark collect functions

ref - functions.scala spark repo

user_47588230

I am not sure why it is not showing you correct result, for me this is as per the expected result -

Input

df_a = spark.createDataFrame([(1,'A','U1','12345'),(1,'A','A1','549BZ4G')],[ "col_1","col_2","col_3","col_4"])

Logic

from pyspark.sql import functions as F
df_a = df_a.groupBy('col_1','col_2').agg(F.collect_list('col_3').alias('col_3'), F.collect_list('col_4').alias('col_4'))

Output

df_a.show()
+-----+-----+--------+----------------+
|col_1|col_2|   col_3|           col_4|
+-----+-----+--------+----------------+
|    1|    A|[U1, A1]|[12345, 549BZ4G]|
+-----+-----+--------+----------------+
user_51082200
Related