如何使用 Salting 技术解决 Apache Spark 中的数据倾斜问题

如何使用 Salting 技术解决 Apache Spark 中的数据倾斜问题

在处理大规模数据集时，Apache Spark 中一个常见的性能问题是数据倾斜（Data Skew）。这通常发生在某些键（key）占据了数据的大部分分布，导致分区不均，从而影响查询性能。这类问题主要出现在需要 shuffle 的操作中，例如 join 或常规的 聚合操作（aggregation）。

一种实用的解决数据倾斜的方法是 Salting（加盐）。它的基本思想是人为地将“热点 key”分散到多个分区中。本文将通过一个实际示例来介绍如何使用该技术。

---

Salting 如何解决数据倾斜问题

通过在 join 的 key 上添加一个随机数（salt），构造一个新的复合 key，Spark 可以更均匀地将原本聚集在一个 key 上的数据分布到多个分区。这样可使数据更加均衡地分布在多个执行器（executors）中，避免某个节点负载过重而其他节点空闲的情况。

---

使用 Salting 的好处

• 降低数据倾斜：将大 key 分散处理，避免少数执行器过载，提高整体资源利用率。

• 提升性能：通过 workload 平衡，加快 join 和聚合操作的执行速度。

• 减少资源争用：防止因某个分区过大导致的内存溢出等异常。

---

何时使用 Salting

当你发现某些 join 或 aggregation 操作耗时较长，或发生 executor 异常失败，很可能是数据倾斜导致的。此时使用 salting 是一种有效的解决方法。它也适用于流处理（streaming）场景中需要精细化分区的数据处理。

---

Salting 示例（Scala）

假设我们有两个要 join 的数据集，一个很大，一个很小，且大表中某些 key 的分布非常不均。

scalaimport org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._

// 模拟大数据集
val largeDF = Seq(
  (1, "txn1"), (1, "txn2"), (1, "txn3"), (2, "txn4"), (3, "txn5")
).toDF("customer_id", "transaction")

// 小数据集
val smallDF = Seq(
  (1, "Ahmed"), (2, "Ali"), (3, "Hassan")
).toDF("customer_id", "name")

第一步：为大数据集添加 salt 列

scalaval numBuckets = 3
val saltedLargeDF = largeDF.
  withColumn("salt", (rand() * numBuckets).cast("int")).
  withColumn("salted_customer_id", concat($"customer_id", lit("_"), $"salt"))

结果示例：

pgsql+-----------+-----------+----+------------------+
|customer_id|transaction|salt|salted_customer_id|
+-----------+-----------+----+------------------+
|          1|       txn1|   1|               1_1|
|          1|       txn2|   1|               1_1|
|          1|       txn3|   2|               1_2|
|          2|       txn4|   2|               2_2|
|          3|       txn5|   0|               3_0|
+-----------+-----------+----+------------------+

第二步：扩展小表，生成所有可能的 salted key

scalaval saltedSmallDF = (0 until numBuckets).toDF("salt").
  crossJoin(smallDF).
  withColumn("salted_customer_id", concat($"customer_id", lit("_"), $"salt"))

结果示例：

diff+----+-----------+------+------------------+
|salt|customer_id|  name|salted_customer_id|
+----+-----------+------+------------------+
|   0|          1| Ahmed|               1_0|
|   1|          1| Ahmed|               1_1|
|   2|          1| Ahmed|               1_2|
|   0|          2|   Ali|               2_0|
...

第三步：执行 join

scalaval joinedDF = saltedLargeDF.
  join(saltedSmallDF, Seq("salted_customer_id", "customer_id"), "inner").
  select("customer_id", "transaction", "name")

输出：

pgsql+-----------+-----------+------+
|customer_id|transaction|  name|
+-----------+-----------+------+
|          1|       txn2| Ahmed|
|          1|       txn1| Ahmed|
|          1|       txn3| Ahmed|
|          2|       txn4|   Ali|
|          3|       txn5|Hassan|
+-----------+-----------+------+

---

Salting 示例（Python）

pythonfrom pyspark.sql import SparkSession
from pyspark.sql.functions import col, rand, lit, concat
from pyspark.sql.types import IntegerType

# 模拟数据
largeDF = spark.createDataFrame([
    (1, "txn1"), (1, "txn2"), (1, "txn3"), (2, "txn4"), (3, "txn5")
], ["customer_id", "transaction"])

smallDF = spark.createDataFrame([
    (1, "Ahmed"), (2, "Ali"), (3, "Hassan")
], ["customer_id", "name"])

# 添加 salt 列
numBuckets = 3
saltedLargeDF = largeDF.withColumn("salt", (rand() * numBuckets).cast(IntegerType())) \
    .withColumn("salted_customer_id", concat(col("customer_id"), lit("_"), col("salt")))

# 扩展小表
salt_range = spark.range(0, numBuckets).withColumnRenamed("id", "salt")
saltedSmallDF = salt_range.crossJoin(smallDF) \
    .withColumn("salted_customer_id", concat(col("customer_id"), lit("_"), col("salt")))

# 执行 join
joinedDF = saltedLargeDF.join(
    saltedSmallDF,
    on=["salted_customer_id", "customer_id"],
    how="inner"
).select("customer_id", "transaction", "name")

---

调优建议：如何选择 numBuckets

• 如果设置 numBuckets = 100，每个 key 会被分成 100 个子分区。

• 然而，不要设置过大，否则会对本身就很小的数据产生额外开销。

• 通常建议从 10~20 开始，观察是否缓解 skew 问题，然后再逐步调整。

---

进阶技巧：仅对倾斜 key 使用 salting

你可以先识别出数据倾斜的 key，然后只对这些 key 进行 salting，其他保持默认：

scalaval saltedLargeDF = largeDF.
  withColumn("salt", when($"customer_id" === 1, (rand() * numBuckets).cast("int")).otherwise(lit(0))).
  withColumn("salted_customer_id", concat($"customer_id", lit("_"), $"salt"))

val saltedSmallDF = (0 until numBuckets).toDF("salt").
  crossJoin(smallDF.filter($"customer_id" === 1)).
  select("customer_id", "salt", "name").
  union(smallDF.filter($"customer_id" =!= 1).withColumn("salt", lit(0)).select("customer_id", "salt", "name")).
  withColumn("salted_customer_id", concat($"customer_id", lit("_"), $"salt"))

---

结语

当传统的分区策略或 Spark 的 SKEWED JOIN 提示不能解决数据倾斜时，Salting 是一种非常实用而有效的方法。合理调参和监控执行计划，可以大大提升 Spark 在处理数据倾斜场景下的性能表现。