Pyspark: devide one row by another in groupBy - group-by

Question

I have a pyspark data frame and I'd like to divide one row by another within groups. Within groups there will be two rows: one with a count value where removal == 1 and the other with a count value where removal == 0.
How do I divide one count by the other to get the ratio for each group in a new column? The groupBy is on limit and test_id
columns = ['removal', 'limit', 'test_id', 'count']
vals = [
(1, 'UL', 'AB', 141),
(0, 'UL', 'AB', 140),
(1, 'LL', 'AB', 21),
(0, 'LL', 'AB',12),
(0, 'UL', 'EF', 200),
(1, 'UL', 'EF',12)
]
What I want: (or in a similar layout)
columns = ['limit', 'test_id', 'ratio', count_1, count_0]
vals = [
('UL', 'AB', 1.007, 141, 140)
('LL', 'AB', 1.75, 21, 12),
('UL', 'EF', 0.06, 12, 200)
]
I know ways to do it by splitting and then merging the data again, but I'd rather have a nicer agg function.

Answer 1

Since there is only one row per value of removal, the straightforward way is to use where to filter for each distinct value and join:
from pyspark.sql.functions import col
df.where("removal = 1").alias("a")\
.join(df.where("removal = 0").alias("b"), on=["limit", "test_id"])\
.select(
"limit",
"test_id",
(col("a.count") / col("b.count")).alias("ratio"),
col("a.count").alias("count_1"),
col("b.count").alias("count_0")
).show()
#+-----+-------+------------------+-------+-------+
#|limit|test_id| ratio|count_1|count_0|
#+-----+-------+------------------+-------+-------+
#| UL| AB|1.0071428571428571| 141| 140|
#| LL| AB| 1.75| 21| 12|
#| UL| EF| 0.06| 12| 200|
#+-----+-------+------------------+-------+-------+