Aggregate ArrayType column to get another ArrayType column without UDF

I have a dataframe like this, that has one ArrayType column

simpleData = [("202305","Sales","NY",[1,2,3]),
    ("202306","Sales","NY",[4,2,3]),
    ("202305","Sales","CA",[4,5,3]),
    ("202306","Finance","CA",[4,5,6]),
    ("202305","Finance","NY",[5,6,7]),
    ("202306","Finance","NY",[6,7,8]),
  ]

schema = ["month","department","state","lis"]
df = spark.createDataFrame(data=simpleData, schema = schema)

What I would like to do is find the median value of the lis column, but grouped by some columns. Pretty much like what we would do with percentile_approx() over a numeric column, but here I want that to be done such that the median is computed on the unioned list for each group.

Desired output:

df1 = df.groupBy('department').agg(desired_function(col('lis')).alias('lis_median'))

gives

+----------+----------+
|department|lis_median|
+----------+----------+
|     Sales|         3|
|   Finance|         6|
+----------+----------+

I tried doing

df1 = df.groupBy('department').agg(collect_list(col('lis')).alias('lis_median'))

to atleast get the aggregated list over the group, but it only generates a list of list.

I'd like some way to implement this. I dont want to use UDF due to execution speed concerns

Try with higher order functions for array i.e. flatten,array_sort for this case.

• if the size of array is even then get the median value from midvalue + midvalue -1/2
• if size of array is odd then get the mid value as the median

Example:

simpleData = [("202305","Sales","NY",[1,2,3]),
    ("202306","Sales","NY",[4,2,3]),
    ("202305","Sales","CA",[4,5,3]),
    ("202306","Finance","CA",[4,5,6]),
    ("202305","Finance","NY",[5,6,7]),
    ("202306","Finance","NY",[6,7,8]),
  ]

schema = ["month","department","state","lis"]
df = spark.createDataFrame(data=simpleData, schema = schema)

df.groupBy("department").agg(flatten(collect_set(col("lis"))).alias("cs")).\
  withColumn('col', array_sort(expr('filter(cs, x -> x is not null)'))).\
  withColumn('lis_median', when(size(col('cs')) % 2 == 0,
                                            (expr('col[int(size(cs)/2)]') + expr('col[int(size(cs)/2)-1]'))/2
                                            ).otherwise(expr('col[int(size(cs)/2)]'))).\
                                              drop(*['col','cs']).\
  show(10,False)
#+----------+----------+
#|department|lis_median|
#+----------+----------+
#|Sales     |3.0       |
#|Finance   |6.0       |
#+----------+----------+

I agree with @notNull, flatten the list and get median according to the size.
Below is one of the approaches you can also follow.
You can use posexplode function which gives position and value of that list,
and do aggregate and find median directly as below.

display(df.selectExpr("*", "posexplode(lis) as (pos, lis_value)").groupBy('department').agg(median("lis_value").alias('lis_median')))