如何为多列在pyspark数据框中的一列中计算每个分类变量的频率？

Question

我想计算一列中每个类别的频率，并用频率计数替换该列中的值。 我想为pyspark数据帧在pyspark中的多个列上执行此操作。

例如，考虑以下数据框：

+-------+-------+-------+  
| col_1 | col_2 | col_3 |
+-------+-------+-------+  
|   a   |   f   |   g   |  
|   c   |   e   |   a   |  
|   a   |   d   |   g   |  
|   a   |   d   |   g   |  
|   b   |   f   |   b   |  
|   c   |   d   |   g   |  
|   b   |   d   |   c   |  
|   a   |   d   |   g   |  
|   b   |   f   |   g   |  
+-------+-------+-------+  

我想将此pyspark数据帧转换为以下内容：

+-------+-------+-------+  
| col_1 | col_2 | col_3 |
+-------+-------+-------+ 
|   4   |   3   |   6   |
|   2   |   1   |   1   |
|   4   |   5   |   6   |
|   4   |   5   |   6   |
|   3   |   2   |   1   |
|   2   |   5   |   6   |
|   3   |   5   |   1   |
|   4   |   5   |   6   |
|   3   |   2   |   6   |
+-------+-------+-------+  

我有以下代码：

spark = SparkSession.builder.getOrCreate()

df = spark.read.parquet(data)
df.show()

+-------+-------+-------+
| col_1 | col_2 | col_3 |
+-------+-------+-------+
|   a   |   f   |   g   |
|   c   |   e   |   a   |
|   a   |   d   |   g   |
|   a   |   d   |   g   |
|   b   |   f   |   b   |
|   c   |   d   |   g   |
|   b   |   d   |   c   |
|   a   |   d   |   g   |
|   b   |   f   |   g   |
+-------+-------+-------+

我可以使用以下代码使用for循环计算每一列的频率：

df.groupby('col_1').count().toDF('category', 'count').show()

我知道我可以对每个列进行此操作并将结果粘合在一起。 我想知道是否有更好的方法可以做到这一点。

Answer 1

您可以使用窗口函数来实现：

import pyspark.sql.functions as F
from pyspark.sql import Window

l = [
(   'a'   ,   'f'   ,   'g'   ),
(   'c'   ,   'e'   ,   'a'   ),
(   'a'   ,   'd'   ,   'g'   ),
(   'a'   ,   'd'   ,   'g'   ),
(   'b'   ,   'f'   ,   'b'   ),
(   'c'   ,   'd'   ,   'g'   ),
(   'b'   ,   'd'   ,   'c'   ),
(   'a'   ,   'd'   ,   'g'   ),
(   'b'   ,   'f'   ,   'g'  )
]

columns = ['col_1', 'col_2','col_3']

df=spark.createDataFrame(l, columns)


for column in columns:
    df = df.withColumn(column, F.count(column).over(Window.partitionBy(column)))

df.show()

输出：

+-----+-----+-----+ 
|col_1|col_2|col_3| 
+-----+-----+-----+ 
|    4|    3|    6| 
|    3|    3|    6| 
|    4|    5|    6| 
|    4|    5|    6| 
|    4|    5|    6| 
|    2|    5|    6| 
|    3|    5|    1| 
|    3|    3|    1| 
|    2|    1|    1|
+-----+-----+-----+