Pyspark-取一列的总和，然后使用该值除以另一列

Question

我有一个数据框df

>>> df = spark.createDataFrame([[1,0], [2,1], [3,1], [4,0], [5,1]], ['a', 'b'])
>>> df.show()
+---+---+
|  a|  b|
+---+---+
|  1|  0|
|  2|  1|
|  3|  1|
|  4|  0|
|  5|  1|
+---+---+

和

>>> nrows = df.count()

我使用df创建了一个新数据框a，它是df的集合。

>>> a = df.groupby('b').count()
>>> a.show()
+---+-----+
|  b|count|
+---+-----+
|  0|    2|
|  1|    3|
+---+-----+

我需要在a中创建一个名为ev的新列。第ev行的i的值由

给出

这是我期望的输出

+---+-----+------------------+
|  b|count|           ev_norm|
+---+-----+------------------+
|  0|    2|              1.25|
|  1|    3|0.8333333333333334|
+---+-----+------------------+

但是我首先到达这里，为分子num创建了一个新列(nrows-count)，在两行中产生了3,2。然后，我继续计算所有行都相同的分母（denom=0.48）。最后，我使用.withColumn创建了一个新列，该列将num列除以denom。

有没有一种方法可以即时执行denom的计算而不必预先计算并可以一步完成上述操作？

我尝试过类似的

a = a.withColumn('ev_norm', (nrows - F.col('count'))/F.sum(F.col('count')*(nrows - F.col('count'))))`

和

a = a.withColumn('ev_norm', (sum([F.col('count')*(nrows-F.col('count'))]))

但是它给我一个错误或错误的答案。