Python使用队列实现读写分离

场景：从文件读取数据，处理后写入文件。考虑到的处理方式有两种。这里使用的队列是最常见的先进先出队列。

方式一：

读取并处理数据 ——> 存入Queue ——> 从Queue取出并写入文件

对于处理过程比较简单的情况可以在取出数据的时候直接处理，然后将数据存入Queue。在读取数据时可以使用多进程来提高速度。

方式二：

读取数据并预处理 ——> 存入Queue1 ——> 从Queue1取出数据处理  ——>  存入Queue2  ——>  从Queue2取出并写入文件

对于处理过程比较繁琐的情况，可以在读取数据时预处理，然后将数据存入Queue1，再从Queue1取出数据处理后存入Queue2，最后将Queue2取出的数据写入文件。在读取数据时同样可以使用多进程来提高速度。

代码

下面是具体实现。多进程使用的是multiprocessing.Pool.apply_async。省去了数据处理过程，只是用了一个队列。在写入文件的时候，防止一个文件数据量多过，写入同一文件的数据条数达到设定值后会打开一个新的文件写入数据。

import json
import os
import time
from threading import Thread
from multiprocessing import Pool, Manager


class ReadAndWrite():
    def __init__(self, input_path, output_path):
        self.queue_1 = Manager().Queue()
        self.read_end = Manager().list()
        self.input_path = input_path
        self.output_path = output_path

    def read_file(self, pool):
        '''并行读取文件'''
        for _path, _, _files in os.walk(self.input_path):
            for _file_name in _files:
                if _file_name.endswith('.json'):
                    file_path = os.path.join(_path, _file_name)
                    pool.apply_async(self._read_file, (file_path,))

    def _read_file(self, file_path):
        '''读取单个文件，将内容放入队列'''
        with open(file_path, 'r+', encoding='utf8', errors='ignore') as f:
            self.queue_1.put(json.loads(f.read()), block=True)

    def write_file(self, file_size=1000):
        '''将队列中的数据写入文件'''
        _file_count = 1
        _file_size = 0
        _open_new_file = False
        if not os.path.exists(self.output_path):
            os.makedirs(self.output_path)
        f = open(os.path.join(self.output_path, '%03d.json'%_file_count), 'w', encoding='utf8')
        while True:
            if not self.queue_1.empty():
                if _open_new_file:
                    f.close()
                    _open_new_file = False
                    f = open(os.path.join(self.output_path, '%03d.json'%_file_count), 'w', encoding='utf8')
            if _file_size < file_size:
                f.write(json.dumps(self.queue_1.get(), ensure_ascii=False, sort_keys=True)+ '\n')
                _file_size += 1
            else:
                _file_size = 0
                _open_new_file = True
            if self.queue_1.empty() and self.read_end:
                break
        f.close()

    def execute(self):
        '''执行入口'''
        t = Thread(target=self.write_file(), args=())
        t.start()
        pool = Pool(3)
        self.read_file(pool)
        pool.close()
        pool.join()
        self.read_end.append(1)


if __name__ == '__main__':
    start_time = time.time()
    input_path = r'/home/test/input'
    output_path = r'/home/test/output'
    read_and_write = ReadAndWrite(input_path, output_path)
    print('耗时：%s 秒'%(time.time() - start_time))

我这里使用了5g大小的json数据，测试对比，使用multiprocessing.Pool.map比multiprocessing.Pool.apply_async的耗时少很多，这里只修改了read_file()这个方法，其它部分不变，修改的代码如下：

class ReadAndWrite():
    def read_file(self, pool):
        '''并行读取文件'''
        file_path = []
        for _path, _, _files in os.walk(self.input_path):
            for _file_name in _files:
                if _file_name.endswith('.json'):
                    file_path.append(os.path.join(_path, _file_name))
        pool.map(self._read_file, file_path)