колонны в писпарке

Question

У меня есть такой набор данных

+--------------+---------+
|flavor        |type     |
+--------------+---------+
|chocolate     |ice cream|
|strawberry    |  cake   |
|chocolate     |  cake   |
|chocolate     |ice cream|
|strawberry    |ice cream|
+--------------+---------+

Мне нужно создать еще один столбец ("итого") с общим количеством

• разновидностей независимо от типа и
• аромат и тип в сочетании

Все это в одной таблице, которая будет иметь 3 столбца, как показано ниже:

+--------------+---------+--------+
|flavor        |type     |total   |
+--------------+---------+--------+
|chocolate     |ice cream|   2    |
|chocolate     |  cake   |   1    | 
|chocolate     |  null   |   3    |
|strawberry    |  cake   |   1    |
|strawberry    |ice cream|   1    |
|strawberry    | null    |   2    |
+--------------+---------+--------+

Я пробовал это код df.groupBy("flavor,"type").agg(F.count(df['flavor']).alias("total"))

Answer 1

Вы можете использовать rollup

df.rollup('flavor','type').count() \
    .where('flavor is not null or type is not null') \
    .selectExpr('flavor', 'type', 'count as total') \
    .show()
# Where cluase is to remove total count=5 (flavor and type both null)

+----------+--------+-----+
|    flavor|    type|total|
+----------+--------+-----+
| chocolate|    cake|    1|
| chocolate|    null|    3|
|strawberry|    cake|    1|
| chocolate|icecream|    2|
|strawberry|    null|    2|
|strawberry|icecream|    1|
+----------+--------+-----+