最近在工作中遇到了一个需求，在方案设计上，考虑存在并行处理任务的需要，于是尝试学习和调用 multiprocessing 库，多进程操作任务。此前背诵八股文，问及并行和并发相关的问题，肌肉记忆脱口而出。不过真正需要用到相应知识的时候，只能感叹一声：这知识啊，它不进脑子啊！刚好，趁这次机会，学会并行编程在工程中落地，在此简单记录一下，算是温故而知新吧！

问题

基于 PySpark 处理数据，输入数据源及其某一列数据的转换条件，依据转换条件修改数据源中对应列的数据，保存并输出。同时，要求既可以输入单个数据源，也可以输入多个数据源。

（作者注：出于合规考虑，本文只描述技术细节，工作业务相关内容隐去不表...)

问题分析

我们首先需要考虑两个问题，其一是使用 PySpark 转换数据，其二是处理多个数据源输入。对于数据处理，PySpark 提供了丰富的 API，我们既可以使用类似 SQL Case-When 语句的相关方法，也可以通过创建 RDD 使用 mapPartitions 方法转换数据。相较之下，后者更为灵活。我们可以通过两个简单的例子快速过一下二者的使用姿势。

Spark 数据转换

首先，创建一个简单的 DataFrame 作为数据示例，代码如下所示：

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("playground").getOrCreate()
data = [("James","M",60000),("Michael","M",70000),
        ("Robert",None,400000),("Maria","F",500000),
        ("Jen","",None)]
columns = ["name","gender","salary"]
df = spark.createDataFrame(data = data, schema = columns)
df.show()

借助 Jupyter Notebook 等工具运行上述代码，可以打印得到一个 DataFrame：

+---------+--------+--------+
|    name | gender | salary |
+---------+--------+--------+
|   James |      M |  60000 |
| Michael |      M |  70000 |
|  Robert |   null | 400000 |
|   Maria |      F | 500000 |
|     Jen |        |   null |
+---------+--------+--------+

此时，若是我们希望将性别所代表的缩写值修改为性别词，那么就可以通过 when 方法，来达到 SQL 语言中类似 Case-When 语句的效果：

from pyspark.sql.functions import when


df2 = df.withColumn("new_gender", when(df.gender == "M","Male")
                                 .when(df.gender == "F","Female")
                                 .when(df.gender.isNull() ,"")
                                 .otherwise(df.gender))
df2.show()

查看打印输出，可得：

+-------+------+------+----------+
|   name|gender|salary|new_gender|
+-------+------+------+----------+
|  James|     M| 60000|      Male|
|Michael|     M| 70000|      Male|
| Robert|  null|400000|          |
|  Maria|     F|500000|    Female|
|    Jen|      |  null|          |
+-------+------+------+----------+

方法 when 接受两个参数，其一为数据转换判定条件，其二是数据转换目标值。熟悉 SQL 的旁友对于 Spark 中的 when 想必接受很快。另一种数据转换的方式是通过 RDD 的 mapPartitions 方法，此时需要定义一个函数，用于描述转换规则。比如同样是如上转换，可以这样：

def transform(rows):
    mapping = {"M": "Male", "F": "Female"}
    
    for row in rows:
        if not isinstance(row, dict):
            instance = row.asDict(recursive=True)
        else:
            instance = row
            
        if row["gender"] is None:
            row["new_gender"] = ""
            yield row
        
        if row["gender"] in ["M", "F"]:
            row["new_gender"] = mapping[row["gender"]]
            yield row
            
        row["new_gender"] = row["gender"]
            
        yield row

df3 = df.rdd.mapPartitions(transform)

上述代码同样能够打印输出与 df2 一致的内容。了解了 Spark 数据转换的部分使用姿势之后，我们再来考虑如何处理多个数据源输入这一问题。

多个数据源处理

直觉上，我们很自然地会想到将输入定义为一个数据源或者一个数据源列表（假定存在一个函数，使得我们能够根据数据源名称，获取数据源内容），则输入类型定义如下：

from typing import List, Union

DataSource = Union[str, List[str]]

为了能够同时处理一个或多个数据源，我们会统一将数据源转换为 List 类型值，而后遍历每一个数据源，并完成数据转换，比如：

def convert(data_source: DataSource) -> DataFrame:
        data_source = data_source if isinstance(data_source, list) else [datasource].
        
        for source in data_source:
            # ...code
        
        return # ...code

由上可见，我们能够轻而易举地完成数据转换，然而一个要命的因素是任务耗时。若是存在大量数据源，需要做同样的数据转换，则需要持续等待每一个任务完成。而这恰好是并行编程的用武之地。

Python 多进程

碰到运行多个任务且单个任务耗时长的场景，我们脑中马上会蹦出一个词 ——「并行」，与之相伴的，还有另一个词 —— 「并发」。二者看起来十分相似，望文生义之下也十分容易相互混淆。不过二者之间的本质区别是，前者表示不同任务同时运行，后者表示不同任务交替运行。无论是哪一种，比起一个接一个运行任务，在降低耗时这方面都是巨大的进步。但如果资源充裕的情况下，任务能够同时运行，何乐而不为呢？

就 Python 而言，我们可以通过内置的 multiprocesssing 库来实现并行编程。首先，我们来看一下如何启动多个进程：

from multiprocessing import Process

def proc(i: int) -> None:
    print("I'm Process {}".format(i))
    
if __name__ == "__main__":
    for i in range(10):
        Process(target=proc, args=(i,)).start()

如上所示，我们创建了一个非常小的任务 proc，接受一个数字类型的参数，并打印一段字符串。接下来，我们通过实例化 Process 类来创建进程，传入 target 和 args 参数，分别表示进程内运行的任务，及任务运行所需的参数；而后调用 start 方法启动进程。

那么问题来了，假如我们希望多个进程能够互相通信，要怎么办呢？multiprocesssing 库提供了另一个类，用于进程之间的通信，即队列（Queue）。我们来简单看一下它的用法：

from multiprocessing import Process, Queue

def proc(i: int, queue: Queue) -> None:  
    message = f'I am Process {i}'  
    queue.put(message)

if __name__ == "__main__":
    queue = Queue()
    
    for i in range(10):
        Process(target=proc, args=(i, queue,)).start()
        
    for i in range(10):  
        message = queue.get()  
        print(message)

首先，我们创建了一个队列实例，并将其作为参数传给任务 proc，而在任务运行时，使用 put 方法将消息加入队列；同样地，我们可以使用 get 方法将消息从队列中取出来。需要注意的是，消息加入队列这一操作是在不同的进程中进行的，属于异步操作；而将消息取出队列则是同步操作，防止消息丢失。

看起来万事俱备，然而我们尚欠一股东风。很显然，我们不可能同时启动不限数量的进程运行任务，那样多半是要出问题的。不过好在 multiprocesssing 库提供了 Pool，它允许我们创建无数进程，而令部分进程始终保持活跃状态，当然这取决于 Pool 的大小，比如：

from multiprocessing import Pool, Process, Queue

def proc(i: int, queue: Queue) -> None:  
    message = f'I am Process {i}'  
    queue.put(message)

def consume(i: int, queue: Queue) -> None:
    message = queue.get()
    print(message)
    
if __name__ == "__main__":
    queue = Queue()
    
    for i in range(10):
        Process(target=proc, args=(i, queue,)).start()
        
    pool = Pool(10)
    
    consumers = []
    for i in range(10):  
        consumers.append(pool.apply_async(consume, (i, queue,)))
        
    [r.get() for r in readers]

看起来相当优雅，然而真正运行上述代码，那么就会得到一个 Exeption:

RuntimeError: Queue objects should only be shared between processes through inheritance.

也就是说，直接创建并启动多个不同的进程，进程之间的通信在这里是行不通的。莫慌，我们可以引入一个新的类 Manager，专门用于创建不同进程之间共享的数据，于是我们可以稍作修改：

from multiprocessing import Manager, Pool, Process, Queue

def proc(i: int, queue: Queue) -> None:  
    message = f'I am Process {i}'  
    queue.put(message)

def consume(i: int, queue: Queue) -> None:
    message = queue.get()
    print(message)
    
if __name__ == "__main__":
    manager = Manager()
    
    queue = manager.Queue()
    for i in range(10):
        Process(target=proc, args=(i, queue,)).start()
        
    pool = Pool(10)
    
    consumers = []
    for i in range(10):  
        consumers.append(pool.apply_async(consume, (i, queue,)))
        
    [r.get() for r in readers]

如此，我们就能够使用多个不同的进程跑相同的任务，并共享任务数据啦！回到我们之前的问题，对于多个数据源来说，数据转换的条件是一致的，我们完全可以启动多个进程同时处理数据，由此能够大大降低任务耗时。

PySpark 多进程数据处理

前文假定我们能够通过数据源名称，获取数据内容。此处我们可以定义得更具体些，通过数据源名称获取数据对应的存储位置（甭管怎么获取的）:

def get_datasource_path(datasource_name: str) -> str:
    # ...
    return path

定义任务

既然是多进程处理数据，那我们先把任务本身定义清楚，即基于 PySpark 库，给定转换条件，输出转换后的数据。此处省略数据转换相关逻辑，本身也不复杂，基于 PySpark 提供的方法，活学活用即可：

from pyspark.sql import DataFrame, SparkSession

spark = SparkSession.builder.appName("DataProcessing").getOrCreate()

def task(path: str) -> DataFrame:
    df = spark.read.json(path)
    
    # ...
    return df

多进程处理任务

结合上文 Python 多进程相关知识，我们需要编写三个函数，其一创建并启动多进程处理任务，其二收集任务运行结果，其三管理前俩函数。额外需要注意的是，由于我们使用多进程处理不同参数的同一任务，最后若是需要将所有任务结果合并输出，我们或需要再写一个函数，用于合并任务输出结果。

from multiprocessing import Manager, Pool, Process, Queue, cpu_count
from pyspark.sql import DataFrame

def dispatch(queue: Queue, path: str) -> None:
    df = task(path)
    queue.put(df)
    
def listen(queue: Queue) -> DataFrame:
    df = queue.get()
    return df
    
def execute() -> DataFrame:
    manager = Manager()
    
    queue = manager.Queue()
    
    datasource_path_list = [
        get_datasource_path(datasource) for datasource in datasource_list
    ]
    for path in datasource_path_list:
        process = Process(target=dispatch, args=(queue, path,))
        process.start()
    
    pool = Pool(cpu_count())
    listeners = [
        pool.apply_async(listen, (queue,)) for _ in range(len(datasource_list))
    ]
    
    try:
        results = [listener.get() for listener in listeners]
        return merge(results)
    except:
        pool.terminate()
        pool.join
        
    return

此处，我们使用 multiprocesssing 提供的 cpu_count 方法获取 CPU 数量，作为 Pool 大小，在追求任务性能的同时，充分利用资源。但若我们直接运行上述代码，结果并不如人意，会得到大量 Spark 报错，部分任务完全执行不了。坦白讲，我并没有找到报错的真实原因，只能猜想不能同时启动 Spark 任务，于是在启动之前，给一个等待时间，错开任务启动，虽然耗时有所延长，但比起任务失败，能跑起来且近乎同时运行这一结果，还是很可以接受的：

import time

def execute() -> DataFrame:
    # ...
    for path in datasource_path_list:
        time.sleep(15)
        process = Process(target=dispatch, args=(queue, path,))
        process.start()
    # ...

好啦，以上就是一次 PySpark 多进程编程实践的记录。其实一开始困于多进程启动 Spark 任务部分能跑部分不能跑的问题，一直在搜如何使用 Pool 启动多个 Spark 任务。最后虽然也没有找到原因，但摸索出来能用的办法，老实用主义了。希望能帮到有需要的同学，以上！