在分布式环境Spark中关闭jieba延时加载等优化方法 （3）

在分布式环境中关闭jieba延时加载的方法

这篇博客是记录在使用spark+jieba的过程中，每个task调动jieba分词的时候需要进行延时加载或者从缓存中读取模型文件从而浪费时间问题：

Building prefix dict from the default dictionary ...
Loading model from cache C:\Users\C\AppData\Local\Temp\jieba.cache
Loading model cost 0.555 seconds.
Prefix dict has been built succesfully.

1. 提前加载模型

根据官方的方法，可以使用下面的方法关闭延时加载：

import jieba
jieba.initialize()  # 手动初始化（可选）

这样就会提前将词典写入文件中，下次在调用的时候就可以直接从文件中读取模型了，这样可以节约一部分时间，但是从文件中读取模型任然需要花时间，而且对于一个分布式程序而言，我们很多时候是不会动态的继续加词典的，那么我们是否可以将模型文件放在内存中，每个task在调用分词的时候，直接从内存中加载模型文件，而不是从磁盘中加载呢。

2. 从内存中加载模型

通过阅读jieba的__init__.py 文件我们可以看到加载的模型文件主要是一个词典，和词典的大小，分别是一个dict和int类型的数据，我们可以通过上面的方法提前加载模型，然后让initialize方法返回这两个数据，在新的task中调用cut方法的时候，我们把这两个模型数据直接传入到对象中，从而避免了对象在initlizale方法中从磁盘中加载模型了。修改__init__.py代码很少，主要在init函数和initialize方法中（修改前先备份一下吧，不然。。。）：

init 方法：

    def __init__(self, dictionary=DEFAULT_DICT, FREQ=None, total=None):
        self.lock = threading.RLock()
        if dictionary == DEFAULT_DICT:
            self.dictionary = dictionary
        else:
            self.dictionary = _get_abs_path(dictionary)
        if FREQ is None:
            self.FREQ = {}
        else:
            self.FREQ = FREQ
        if total is None:
            self.total = 0
        else:
            self.total = total
        
        self.user_word_tag_tab = {}
        
        if FREQ is not None and FREQ is not None:
            self.initialized = True
        else:
            self.initialized = False
        self.tmp_dir = None#os.path.abspath('.')
        self.cache_file = None

可以看到添加了两个模型文件，可以通过外部传入。
initialize方法几乎没改，在最后添加三个返回值：

            self.initialized = True
            default_logger.debug(
                "Loading model cost %.3f seconds." % (time.time() - t1))
            default_logger.debug("Prefix dict has been built succesfully.")
        #添加三个返回值，上面的代码啥也没改...
        return self.FREQ, self.total,self.initialized

使用Spark调用

test中得数据可以参考一下之前的博文PySpark TF-IDF计算（2）

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

"""

 @Time    : 2019/2/21 14:23
 @Author  : MaCan (ma_cancan@163.com)
 @File    : jieba_test.py
"""

from pyspark import SparkContext, SparkConf, Row
from pyspark.sql import SparkSession


import os

import time
import random
import jieba
# initialize segment model
cut_model = jieba.Tokenizer()
FREQ, total = cut_model.initialize()


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


def seg(data):
    """
    使用原始的jieba
    :param data:
    :return:
    """
    return [w for w in jieba.cut(data[0].strip(), cut_all=False) if len(w) > 1]


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:
        # 使用两个广播变量来广播jieba的两个模型数据
        broadcast_FREQ = sc.broadcast(FREQ)
        broadcast_total = sc.broadcast(total)
        rdd_data = spark.read.text('test').rdd.flatMap(seg_with_model)
        print(rdd_data.collect())
    except Exception as e:
        print(e)
        exit(-1)
    finally:
        print('work time used:{}'.format(time.perf_counter() - start_t))
        spark.stop()

输出：

['来到', '北京', '清华大学', '来到', '网易', '杭研', '大厦', '来到', '北京', '清华大学', '来到', '网易', '杭研', '大厦', '来到', '北京', '清华大学', '来到', '北京', '清华大学']
work time used:7.5783597

我们测试一下不适用这样方式的方法,我们新写了一个分词的函数seg，方法里面是直接调用jieba.cut进行分词。
同时在mian入口函数中flatMap中得seg_with_model替换为seg,注释掉两个广播变量

def seg(data):
    """
    使用原始的jieba
    :param data:
    :return:
    """
    return [w for w in jieba.cut(data[0].strip(), cut_all=False) if len(w) > 1]

运行输出：

[Stage 0:>     (0 + 1) / 1]Building prefix dict from the default dictionary ...
Loading model from cache C:\Users\C\AppData\Local\Temp\jieba.cache
Loading model cost 0.592 seconds.
Prefix dict has been built succesfully.
['来到', '北京', '清华大学', '来到', '网易', '杭研', '大厦', '来到', '北京', '清华大学', '来到', '网易', '杭研', '大厦', '来到', '北京', '清华大学', '来到', '北京', '清华大学']
work time used:8.2881966

我们看到在stage中打印了jieba日志：加载模型的log输出，同时运行时间多了大约0.7s,注意看jieba的日志中它load 模型花了0.59s。再加上jieba中initiale方法中得其他逻辑操作的时间消耗，个人感觉时间还是可以对的上的。如果数据量很大的话，每个task都从磁盘中加载jieba的模型文件，这个时间将会更加的明显。