Spark分区器详解:从内置分区到自定义Partitioner,彻底解决数据倾斜
但凡被线上Spark作业的耗时折磨过的人应该都见过这种场景集群几十个executor好几个在空转某一个CPU被打满日志里同一个stage的task快的一秒跑完慢的要等半个多小时。排查了半天SQL也没看出逻辑问题根源往往不在写法而在分区——更准确地说在分区的规则和分区器上。Spark里的分区规则严格来讲分两个层次一是框架内置的Hash、Range两种默认分区方式二是允许我们自己继承Partitioner抽象类实现任意分桶逻辑。而分区器Partitioner就是决定每条数据去哪个分区的那个“裁判”。搞懂了分区器很多数据倾斜问题根本不需要靠加资源硬扛改几行代码就能把耗时拉下来。这篇内容我会从分区和并行度的关系讲起把内置分区器源码拆开看一遍再手把手写一个自定义Partitioner最后把实践里踩过的坑一次性交代清楚。适合刚接触Spark但已经被数据倾斜折磨过的开发也适合想系统搞明白repartition、partitionBy、sortByKey背后到底发生了什么的同学。1. 分区是什么为什么它决定Spark作业的生死1.1 分区与并行度一个Task处理一个分区先纠正一个非常常见的误解不少人以为RDD的partition数等于数据条数或者以为分区越多数据就越“散”其实都不是。分区是Spark计算调度的最小单位。一个作业会被切分成多个stage每个stage里有一批task而每个task恰好处理一个partition。所以并行度的上限就是当前stage的分区数。假如你有一个10亿条的RDD却只有4个分区那么无论集群里有100个executor还是1000个executor真正同时干活的计算任务最多只有4个。剩下96个executor全在空转资源白费。反过来如果数据量很小却分了10万个大区光task调度和序列化的开销就能把作业拖垮。这个关系我建议所有Spark开发都刻在脑子里分区数决定了并行度并行度决定了这个stage能用多少资源而分区的数据分布决定了每个task是否均匀。三者凑在一起基本就决定了作业跑得快不快。1.2 分区从哪来输入、Shuffle、输出三个环节要排查分区问题首先得知道分区是在哪个环节产生的。按我自己的排查经验Spark里分区的来源主要有三个。第一是输入分区。比如用textFile读取HDFS上的文件时分区数由InputFormat决定通常是文件块大小决定的一个块一个分区。这跟你RDD设置了多少分区没有直接关系它只决定原始数据被切成多少片进入计算。第二是Shuffle分区。这是数据在算子之间跨节点流动时产生的。但凡涉及groupByKey、reduceByKey、join、distinct这类需要把相同key汇总到一起的操作一定会触发shuffle。而shuffle后每个分区有多少数据、key怎么分布就完全由分区器决定。绝大多数数据倾斜都发生在这个环节。第三是输出分区。写入文件时每个分区会对应生成一个文件这也决定了最终小文件的数量。比如你用df.write保存一张表分区数就等于最后文件数。很多人以为文件多是因为数据量大其实往往只是shuffle分区没调好。这三个环节的分区数可以各自独立设置但它们之间会互相影响。我见过不少人把repartition和写入分区数搞混导致输出了一堆只有几KB的小文件下游读起来想哭。后面我会专门讲这个坑。2. 内置分区规则Hash、Range和默认行为2.1 Hash分区默认规则快但险Hash分区是Spark用得最多的分区方式也是默认分区方式。原理非常简单对key做hashCode然后对分区数取模得到一个0到numPartitions-1之间的整数这个整数就是分区编号。// 伪代码示意 partitionIndex nonNegativeMod(key.hashCode, numPartitions)之所以叫nonNegativeMod是因为Java的hashCode可能返回负数如果不处理取模之后可能出现负分区号。Spark的做法是先对hashCode取绝对值或者加偏移量确保结果落在合法范围内。Hash分区的优点是效率高每个key只需要算一次hash不需要像Range分区那样做采样和排序。缺点是它只保证同一个key进同一个分区但并不保证数据均匀。如果你的key分布本身是倾斜的比如某个热门商品的ID出现频率特别高那么hash之后这个ID仍然会集中到同一个分区数据倾斜就出来了。2.2 Range分区为排序和范围查询而生Range分区是另一种内置规则它不像Hash那样把key散列到分区而是把key按大小顺序切分成若干个连续区间。分区0里的key都小于分区1里的key分区1里的key都小于分区2里的key以此类推。Spark实现Range分区的核心步骤是先对所有key做采样然后根据采样结果计算出一组分区边界值再把真实数据按边界值划分到不同的分区。所以Range分区的质量高度依赖采样质量。如果采样不均匀边界值就选得不准同样会产生倾斜。Range分区典型用在两个场景。一是全局排序比如sortByKey它需要每个分区的数据有序且分区之间的key范围也是递增的。二是repartitionByRange这个在Spark SQL里用得比较多用于把DataFrame按某个或多个字段的范围重新分区后续如果有按范围过滤的查询能显著减少扫描量。2.3 组件默认行为groupByKey、join、repartition各自用什么很多新手会困惑同一个RDD执行groupByKey和repartition用的分区器是不是同一个。答案是不一定。先看groupByKey、reduceByKey、aggregateByKey这类聚合算子。它们在做shuffle聚合时如果没有显式指定分区器Spark默认用HashPartitioner分区数取max父RDD分区数的较大值中的较大值以及spark.default.parallelism里的配置。这里对初学者很友好的一面是分区数不需要你手算Spark会尽量取一个合理值。再看sortByKey和repartitionAndSortWithinPartitions。sortByKey必须使用RangePartitioner因为只有Range分区才能保证全局有序。repartitionAndSortWithinPartitions则是一个很有用的组合算子它在重分区的同时按key排序特别适合需要在每个分区内做排序的场景。然后是repartition和coalesce。repartition(numPartitions)底层走的是shuffle机制重新散列数据所以它本质上也是Hash分区。coalesce(numPartitions, shuffle false)比较特殊它会尽量在不shuffle的情况下合并相邻分区好处是省了shuffle开销坏处是如果原始分区数据本身不均合并后会更不均。我做个简单的表格帮大家区分算子默认分区器是否触发Shuffle典型用途groupByKey / reduceByKeyHashPartitioner是按键聚合sortByKeyRangePartitioner是全局排序repartitionHashPartitioner是增加/减少分区coalesce无专门分区器默认否减少分区节省shufflepartitionBy手动指定是使用自定义分区器3. 内置分区器源码拆解HashPartitioner与RangePartitioner3.1 HashPartitioner几行代码里的门道HashPartitioner的源码非常短核心逻辑就三块numPartitions属性、getPartition方法、equals方法。我直接给出简化的结构class HashPartitioner(partitions: Int) extends Partitioner { require(partitions 0, sNumber of partitions must be positive but got $partitions) def numPartitions: Int partitions def getPartition(key: Any): Int key match { case null 0 case _ Utils.nonNegativeMod(key.hashCode, numPartitions) } override def equals(other: Any): Boolean other match { case h: HashPartitioner h.numPartitions numPartitions case _ false } }有几个细节值得注意。null key被特殊处理直接进0号分区。这其实是一个很关键的约定因为如果key为nullkey.hashCode会抛异常所以Spark干脆统一规定null key进第一个分区。nonNegativeMod的实现是(key.hashCode % numPartitions numPartitions) % numPartitions。很多自己写分区器的人容易忽略负数问题直接用hashCode % numPartitions结果线上会出现IllegalArgumentException报错说partition index out of range。equals方法的作用特别容易被低估。Spark在调用partitionBy或groupByKey时会判断RDD已有的分区器是否等于你传入的分区器。如果相等就不会触发shuffle。HashPartitioner的equals只比较分区数所以两个分区数相同的HashPartitioner会被认为是同一个分区器从而跳过shuffle。这个设计对性能影响非常大后面讲自定义分区器时我会强调为什么自己的实现也一定要重写equals。3.2 RangePartitioner采样与边界计算RangePartitioner比HashPartitioner复杂得多因为它需要知道数据的大致分布才能把key切成合理的范围。整个过程分两步第一步是采样第二步是根据采样结果计算边界。采样用的是水塘采样算法reservoir sampling主要目的是在不遍历全部数据的情况下随机抽取一批key作为代表。Spark把RDD划分成多个分区每个分区内部随机抽一定数量的样本最后汇总起来按key排序再从排序后的样本序列中均匀地挑出numPartitions-1个作为边界。class RangePartitioner(partitions: Int, rdd: RDD[_], private var ascending: Boolean true) extends Partitioner { val rangeBounds: Array[K] { val (_, sorted) RangePartitioner.sketch(rdd.mapPartitions(_.map(v (v, v))), partitions) // 从排序后的样本中取 numPartitions - 1 个边界 } def getPartition(key: Any): Int { // 二分查找 key 属于哪个区间 val partition binarySearch(rangeBounds, key) ... } }这里有个非常实际的影响RangePartitioner采样是带代价的而且采样结果不一定准。当数据量极大而采样比例不足时边界可能偏离真实分布导致某个分区数据特别多。我曾经遇到过sortByKey之后某个分区比其他分区大10倍的情况最后排查下来就是采样命中率太低。所以如果对排序后整体均匀性有硬要求可以考虑先做一次repartition基于Hash打散或者调整采样参数。不过大部分场景下RangePartitioner的默认表现已经足够。3.3 其他相关分区器Python中的partitionBy函数在PySpark里情况稍微有点不一样。PySpark没有直接暴露Scala的Partitioner抽象类但RDD有一个partitionBy方法它接受一个int类型的numPartitions和一个partitionFunc函数。rdd sc.parallelize([(a, 1), (b, 2), (c, 3)], 3) result rdd.partitionBy(2, lambda k: k[0])这里的partitionFunc其实是对key做一次预处理映射返回一个“新的key”然后用这个新key的hashCode再去取模。所以它本质上还是Hash分区只是允许你先对key做一层转换。比如你可以把手机号的运营商提取出来作为分区键让同一运营商的用户落进同一分区。另外Spark SQL里还有一组DataFrame APIrepartition按指定列做Hash分区repartitionByRange按指定列做Range分区。它们走的是与RDD分区器一样的底层逻辑只是表达层面更高适合在SQL分析法里直接用。4. 手把手实现自定义Partitioner4.1 正确理解Partitioner抽象类自定义分区器说难不难说简单也不简单关键是把抽象类搞懂。org.apache.spark.Partitioner是一个抽象类所有分区器都必须实现两个核心方法abstract class Partitioner extends Serializable { def numPartitions: Int def getPartition(key: Any): Int }numPartitions告诉Spark这个分区器会产生多少个分区。getPartition传入一个key返回该key对应的分区索引范围是0到numPartitions-1。除了这两个必须实现的方法我强烈建议同时重写equals和hashCode。原因前面已经说了Spark判断一个RDD是否已经“按某个分区器分区”时会调用分区器的equals方法。如果你不重写equals那么即使你两次用的是逻辑上完全一样的分区器Spark也会认为是两个不同对象进而多触发一次全量shuffle。这个坑我踩过不止一次线上作业多跑了十几分钟最后发现就是equals没写好。4.2 经典案例按用户ID段位分桶我分享一个我常用的案例。假设有一批用户行为日志key是用户ID我需要按用户ID的首字母或按ID区间切分落进不同的分区做后续处理。需求本身并不复杂但默认的Hash分区并不能满足我的要求因为我想让特定范围的ID进特定分区方便下游按分区读取。import org.apache.spark.Partitioner class UserIdPartitioner(numParts: Int) extends Partitioner { override def numPartitions: Int numParts override def getPartition(key: Any): Int key match { case null 0 case id: String if id.startsWith(A) 0 case id: String if id.startsWith(B) 1 case id: String if id.startsWith(C) 2 case _ (math.abs(key.hashCode) % numPartitions) } override def equals(other: Any): Boolean other match { case p: UserIdPartitioner p.numPartitions numPartitions case _ false } override def hashCode(): Int numPartitions }这里有三个细节值得展开。第一个细节是针对前缀A、B、C的数据我强制把它们分到固定的0、1、2号分区。这么做的目的是为了保证下游如果按照前缀做分区读取能够快速定位。而其他前缀的数据则走兜底的Hash分区避免某个分区数据量过大。第二个细节是getPartition里对null的处理。我没有让null直接抛异常而是把它分到0号分区这符合Spark对null key的默认约定。第三个细节也是最重要的分区器必须保证同一个key永远返回同一个分区号否则Spark的聚合、join都会出问题。比如你已经把一个key分到0号分区但下次再调用时因为哈希值变化或者顺序不对把它分到了1号分区那同一个key的数据就被拆散到了不同分区聚合结果就错了。写好后使用方式非常简单val partitioned rdd .map(log (log.userId, log)) .partitionBy(new UserIdPartitioner(4)) .mapPartitions(iter iter.map(_._2))partitionBy返回的RDD会带上自定义分区器后续如果再做groupByKey或joinSpark会优先复用这个分区信息从而减少shuffle。但要注意一旦对RDD执行了map操作且改变了key分区器信息可能会丢失。这一点我在后面坑位清单里会再提到。4.3 在Spark SQL中是否能用自定义分区器很多做数仓的同学会问我能不能在Spark SQL里直接用自定义分区器很遗憾Spark SQL层面目前没有开放自定义Partitioner的入口。DataFrame的repartition、repartitionByRange都只支持内置的Hash和Range规则。但这不代表完全没法用。如果你遇到必须自定义分区逻辑才能解决倾斜的问题有两种变通方案。第一种是把DataFrame转成RDD使用partitionBy之后再转回DataFrame。具体做法是用df.rdd拿到RDD手动分好区后再用toDF转回去。这样做会有一定的性能开销因为RDD和DataFrame的转换涉及内部表示的变化但在倾斜严重的时候收益远大于开销。第二种是在SQL里使用DISTRIBUTE BY或者CLUSTER BY来近似控制分区规则。比如INSERT OVERWRITE TABLE target_table SELECT * FROM source_table DISTRIBUTE BY substr(user_id, 1, 2)DISTRIBUTE BY会根据表达式的Hash值做分区虽然还是Hash但你可以通过调整表达式把高基数的列或业务上更均匀的字段作为分区键间接缓解倾斜。另外还有一个常被混淆的概念df.write.partitionBy(date)。这里的partitionBy和RDD的partitionBy完全是两回事。它控制的是写入时的目录结构或者说Hive风格的分区字段而不是RDD的分区数量。前者决定最终文件怎么归目录后者决定计算过程中数据怎么分流。我见过不少人在同一个项目里把这两个东西搞混导致调了半天文件数量没变化因为根本没调对地方。5. 实践中的坑与调优经验5.1 如何快速判断数据是否倾斜判断数据倾斜我自己有一套固定的排查流程。第一步看Spark UI里的stage详情如果同一个stage下task耗时差异极大比如绝大多数task几秒钟跑完个别task要跑几十分钟基本就是倾斜了。第二步看Shuffle Read Size如果某个task读取的量是其他task的几十倍说明某个分区数据量过大。第三步查看具体从哪个算子开始倾斜一般问题都出在groupByKey、join、distinct、partitionBy这几个环节。倾斜的直接后果不只是慢还有OOM风险。某个task数据量过大会导致executor内存压力猛增进而触发频繁GC甚至直接把executor搞挂然后整个stage重来那可不是多等几分钟的问题而是作业直接失败。所以倾斜排查越早越好。5.2 自定义分区器最容易踩的六个坑接下来把我实践里踩过、帮别人排过的坑集中列一下每一件都是真实线上事故。第一getPartition返回值越界。这个错最容易出现在你动态计算分区号的逻辑里比如某个分支使用了numPartitions作为取模基数但另一个分支返回了一个绝对写死的数字而这个数字可能大于等于numPartitions。一旦发生Spark会直接报index out of range异常。建议在getPartition的最后加一个校验或兜底确保返回值始终落在0到numPartitions-1之间。第二没有重写equals导致partitionBy后总是触发shuffle。虽然现有分区器是同一个逻辑但因为equals比较的是对象引用Spark判断两次分区器不同于是照样全量shuffle。这个问题在数据量大时特别致命因为我们使用partitionBy的初衷就是减少shuffle结果反而增加了一次shuffle。第三key在map操作后丢失分区信息。很多RDD算子比如map会改变元素结构Spark无法确认新RDD的key是否与原有分区器兼容所以会丢弃partitioner。如果你想保留自定义分区器必须继续以同样的key作为分区依据并尽量使用mapValues这类保持key不变的算子。mapValues虽然名字叫map但它只作用在value上不会破坏分区信息这个细节非常实用。第四自定义分区器没有实现Serializable。Partitioner本身继承自Serializable但你的自定义类如果引用了不可序列化的外部资源比如某些连接对象、配置文件句柄序列化时会抛NotSerializableException。解决办法是把外部依赖在执行前处理好只保留纯计算所需的信息。第五分区数过多导致小文件爆炸。自定义分区器可以自由指定numPartitions但如果你设成几千甚至上万写入结果时会生成海量小文件。小文件不仅占用NameNode内存下游读取时也会因为并发task过多反而变慢。分区数应该根据数据量和目标文件大小综合估算而不是拍脑袋定。第六采样与真实分布不一致导致Range分区倾斜。虽然这属于内置RangePartitioner的范畴但如果你自认为用了Range就万事大吉那会吃亏。真实场景中Range分区的采样比例可能不够导致边界偏移。如果遇到这种问题可以先把数据用Hash打散再排序或者增加采样参数不能盲目信赖默认行为。5.3 分区数设置的实用建议最后说分区数怎么设。这个问题没有标准答案但有一些经验公式和参考基准。对于shuffle产生的分区Spark SQL默认spark.sql.shuffle.partitions是200这个值在大多数中小数据量场景下是合理的。如果你的数据量大比如一个表有几百GB200个分区可能偏少。建议以每个分区处理128MB到256MB数据为基准来估算。假设shuffle读的数据量是1GB那用8个到16个分区是比较合适的。对于RDD API可以通过spark.default.parallelism来控制默认并行度通常建议设为集群总核数的2到3倍。比如你的executor总数有40个核那么默认并行度可以设100左右既保证充足并行又不会因为task过多造成调度开销。不过要注意这个值只是默认值具体算子执行时还会被其他因素覆盖。对于写入文件的分区数要根据目标文件大小来反推。如果一个分区写出的文件是1GB但你的HDFS块大小是128MB那下游读这个文件时会被切分成多个block读起来没问题但文件内部无法并行计算。如果一个分区写出的文件只有5MB那明显太小考虑合并分区。总体来说让输出文件的大小接近块大小的整数倍比较合适。5.4 一个完整的调优案例实录我在这里放一个虚构但非常典型的案例帮大家把前面所有概念串起来。假设有一个用户行为日志RDD大概3亿条key是用户ID。业务上需要对每个用户做近30天的行为聚合并且结果要按用户ID的首字母分区落盘方便下游按范围扫描。直接写代码如下rdd .map(log (log.userId, log)) .groupByKey() .mapValues(iter aggregate(iter)) .saveAsTextFile(/tmp/user_agg)这段代码跑了35分钟其中groupByKey阶段就有20分钟且某个task的Shuffle Read量是其他task的30倍。分析发现用户ID分布并不均匀某些前缀的用户数量特别多Hash分区后这些用户集中在少数分区里。优化方案分两步。第一步自定义一个UserIdPartitioner把用户多的前缀细分成多个分区。比如前缀A的用户最多就让它把A开头的用户ID哈希到0到9号这10个分区前缀B次之分到10到14号这5个分区其他前缀按Hash分到剩余分区。这样每个分区的数据量相对均衡。class UserPrefixPartitioner(numParts: Int) extends Partitioner { override def numPartitions: Int numParts override def getPartition(key: Any): Int { val id key.asInstanceOf[String] id.charAt(0) match { case A math.abs(id.hashCode) % 10 // 0~9 case B 10 math.abs(id.hashCode) % 5 // 10~14 case _ 15 math.abs(id.hashCode) % (numParts - 15) } } override def equals(other: Any): Boolean other match { case p: UserPrefixPartitioner p.numPartitions numPartitions case _ false } }第二步在groupByKey之前先做一次partitionBy让Spark知道数据已经按这个分区器分好区后续groupByKey就不会再产生大范围shuffle。val partitioned rdd .map(log (log.userId, log)) .partitionBy(new UserPrefixPartitioner(20)) .cache() partitioned .groupByKey() .mapValues(iter aggregate(iter)) .saveAsTextFile(/tmp/user_agg)优化后作业总耗时降到了8分钟groupByKey阶段只占3分钟。为什么效果这么明显因为同一个用户ID在一开始就被分到了同一个分区且每个分区数据量相对均匀groupByKey阶段根本不需要做大量跨分区数据合并shuffle量大幅下降。这就是自定义分区器的价值所在它不是让你炫技而是让你根据业务特征精准控制数据分布把原来需要硬扛的数据倾斜问题从根源上消掉。6. 关于分区器我最后想多说的几句写到这里我自己也把Spark分区器的知识重新梳理了一遍。其实分区器本身不难理解它无非就是一个函数输入key输出分区编号。难的是在真实的数据分布、业务需求、资源限制之间找到平衡点。根据我的个人经验多花五分钟检查一个作业的分区和分区器设置远比事后调executor内存、加节点来得划算。加资源只是兜底方案兜底完之后数据还是歪的而正确的分区器能让每个executor都吃得很均匀整个集群都稳稳地转。如果你还在困惑一个问题什么时候该用自定义分区器我的判断标准很简单当你发现默认的Hash和Range满足不了业务对局部性、均衡性、顺序性的要求时就该动手写一个。不要为了用而用但一旦用了就得记得处理好null、负数取模、equals、序列化和分区数这些细节。希望这篇内容能让你少走一些弯路。如果哪一天你的Spark作业又莫名其妙地卡在某个stage别急着骂集群先看一眼分区器说不定答案就在那里。