Суммирование столбца подсчета в PySpark

Question

Я пытаюсь получить пропорции в pyspark df. Я агрегировал и посчитал так (где var1 и var2 - строки):

import pyspark.sql.functions as f

df = df.select('var1', 'var2') \
       .groupBy(df.var1) \
       .agg({'var1':'count'}) \
       .withColumnRenamed('count(1)', 'Total') \
       .withColumnRenamed('sum(var1)', 'Sum')

, который выдает:

+--------------------+---------+
|                var1|    Total|
+--------------------+---------+
|                  a1| 12016668|
|                  a2| 22653585|
|                  a3|107313117|
|                  a4|       69|
|                  a5|        5|
|                  a6| 13092243|
|                  a7|      372|
|                  a8|     3167|
|                  a9| 18712794|
|                 a10|  2456488|
|                 a11|  2733665|
|                 a12| 14854475|
+-----------------------+------+

Рядом с Total Я хочу получить столбец с пропорциями, так что-то вроде Total / sum(Total). Я попытался:

exprs = {'var1':'count', 'var1':'sum'}
df = df.select('var1', 'var2') \
       .groupBy(df.var1) \
       .agg(exprs) \
       .withColumnRenamed('count(1)', 'Total') \

Но это выводит null значения, поскольку он пытается суммировать string:

+-----------------------+----+---------+
|                   var1| Sum|    Total|
+-----------------------+----+---------+
|                     a1|null| 12016668|
|                     a2|null| 22653585|
|                     a3|null|107313117|
|                     a4|null|       69|
|                     a5|null|        5|
|                     a6|null| 13092243|
|                     a7|null|      372|
|                     a8|null|     3167|
|                     a9|null| 18712794|
|                    a10|null|  2456488|
|                    a11|null|  2733665|
|                    a12|null| 14854475|
+-----------------------+----+---------+

Я также попытался определить переменную и поделить на нее:

N = df.groupBy().agg(f.sum(f.col('Total')))


df = df.withColumn('Proportion',  f.format_number(df.Total / N, 4))

В этом случае я получаю ошибку, и я также не думаю, что это хорошая идея делить на такие числа с плавающей точкой вместо того, чтобы иметь это как столбец (с повторяющимися значениями или чем-то подобным):

AttributeError: 'DataFrame' object has no attribute '_get_object_id'

Какой хороший способ сделать это?

Answer 1

У тебя почти было это. Вы просто должны были collect[0][0] сумма.

N = df.agg(f.sum(f.col('Total'))).collect()[0][0]


df.withColumn('Proportion',  f.format_number(df.Total / N, 4)).show()

+----+---------+----------+
|var1|    Total|Proportion|
+----+---------+----------+
|  a1| 12016668|    0.0620|
|  a2| 22653585|    0.1169|
|  a3|107313117|    0.5536|
|  a4|       69|    0.0000|
|  a5|        5|    0.0000|
|  a6| 13092243|    0.0675|
|  a7|      372|    0.0000|
|  a8|     3167|    0.0000|
|  a9| 18712794|    0.0965|
| a10|  2456488|    0.0127|
| a11|  2733665|    0.0141|
| a12| 14854475|    0.0766|
+----+---------+----------+