Here is a way to compute all the medians in parallel (approximately of course).
Let's start by collecting all the possible keys:
// generating data (BTW, please provide that code next time you ask a question)
val df = Seq(("aaaaaaaaaaaa", 11), ("aaaaaaaaaaaa", 44),
("bbbbbbbbbbbb", 22), ("aaaaaaaaaaaa", 33))
.toDF("A", "B")
val cols = df.select("A").distinct.collect.map(_.getAs[String](0))
Let's say cols has size N.
One way to go at it would be to iterate over cols and compute the median via N separate jobs. Another answer provide the code for that.
Yet, it is possible to compute all the medians in parallel (and therefore just one job) using a pivot.
val precision = 1e-3
val medians = df
// the index is artificial, it is just meant to treat each line individually
.withColumn("index", monotonicallyIncreasingId)
.groupBy("index")
.pivot("A").agg(first('B))
.stat.approxQuantile(cols, Array(0.5), precision)
val result = cols.indices.map(i => cols(i) -> medians(i)(0)).toMap
It may not be worth it with only a few keys be could be interesting if you have more.
EDIT
The first solution keeps one line per line in the original dataframe and will work, I think, all the time. In case you have many keys, it could be interesting to compute the index with a window like this.
.withColumn("index", row_number() over Window.partitionBy("A").orderBy("B"))
But do not use windows if you have millions of rows per key. It could be very slow or even crash your job.