Pyspark: разделить один ряд на другой в groupBy

Question

У меня есть фрейм данных pyspark, и я хотел бы разделить одну строку на другую в группах. Внутри групп будет две строки: одна со значением счетчика где removal == 1, а другая со значением счетчика где removal == 0.

Как мне разделить один счет на другой, чтобы получить соотношение для каждой группы в новом столбце? groupBy включен limit и 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)

]

Что я хочу: (или в похожей раскладке)

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)
]

Я знаю способы сделать это, разделив и снова объединив данные, но я бы предпочел более привлекательную функцию агг.

Answer 1

Поскольку для каждого значения removal имеется только одна строка, простой способ - использовать where для фильтрации каждого отдельного значения и 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|
#+-----+-------+------------------+-------+-------+