PySpark TopK 问题(分组TopK)（4）

PySpark TopK 问题(分组TopK)

记录几种利用PySpark计算TopK的方法,准备使用两个例子，其中第一个例子是计算不同院系，不同班,不同学科的成绩前K名的分数。第二个例子以文本数据为例，计算在不同文本类别下出现TopK 频率的单词。

1.准备数据

1,111,68,69,90,1班,经济系
2,112,73,80,96,1班,经济系
3,113,90,74,75,1班,经济系
4,114,89,94,93,1班,经济系
5,115,99,93,89,1班,经济系
6,121,96,74,79,2班,经济系
7,122,89,86,85,2班,经济系
8,123,70,78,61,2班,经济系
9,124,76,70,76,2班,经济系
10,211,89,93,60,1班,外语系
11,212,76,83,75,1班,外语系
12,213,71,94,90,1班,外语系
13,214,94,94,66,1班,外语系
14,215,84,82,73,1班,外语系
15,216,85,74,93,1班,外语系
16,221,77,99,61,2班,外语系
17,222,80,78,96,2班,外语系
18,223,79,74,96,2班,外语系
19,224,75,80,78,2班,外语系
20,225,82,85,63,2班,外语系
--------------------- 
作者：wangpei1949 
来源：CSDN 
原文：https://blog.csdn.net/wangpei1949/article/details/66474029 
版权声明：本文为博主原创文章，转载请附上博文链接！

数据来源在上面。每一行一条记录，每个属性使用,分割，分别是id，学号，语文，数学，外语，班级，院系
下面就开始各种计算不同院系，不同学科的topK的方法，数据被保存名为topK的txt文件。先撸出来结果更为好理解这篇博客需要做的TopK：

[('经济系1班', {'语文top3': [99, 90, 89], '数学top3': [94, 93, 80], '英语top3': [96, 93, 90]}), 
 ('经济系2班', {'语文top3': [96, 89, 76], '数学top3': [86, 78, 74], '英语top3': [85, 79, 76]}),
 ('外语系1班', {'语文top3': [94, 89, 85], '数学top3': [94, 94, 93], '英语top3': [93, 90, 75]}), 
 ('外语系2班', {'语文top3': [82, 80, 79], '数学top3': [99, 85, 80], '英语top3': [96, 96, 78]})]

2.计算成绩TopK例子

2.1. 直观TopK

啥都不考虑，直接使用groupByKey，在不同的group下使用sorted函数后取TopK
代码如下：

# -*- coding: utf-8 -*-

"""
 分组topK
 @Time    : 2019/2/20 15:06
 @Author  : MaCan (ma_cancan@163.com)
 @File    : groupTopK.py
"""

from pyspark import SparkContext, SparkConf, Row
from pyspark.sql import SparkSession, SQLContext, Window
from pyspark.sql.types import *
import pyspark.sql.functions as fn

import os
import re
import time

def top_k_1(line):
    """
    排序分组
    :param line: line[0]==>key, line[1]==>key对应原始的每一行value
    :return:
    """
    key = line[0]

    sort_rst_1 = sorted([entity[1] for entity in list(line[1])], reverse=True)[:3]
    sort_rst_2 = sorted([entity[2] for entity in list(line[1])], reverse=True)[:3]
    sort_rst_3 = sorted([entity[3] for entity in list(line[1])], reverse=True)[:3]
    return (key, {'语文top3':sort_rst_1, '数学top3':sort_rst_2, '英语top3':sort_rst_3})


if __name__ == '__main__':
    start_t = time.perf_counter()
    conf = SparkConf().setAppName('text_trans').setMaster("local[*]")
    sc = SparkContext()
    sc.setLogLevel(logLevel='ERROR')
    spark = SparkSession.builder.config(conf=conf).getOrCreate()
    try:
        data = spark.read.text('topK').rdd.map(lambda x: x[0].split(','))\
            .map(lambda x: (x[1], int(x[2]), int(x[3]), int(x[4]), x[5], x[6]))\
            .groupBy(lambda x: x[5] + x[4])\
            .map(top_k_1)
        print(data.collect())

    except Exception as e:
        print(e)
    finally:
        spark.stop()
        print(time.perf_counter() - start_t)

程序员最怕写注释，程序员最怕别人不写注释~~~,恩恩，不慌，我是个好程序员，下面将主要说明一下代码：
x1-x2 表示代码的x1行带x2行，例如下面

35-38: spark环境参数的初始化
40： 读取文本，并且根据,进行切分每个属性
41： 数据格式转化，方便我们以int格式进行分数的sort
42： 自定义key的group,我们也可以在41行中定义好pair，然后就调用groupByKey就好啦。
43： 调用top_k_1 函数，进行每一个组内的排序，具体代码在19-30行
25： line[0]就是key， line[1]是value，其中key是42行中的自定义key,value是每一行数据的tuple
27-29：对三个学科进行排序
30： 组合数据返回
我想我应该写明白了，后面的demo就不写的那么详细了，只有改动地方会提出。

2.2. 考虑分组平衡的TopK

毕竟我们在使用“大数据”处理框架，那么假设我们的数据很大，大到运行上面的代码会出现OOM，为啥会OOM？因为上面的sorted 排序是内存排序，在python 3.6中默认使用的是[TimeSort 排序算法] (https://stackoverflow.com/questions/10948920/what-algorithm-does-pythons-sorted-use) 当我们的数据很大的时候，总会把内存干炸掉的。这里我们先不考虑更换排序算法，只考虑如何通过groupBy partition这种原生的方式来进行优化。
很明显，我们可以多分组，那么组内的数据就会小了，这样在sorted的时候可以一定成都的规避，同时，通过重新分组，我们可以避免一些数据在组内不平衡的问题，例如某一个学院的人很多，这样在partition计算的时候造成这个key下需要计算很长时间，影响整体性能。
下面上代码如何重新排序规避大量数据的内存排序和数据不平衡问题。

# -*- coding: utf-8 -*-

"""
 分组topK
 @Time    : 2019/2/20 15:06
 @Author  : MaCan (ma_cancan@163.com)
 @File    : groupTopK.py
"""

from pyspark import SparkContext, SparkConf, Row
from pyspark.sql import SparkSession, SQLContext, Window
from pyspark.sql.types import *
import pyspark.sql.functions as fn

import os
import re
import time
import random

def top_k_2(line):
    """
    排序分组
    :param line: line[0]==>key, line[1]==>key对应原始的每一行value
    :return:
    """
    key = line[0]

    index_to_label = ['语文Top3', '数学top3', '英语top3']
    value = list(line[1])
    index = value[0][0][1]
    sort_rst = []
    [sort_rst.extend(entity[1]) for entity in value]
    sort_rst = [x[index + 1] for x in sort_rst]
    sort_rst = sorted(sort_rst, reverse=True)[:3]
    return (key[0], index_to_label[index] + ":" + str(sort_rst))


def top_k_3(line):
    """
    对重新分组的一组数据求topK
    :param line:
    :return:
    """
    key = line[0][1]
    rst = []
    top_k = min(len(line[1]), broadcast_top_k.value)
    for i in range(broadcast_top_k.value):
        sort_rst = sorted(list(line[1]), key=lambda x: x[i+1], reverse=True)[:top_k]
        rst.append(((key, i), sort_rst))

    return rst


if __name__ == '__main__':
    start_t = time.perf_counter()
    conf = SparkConf().setAppName('text_trans').setMaster("local[*]")
    sc = SparkContext()
    sc.setLogLevel(logLevel='ERROR')
    spark = SparkSession.builder.config(conf=conf).getOrCreate()
    try:
        broadcast_random_mode = sc.broadcast(2)
        broadcast_top_k = sc.broadcast(3)
        data = spark.read.text('topK').rdd.map(lambda x: x[0].split(','))\
            .map(lambda x: (x[6]+x[5], [x[1], int(x[2]), int(x[3]), int(x[4])]))\
            .map(lambda x: ((random.randint(0, broadcast_random_mode.value), x[0]), x[1]))\
            .groupByKey()\
            .flatMap(top_k_3)\
            .groupBy(lambda x: x[0])\
            .map(top_k_2)
        print(data.collect())
    except Exception as e:
        print(e)
    finally:
        spark.stop()
        print(time.perf_counter() - start_t)

输出结果：

[('经济系1班', '语文Top3:[99, 90, 89]'),
('经济系1班', '数学top3:[94, 93, 80]'),
('经济系1班', '英语top3:[96, 93, 90]'),
('经济系2班', '语文Top3:[96, 89, 76]'), 
('经济系2班', '数学top3:[86, 78, 74]'), 
('经济系2班', '英语top3:[85, 79, 76]'), 
('外语系1班', '语文Top3:[94, 89, 85]'),
('外语系1班', '数学top3:[94, 94, 93]'), 
('外语系1班', '英语top3:[93, 90, 75]'), 
('外语系2班', '语文Top3:[82, 80, 79]'),
('外语系2班', '数学top3:[99, 85, 80]'), 
('外语系2班', '英语top3:[96, 96, 78]')]

懒得写成上面的输出格式，大佬们见谅一下。　在第一次计算topK的时候，进行第一次的数据过滤，因为对于在一个partition中都不是topK 的数据，在整体中必然也不是的。此时数据减小了不少，在第二次进行排序计算的时候，OOM的概率就会小很多了，整个过程导致OOM的概率也小很多。

2.3. 使用窗口函数

窗口函数是啥？ 可以看下这里和这里
直接撸代码：

# -*- coding: utf-8 -*-

"""
 分组topK
 @Time    : 2019/2/20 15:06
 @Author  : MaCan (ma_cancan@163.com)
 @File    : groupTopK.py
"""

from pyspark import SparkContext, SparkConf, Row
from pyspark.sql import SparkSession, SQLContext, Window
from pyspark.sql.types import *
import pyspark.sql.functions as fn

import re
import time


if __name__ == '__main__':
    start_t = time.perf_counter()
    conf = SparkConf().setAppName('text_trans').setMaster("local[*]")
    sc = SparkContext()
    sc.setLogLevel(logLevel='ERROR')
    spark = SparkSession.builder.config(conf=conf).getOrCreate()
    try:
        data = spark.read.text('topK').rdd \
            .map(lambda x: x[0].split(',')) \
            .map(lambda x: Row(labels=x[6]+x[5], id=x[1], score_1=int(x[2]), score_2=int(x[3]), score_3=int(x[4])))
        data = spark.createDataFrame(data)
        data.select("labels", 'score_1',
                    fn.row_number().over(Window.partitionBy("labels").
                                         orderBy(data["score_1"].desc())).alias('index')).where('index <= 3').show()
    except Exception as e:
        print(e)
    finally:
        spark.stop()
        print(time.perf_counter() - start_t)

输出结果：

+------+-------+-----+
|labels|score_1|index|
+------+-------+-----+
| 外语系1班|     94|    1|
| 外语系1班|     89|    2|
| 外语系1班|     85|    3|
| 经济系2班|     96|    1|
| 经济系2班|     89|    2|
| 经济系2班|     76|    3|
| 外语系2班|     82|    1|
| 外语系2班|     80|    2|
| 外语系2班|     79|    3|
| 经济系1班|     99|    1|
| 经济系1班|     90|    2|
| 经济系1班|     89|    3|
+------+-------+-----+

上面除了使用DataFrame的select 还可以使用withColumn:

        data.withColumn("index",
                        fn.row_number().over(Window.partitionBy('labels').orderBy('score_1'))).where("index <= 3").show()

至此写了三种不同的分组排序的方法，下面就使用一个文本分析的实例来体验一把。

3. 找出不同类别文章下频率Top10的单词

数据来源于清华大学的分类数据集，下载链接:http://thuctc.thunlp.org/
我们使用解压后的13个类别的数据，使用jieba进行分词后，统计每个类别下的出现频率最高的前10个单词，使用上面方法二来写，其余的大同小异就不给例子了，运行环境为单机16G内存的PC

# -*- coding: utf-8 -*-

"""
 分组topK
 @Time    : 2019/2/20 15:06
 @Author  : MaCan (ma_cancan@163.com)
 @File    : groupTopK.py
"""

from pyspark import SparkContext, SparkConf, Row
from pyspark.sql import SparkSession, SQLContext, Window
from pyspark.sql.types import *
import pyspark.sql.functions as fn

import os
import re
import time
import random
import jieba


def seg_with_dir_name(dir_name, data):
    """
    分词后返回分词的dataframe
    :param spark:
    :param data:
    :return:
    """
    return [((dir_name, w), 1) for w in jieba.cut(data.strip(), cut_all=False) if len(w) > 1]


def top_k_first(line):
    """
    第一次topK 排序
    :param line:
    :return:
    """
    key = line[0][1]
    top_k = min(len(line[1]), broadcast_top_k.value)
    sort_rst = sorted(list(line[1]), key=lambda x: x[1], reverse=True)[:top_k]
    print((key, sort_rst))
    return (key, sort_rst)


def top_k_second(line):
    """
    第二次合并和topK 排序
    :param line:
    :return:
    """
    key = line[0]
    value = line[1]
    merge_value = []
    [merge_value.extend(list(x)) for x in list(value)]
    # merge_value = [tuple(x) for x in merge_value]
    merge_value = sorted(merge_value, key=lambda x: x[1], reverse=True)[:broadcast_top_k.value]
    print(merge_value)
    return (key, merge_value)


if __name__ == '__main__':
    start_t = time.perf_counter()

    # load file paths
    path = r'F:\data\THUCNews'
    paths = []
    for file in os.listdir(path):
        curr_dir = os.path.join(path, file)
        if os.path.isdir(curr_dir):
            paths.append(file)
    print(paths)
    
    conf = SparkConf().setAppName('text_trans').setMaster("local[*]")
    sc = SparkContext()
    sc.setLogLevel(logLevel='ERROR')
    spark = SparkSession.builder.config(conf=conf).getOrCreate()
    try:
        broadcast_top_k = sc.broadcast(10)
        broadcast_dir_name = sc.broadcast(paths)
        broadcast_random_mode = sc.broadcast(50)
        rdd_data = spark.read.text(os.path.join(path, broadcast_dir_name.value[0])).rdd.flatMap(lambda x: seg_with_dir_name(broadcast_dir_name.value[0], x[0]))
        rdd_data = rdd_data.union(spark.read.text(os.path.join(path, broadcast_dir_name.value[1])).rdd.flatMap(
            lambda x: seg_with_dir_name(broadcast_dir_name.value[1], x[0])))\
            .repartition(12)\
            .combineByKey(lambda v: 1,
                          lambda x, v: x+v,
                          lambda x,y: x+y)\
            .map(lambda x: ((random.randint(0, broadcast_random_mode.value), x[0][0]), (x[0][1], x[1])))\
            .groupByKey()\
            .map(top_k_first)\
            .groupByKey()\
            .map(top_k_second)

        print(rdd_data.collect())
    except Exception as e:
        print(e)
    finally:
        spark.stop()
        print(time.perf_counter() - start_t)

输出结果：

[('星座', [('自己', 9596), ('他们', 6709), ('爱情', 5491), ('一个', 4740), ('星座', 4481), ('测试', 4034), ('可以', 3488), ('我们', 3473), ('没有', 3382), ('心理', 3148), ('喜欢', 3027), ('因为', 2977), ('朋友', 2913), ('就是', 2883), ('如果', 2765), ('男人', 2664), ('什么', 2650), ('对方', 2554), ('时候', 2515), ('工作', 2391)]), ('彩票', [('主场', 24652), ('VS', 18599), ('客场', 16229), ('号码', 15810), ('比赛', 15671), ('球队', 14381), ('彩票', 14368), ('主队', 14364), ('开出', 13926), ('10', 13868), ('推荐', 13108), ('两队', 12599), ('赔率', 12222), ('联赛', 11984), ('双色球', 11426), ('本场', 10332), ('数据', 9814), ('目前', 9451), ('彩民', 9302), ('11', 8894)])]
work time used:497.6807523

时间上主要是文件读取和分词花了很长时间。
确认一下眼神，在做这个的时候，顺便优化了一下spark调用结巴分词，在每个task中都会从磁盘中加载模型文件的问题，每次从磁盘中加载都会消耗接近1s的时间。具体可以看这个：https://blog.csdn.net/macanv/article/details/87860691 如果你想优化这个时间，将上面的seg_with_dir_name函数替换成下面的：

def seg_with_dir_name(dir_name, data):
    """
    分词后返回分词的dataframe
    :param spark:
    :param data:
    :return:
    """
    if data is None:
        return []
    seg_model = jieba.Tokenizer(FREQ=broadcast_FREQ.value, total=broadcast_total.value)
    return [((dir_name, w), 1) for w in seg_model.cut(data.strip(), cut_all=False) if len(w) > 1]

并且在main入口中添加两个广播变量：

broadcast_FREQ = sc.broadcast(FREQ)
broadcast_total = sc.broadcast(total)

4. 不考虑分组的排序

写了那么多的分组排序，那么不需要分组的排序呢？
直接上代码：

data = spark.read.text('topK').rdd \
    .map(lambda x: x[0].split(',')) \
    .map(lambda x: Row(labels=x[6] + x[5], id=x[1], score_1=int(x[2]), score_2=int(x[3]), score_3=int(x[4])))
data = spark.createDataFrame(data)
data.sort(fn.desc('score_1')).show()

看到木有，直接调用DF的sort方法即可，desc是降序的意思，fn在上面代码的import中，在sort内的col按照优先级排，例如：

data.sort([fn.desc('score_1'),fn.desc('score_2'),fn.asc('score_3')]).show()

表示优先按照score_1降序，如果score_1 相等的时候，按照score_2降序，如果score_1, score_2都相等的时候，按照socre_3的升序排列。