PySpark大数据处理性能优化指南

Apache Spark（https://spark.apache.org/）凭借其在分布式数据处理方面的强大能力，近年来一直是最领先的分析引擎之一。PySpark，即Spark的Python API，经常被用来解决个人和企业项目中的数据挑战。例如，我们可以通过PySpark高效地进行时间序列数据的特征工程，包括数据摄入、提取和可视化。然而，在面对诸如极端数据分布和复杂数据转换流程等情况时，尽管它能够处理大量数据集，仍然可能遇到性能瓶颈。

本文将探讨在Databricks上使用PySpark进行数据处理时常见的性能瓶颈，并介绍各种优化策略以实现更快的运行速度。

这张照片由 Veri Ivanova 在 Unsplash 上拍摄。

因此，你计划通过分析现有的交易数据，更好地服务现有顾客并吸引更多新顾客。这让你决定花大量时间来整理交易记录，作为准备工作。

我们首先在一个CSV文件中模拟了100万条交易记录（在现实的大数据场景中，我们预计会处理更大规模的数据集）。每条记录包含客户ID、购买商品以及支付方式和总金额等交易详情。特别值得注意的是，客户ID为#100的代理商拥有大量客户，因此，这位代理商为您店铺的代发货业务贡献了很大一部分订单。

下面是一些代码示例，展示该场景：

    import csv  
    import datetime  
    import numpy as np  
    import random  

    # 如果有已存在的 'retail_transactions.csv' 文件，删除它  
    ! rm -f /p/a/t/h retail_transactions.csv  

    # 设置交易次数以及其他设置  
    no_of_iterations = 1000000  
    data = []  
    csvFile = 'retail_transactions.csv'  

    # 打开文件csvFile以写入模式，并设置newline=''参数:  
    with open(csvFile, 'w', newline='') as f:  

       # 定义字段名为：orderID, customerID, productID, state, paymentMthd, totalAmt, invoiceTime  
       fieldnames = ['orderID', 'customerID', 'productID', 'state', 'paymentMthd', 'totalAmt', 'invoiceTime']  
       writer = csv.DictWriter(f, fieldnames=fieldnames)  
       writer 写入表头  

       for num in range(no_of_iterations):  
         # 生成一条具有随机值的交易记录  
         new_txn = {  
         'orderID': num,  
         'customerID': random.choice([100, random.randint(1, 100000)]),  
         'productID': np.random.randint(10000, size=random.randint(1, 5)).tolist(),  
         'state': random.choice(['CA', 'TX', 'FL', 'NY', 'PA', 'OTHERS']),  
         'paymentMthd': random.choice(['Credit card', 'Debit card', 'Digital wallet', 'Cash on delivery', 'Cryptocurrency']),  
         'totalAmt': round(random.random() * 5000, 2),  
         'invoiceTime': datetime.datetime.now().isoformat()  # 将当前时间格式化为ISO格式  
         }  

         data.append(new_txn)  

      writer 写入数据行

在对数据进行模拟后，我们使用Databricks提供的Jupyter Notebook将CSV文件加载到PySpark DataFrame中。

    # 设置文件的位置和类型  
    file_location = "/FileStore/tables/retail_transactions.csv"  
    file_type = "csv"  

    # 定义 CSV 的选项  
    schema = "orderID INTEGER, customerID INTEGER, productID INTEGER, state STRING, 支付方式 STRING, totalAmt DOUBLE, invoiceTime TIMESTAMP"  
    first_row_is_header = "True"  
    delimiter = ","  

    # 将 CSV 文件读入 DataFrame  
    df = spark.read.format(file_type) \  
      .schema(schema) \  
      .option("header", first_row_is_header) \  
      .option("delimiter", delimiter) \  
      .load(file_location)  

我们还创建了一个可重复使用的装饰器工具，用于测量并比较不同策略在各个函数中的执行时间。

    import time  

    # 测量给定函数的执行时间  
    def time_decorator(func):  
      def wrapper(*args, **kwargs):  
        begin_time = time.time()  
        output = func(*args, **kwargs)  
        end_time = time.time()  
        print(f"函数 {func.__name__} 的执行时间是大约 {round(end_time - begin_time, 2)} 秒。")  
      return wrapper  
    return wrapper

好了，所有的准备完成了。我们来看看接下来的部分中可能遇到的执行性能挑战。

Spark 使用 弹性分布式数据集（RDD） 作为其核心构建模块，数据通常默认存放在内存中。无论是执行计算（如连接和聚合等）还是在集群中存储数据，所有操作都会增加统一的内存使用区域中的内存占用。

一个具备执行内存（RAM）和存储内存（ROM）的统一区域（图片来源：作者）

如果我们设计不合理，可用内存可能不够用，这样会导致额外的数据被写入磁盘，从而导致性能变慢。

缓存和持久化中间结果或频繁访问的数据集是一种常见做法。虽然缓存和持久化都服务于相同的目的，但它们在存储级别上可能有所不同。应优化资源使用，确保读写操作的高效性。

例如，如果转换数据将在不同后续阶段反复用于计算和算法，最好将该数据缓存起来。

代码示例： 假设我们想要以数字钱包作为支付手段来研究不同交易记录的子集。

• 效率不高 — 没有缓存

    导入 pyspark.sql.functions 中的 col 函数

    @time_decorator
    # 定义一个不使用缓存的函数来过滤数据
    def without_cache(data):
      # 第一次过滤数据
      df2 = data.where(col("paymentMthd") == "Digital wallet")
      count = df2.count()
      # 计算df2的条目数

      # 第二次过滤数据
      df3 = df2.where(col("totalAmt") > 2000)
      count = df3.count()
      # 计算df3的条目数

      return count

    显示(without_cache(df))

• 高效 — 对重要数据集缓存

    从 pyspark.sql.functions 导入 col  

    @time_decorator  
    def after_cache(data):  
      # 第一次过滤并缓存  
      df2 = data.where(col("paymentMthd") == "Digital wallet").cache()  
      count = df2.count()  

      # 进行第二次过滤  
      df3 = df2.where(col("totalAmt") > 2000)  
      count = df3.count()  

      return count  

    显示(after_cache(df))

缓存后，即使我们想用不同的交易金额或其他数据维度来过滤转换后的数据集，执行时间也会更加容易管理。

当我们执行像连接 DataFrames 或按数据字段分组这样的操作时，会发生数据重新排序。这就需要将所有记录重新分配到集群中的各个节点，并确保相同键的记录被分配到同一节点上。这反过来使同时处理和合并结果成为可能。

洗牌（即数据混洗联接）（图片来源：作者）

然而，这个乱序操作成本高昂——执行时间较长，并且数据在节点间移动造成的额外网络开销。

为了减少洗牌次数，这里有几个方法：例如，

(1) 使用广播变量来分发小数据集，给每个工作节点发送一份只读副本，用于本地处理

虽然通常将“小”数据集定义为每个执行器的最大内存限制为8GB，但广播数据的理想规模应通过具体情况进行实验来确定。

广播接入（图片由作者提供）

尽早进行过滤，以尽早减少需要处理的数据量。

(3) 控制分区数以确保最佳性能

代码示例： 比如我们要返回与状态清单匹配的交易明细以及它们的完整名称

• 大数据集和小数据集之间的shuffle join效率低

    from pyspark.sql.functions import col  

    @计时装饰器  
    def no_broadcast_var(data):  
      # 创建小数据集  
      small_data = [("CA", "California"), ("TX", "Texas"), ("FL", "Florida")]  
      small_df = spark.createDataFrame(small_data, ["state", "stateLF"])  

      # 执行连接  
      result_no_broadcast = data.join(small_df, "state")  

      return result_no_broadcast.计数()  

    显示(no_broadcast_var(df))

• 高效 — 用广播变量把小数据集和大数据集连接起来

    from pyspark.sql.functions import col, broadcast  # 导入 col 和 broadcast  

    @time_decorator 装饰器  
    def have_broadcast_var(data):  
      small_data = [("CA", "加利福尼亚"), ("TX", "德克萨斯"), ("FL", "佛罗里达")]  
      small_df = spark.createDataFrame(small_data, ["state", "stateFullName"])  

      # 创建广播变量并执行连接操作  
      result_have_broadcast = data.join(broadcast(small_df), "state")  # 执行连接操作  

      return result_have_broadcast.count()  

    display(have_broadcast_var(df))  # 显示结果  

数据有时分布不均，特别是在作为处理关键的数据字段上。这会导致分区大小不均衡，有些分区的大小明显大于或小于平均值。

由于执行性能受限于执行时间最长的任务，必须处理过载的节点。

一种常见的方法是加盐。这种方法通过向偏斜的键添加随机数来实现，从而使各个分区中的分布更加均匀。当我们根据偏斜键聚合数据时，我们将使用加盐后的键进行聚合，然后再根据原始键进行聚合。另一种方法是增加分区数量，以使数据分布更均匀。

数据分布情况 — 加盐处理前后 （图片由作者提供）

代码示例: 我们想要汇总一个主要因为客户ID #100 而变得不均衡的不对称数据集。

• 效率低下 — 直接用了歪斜的密钥

    从pyspark.sql.functions导入col和desc  

    @time_decorator  
    定义no_salting(data):  
      # 执行聚合  
      agg_data = data.groupBy("customerID").agg({"totalAmt": "sum"}).orderBy(desc("sum(totalAmt)"))  
      返回agg_data  

    打印(no_salting(df))

• 高效 — 使用加盐偏移键来聚合数据

    from pyspark.sql.functions import col, lit, concat, rand, split, desc  

    @time_decorator  
    def 进行加盐处理(data):  
      # 对customerID进行加盐处理，通过添加后缀  
      salted_data = data.withColumn("salt", (rand() * 8).cast("int")) \  
                    .withColumn("saltedCustomerID", concat(col("customerID"), lit("_"), col("salt")))  

      # 执行聚合操作  
      agg_data = salted_data.groupBy("saltedCustomerID").agg({"totalAmt": "sum"})  

      # 移除加盐部分，以便进行进一步的聚合操作  
      final_result = agg_data.withColumn("customerID", split(col("saltedCustomerID"), "_")[0]).groupBy("customerID").agg({"totalAmt": "sum"}).sort(desc("sum(totalAmt)"))  

      # 按totalAmt总和降序排序  
      return final_result  

    显示(进行加盐处理后的结果(df))

对于偏斜的键值，随机添加前缀或后缀都可以。通常，5到10个随机值（如5到10个）是一个不错的开始点，能够在分散数据和保持高复杂度之间找到平衡。

人们常常更喜欢使用user-defined functions (UDFs)，因为它可以灵活地定制数据处理逻辑，这让用户可以根据需要调整数据处理方式。然而，UDFs是针对每一行数据进行操作的。这些代码会被Python解释器序列化，然后发送到执行器的JVM中进行反序列化。这不仅增加了序列化的成本，还使得Spark无法对代码进行优化处理，从而影响了处理效率。

尽量不要用UDF，这样更直接简单。

我们应该首先考虑使用内置的Spark功能，这些功能可以处理诸如聚合、数组/映射操作、日期/时间戳以及JSON数据处理等任务。如果内置功能确实无法满足您的需求，我们可以考虑使用pandas UDFs。这些UDFs基于Apache Arrow构建，与普通的UDF相比，它们在开销上更低，性能上更高。

代码示例: 交易价格会根据原州折扣一下。

• 效率低 — 用了 UDF

    from pyspark.sql.functions import udf  
    from pyspark.sql.types import DoubleType  
    from pyspark.sql import functions as F  
    import numpy as np  

    # UDF来计算折扣金额  
    def calculate_discount(state, amount):  
      if state == 'CA':  
        return amount * 0.90  # 10%折扣  
      else:  
        return amount * 0.85  # 15%折扣  

    discount_udf = udf(calculate_discount, DoubleType())  

    @time_decorator  
    def have_udf(data):  
      # 使用 UDF  
      discounted_data = data.withColumn("discountedTotalAmt", discount_udf("state", "totalAmt"))  

      # 显示结果：  
      discounted_data.select("customerID", "totalAmt", "state", "discountedTotalAmt").show()  

    display(have_udf(df))  

• 高效 — 利用内置的 PySpark 函数

    from pyspark.sql.functions import when  

    @time_decorator  
    def no_udf(data):  
      # 定义一个不使用 UDF 的函数  
      # 使用 when 和 otherwise 根据不同条件对总金额进行折扣  
      discounted_data = data.withColumn(  
      "discountedTotalAmt",  
      when(data.state == "CA", data.totalAmt * 0.90)  # 10% 的折扣  
      .otherwise(data.totalAmt * 0.85))  # 15% 的折扣  

      # 显示结果如下  
      return discounted_data.select("customerID", "totalAmt", "state", "discountedTotalAmt").show()  

    display(no_udf(df))  # 显示处理后的数据

在这个例子中，我们使用内置的PySpark函数“when和otherwise”来依次检查多个条件。这样的例子不胜枚举。pyspark.sql.functions.transform 是自PySpark 3.1.0版本起提供的函数，用于帮助对输入数组中的每个元素进行转换。

如存储部分所述，当内存中无法容纳所有所需数据时，会通过将临时数据从内存写入磁盘来发生溢出。我们讨论过的许多性能问题都与溢出有关。例如，操作在分区间洗牌大量数据时，这很容易导致内存不足和溢出。

内存不足导致的不同错误场景（图片来源：作者）

在Spark UI中检查性能指标非常重要。如果我们注意到内存溢出和磁盘溢出的统计数据，这意味着溢出可能是导致任务长时间运行的原因。要解决这个问题，可以尝试增加每个工作节点的内存，例如，通过调整spark.executor.memory配置值来增加执行器内存大小；或者，我们可以通过配置spark.memory.fraction来调整分配给执行和存储的内存比例。

我们遇到了一些影响 PySpark 性能下降的常见原因以及可能的改进。

• 存储：使用缓存和持久化来保存经常访问的中间结果
• 数据重组：对于小数据集，使用广播变量来方便Spark的本地处理
• 偏斜：执行添加盐或重新分区数据以更均匀地分布偏斜的数据
• 序列化：优先使用内置的Spark函数来提高性能
• 溢出处理：调整配置值以明智地分配内存

最近一段时间，自适应查询执行（AQE） 针对动态规划和重新规划进行了优化。这支持查询执行过程中不同的再优化功能，形成了一种非常有效的优化技术。然而，在设计初期理解数据特性仍然很重要，因为它有助于编写更有效的代码和查询，从而利用 AQE 进行微调。

如果你喜欢这篇读物，我邀请你订阅我的Medium页面和关注领英页面的动态。你可以随时了解与数据科学项目实践、MLOps示例以及项目管理方法相关的精彩内容。