1. 从一次任务序列化失败倒推 RDD 的五个钩子第一次认真翻 RDD 源码是因为一个跑了两小时卡在 99% 的作业。日志里只有零星几条拉取失败的警告我第一反应是数据倾斜把并行度翻倍没用加了一轮重分区还是没用。最后把toDebugString打出来对着 Stage 图看才发现真正的原因出在一段mapPartitions里里面每次都会新建一个连接池对象而compute的调用次数是按 Task 算的不是按分区里的记录数算的。这件事让我意识到Spark 的算子 API 太顺手了顺手到我们习惯把 RDD 当成一个普通集合来用而它本质上是一份描述怎么算的元数据加一个在 Executor 上被反复调用的compute方法。RDD 源码解析这件事讲的不是背 API而是搞清楚这几个钩子分别在什么时机被谁调用。1.1 RDD 类签名里那些 transient 到底在保护什么把 RDD 的类声明摊开看最有信息量的是那几个注解。核心两个字段_scSparkContext和deps依赖序列都被标了transient同时类本身extends Serializable。这个组合不是随手写的RDD 对象需要被序列化后送到 Executor 上参与反序列化但 SparkContext 绝对不能被序列化——里面装着调度器、监听总线、各种线程池序列化它等于把整个 Driver 进程状态打包发出去。所以做法是类可序列化但这两个字段在序列化时跳过反序列化之后 Executor 上这两个字段是null代码路径也就永远走不到context上。这个细节直接解释了生产里一个高频报错在算子闭包里调用rdd.context、或者调用某个依赖SparkContext的方法Driver 端跑得好好的到 Executor 上就是空指针。同理闭包里塞进一个持有 SparkContext 引用的自定义工具类ClosureCleaner会尝试把它清理掉清理不掉就直接抛Task not serializable。我后来排查这类问题有个固定套路先看报错的闭包类里有没有非序列化的字段再看这个字段是不是顺着某条引用链能摸到 Driver 端对象。另一个值得说的是partitions_和dependencies_这两个缓存字段同样是transient的配合stateLock做双重检查加锁。也就是说RDD 的分区信息和依赖信息在 Executor 侧是懒加载的第一次访问时才通过getPartitions/getDependencies现场算出来。partitions方法里还有一行很容易被忽略的校验拿到分区数组后会用zipWithIndex逐个检查partition.index index不满足就直接抛异常。这条校验就是自定义 RDD 的第一个雷区——分区索引必须从 0 开始严格递增不能有空洞也不能乱序。1.2 getPartitions 与 compute两个真正必须实现的方法官方文档常说 RDD 有五个抽象方法但看源码会发现严格来说只有两个是抽象的getPartitions和compute。另外三个——getDependencies、getPreferredLocations、partitioner——都有默认实现属于可覆写的钩子。这个区分很重要因为它对应着两类完全不同的扩展场景前者决定数据怎么切后者决定任务怎么排。getPartitions返回Array[Partition]每个Partition只有一个index字段其hashCode就等于index。这意味着分区对象本身不携带任何数据定位信息真正的定位信息由具体子类自己扩展。比如读写分布式文件系统的那类 RDD分区里会额外包一个输入分片对象分片里才有偏移量和所在节点。这也解释了为什么partitions数组在 Driver 端就要算出来——调度器要靠它决定起多少个 Task。compute的签名是(split: Partition, context: TaskContext) Iterator[T]注意返回值是迭代器而不是集合。这个设计是整条流水线能撑住大数据量的根基上游不需要把结果全部物化下游拿到迭代器就开始逐条消费中间的内存占用是常数级的。反过来说如果你在自己的实现里先toList再返回迭代器等于把一个分区的数据全量堆进堆内存分区一大就是 OOM。我见过不少自己写个 RDD 读消息队列的代码栽在这里写法看着没问题压测一上量就崩。1.3 partitioner 与 getPreferredLocations两个决定数据放哪、任务去哪跑的钩子partitioner是个Option[Partitioner]默认None。它的作用不是怎么分区而是告诉下游我的数据已经按某个规则分好了。这个元信息会被下游的MapPartitionsRDD之类原样传递preservesPartitioning为 true 时从而避免一次额外的重分区。典型的坑是自己做了一次map之后partitioner就丢了因为MapPartitionsRDD只有在算子明确声明保持分区时才透传普通map传的是false。所以rdd.map(...).reduceByKey(...)和rdd.mapValues(...).reduceByKey(...)在 shuffle 数据量上可能差出一大截前者要多一次分区重算。getPreferredLocations(split: Partition)返回一个主机名列表调度器在分配 Task 时会优先往这些节点上放。默认返回Nil也就是随便放。文件类 RDD 会覆写它从输入分片里取数据块所在的主机列表。这个钩子的价值在于把计算推到数据旁边减少跨节点传输。如果你的数据源是某种自带本地性的存储比如按节点分片的本地缓存、某种就近接入的采集服务把这个信息透出来能省掉相当一部分网络开销——这也是自定义 RDD 最容易被漏掉的优化点很多人把getPartitions和compute写完就收工了。2. compute 返回 Iterator 这个设计决定了你的作业会不会 OOMRDD 最容易被误解的地方是把它当成数据容器。实际上它更像一份菜谱compute是怎么做这道菜的步骤真正的食材在 Executor 上现切现炒。理解了这一点很多看似奇怪的行为就通了——比如为什么cache之后第一次执行依然很慢第一次是现做做完才存起来为什么count和collect触发的计算过程几乎一样都是从头走一遍菜谱只是最后装盘方式不同。2.1 一次 count 调用展开的 compute 调用链拿最常见的rdd.map(f).filter(g).count()举例。count的实现是sc.runJob(this, Utils.getIteratorSize _).sum也就是把数迭代器里有多少个元素这个函数发到每个分区上跑最后把各分区结果求和。关键在runJob之前那些步骤DAGScheduler 拿到最终 RDD那个 filter 产生的MapPartitionsRDD后沿依赖往上找发现全是窄依赖于是只有一个 ResultStage分区数等于最底层 RDD 的分区数起 N 个 ResultTask。每个 Task 在 Executor 上执行时调的是rdd.iterator(partition, context)而不是直接调compute。iterator是个final方法它做的第一件事是判断storageLevel ! StorageLevel.NONE如果设了持久化级别就走getOrCompute否则走computeOrReadCheckpoint。computeOrReadCheckpoint再判断当前 RDD 是否已经物化过检查点是的话从检查点的第一个父 RDD 拿迭代器不是的话才调自己的compute。所以一次count在 Executor 上的实际调用栈是这样的最外层MapPartitionsRDD.compute调用f(context, split.index, firstParent.iterator(split, context))这一句里的firstParent.iterator又触发上一层MapPartitionsRDD.compute一层层递归下去直到最底层的源 RDD 返回一个真正的数据迭代器。整条链在构造迭代器阶段就全部展开完毕但一条记录都还没读——真正的读取发生在最外层getIteratorSize开始hasNext/next的时候数据像水流一样从底层一路被拉上来。这就是惰性的物理实现不是靠某种魔法而是靠迭代器的next()调用沿调用栈向上传播。理解了这个传播方向你就明白为什么在map里做副作用写数据库、发消息会出问题——副作用发生在next()被调用时而next()什么时候被调、被调多少次取决于下游算子你控制不了。take(10)可能只触发前 10 条记录的副作用count会触发全部中间还可能因为重算再来一遍。2.2 map 与 mapPartitions 的源码差异以及它带来的 GC 差异map内部就是mapPartitions的封装最终都生成MapPartitionsRDD。区别只在传进去的函数。map传的是(it: Iterator[T]) it.map(f)mapPartitions直接传你自己的Iterator[T] Iterator[U]。源码层面差别就这么点但资源开销差别不小。it.map(f)在 Scala 里返回的是一个包装迭代器每取一个元素都要多一层方法调用和一次闭包捕获。单条记录的开销可以忽略但一个分区几百万条记录时这层包装带来的对象创建和方法调用累加起来很可观。更关键的是mapPartitions让你有机会做分区级的初始化建连接、开文件句柄、预加载维表。map做不到这件事因为它的函数是按记录调用的。我自己的经验法则是纯字段变换用map可读性好涉及外部资源、批量提交、需要按分区做聚合初始化的一律用mapPartitions。曾经有个写外部存储的作业用map每条记录开一次连接QPS 上不去还把连接池打满了换成mapPartitions在函数开头建一次、结尾关一次吞吐直接翻了若干倍GC 日志里的短命对象也少了一大截。这个改动没有任何高级技巧纯粹是理解了compute的调用粒度是按分区而不是按记录。2.3 自己实现 RDD 时的迭代器陷阱真正动手写自定义 RDD 的时候有几个坑几乎人人都会踩一遍。第一个是迭代器只能消费一次compute返回的Iterator被下游消费完就空了如果同一份迭代器被两处引用比如你在compute里存了个成员变量指向它第二处拿到的就是空迭代器。第二个是compute可能被调多次没有缓存的情况下一个 RDD 被两个 Action 用到compute就会跑两遍。这不是 bug是设计——所以任何在compute里做副作用的写法都是危险的。第三个坑更隐蔽跟hasNext的语义有关。自己实现迭代器时如果hasNext内部做了预读下一条的逻辑但没有正确缓存会出现hasNext返回 true 但next抛异常的情况反过来如果hasNext有副作用比如消费了缓冲区连续调用两次hasNext就会丢数据。Spark 内部大量代码会连续调用hasNext所以这个约束必须严格遵守。我后来写这类迭代器一律用懒初始化 状态标志的模板宁可多几行代码也不省那点逻辑。第四个是资源释放。compute里开的文件句柄或者网络连接正常路径下消费完自然关闭但如果 Task 被 kill 或者中途抛异常就需要靠TaskContext.addTaskCompletionListener注册回调来兜底。这个 API 在源码里到处都在用它注册的回调在 Task 无论成功失败结束时都会执行是自定义 RDD 里唯一可靠的清理入口。3. 宽依赖在哪里断开DAGScheduler 切 stage 的真实判定过程很多人能背出窄依赖不需要 shuffle宽依赖需要 shuffle但问到具体是哪一行代码做了这个判断就答不上来了。搞清楚这个过程的价值在于你能预判一个作业会被切成几个 Stage、每个 Stage 起多少个 Task、shuffle 文件会写多少份。这些数字直接决定作业的稳定性和资源占用而不是等跑起来看监控才发现问题。3.1 NarrowDependency 与 ShuffleDependency 的实现差异依赖体系的根是Dependency[T]只有两个直接子类NarrowDependency和ShuffleDependency。NarrowDependency是抽象类核心是一个方法getParents(partitionId: Int): Seq[Int]——给定下游分区号返回它依赖哪些上游分区号。最常用的OneToOneDependency直接返回List(partitionId)也就是一一对应。另一类是范围依赖一个下游分区依赖上游一段连续分区典型用在合并多个 RDD 的场景。ShuffleDependency的构造参数信息量大得多分区器、序列化器、可选的排序规则、可选的聚合器、是否在 map 端做合并还有一个在构造时自动分配的shuffleId来自sparkContext.newShuffleId()。这几个参数其实就定义了 shuffle 的全部行为分区器决定数据往哪个 reduce 分区走聚合器和 map 端合并开关决定 map 端要不要先做一次局部聚合这就是为什么reduceByKey比groupByKey省内存——前者构造依赖时传了聚合器和mapSideCombine true后者没有。判定逻辑本身非常简单DAGScheduler 从最终 RDD 出发做一次广度优先遍历遇到ShuffleDependency就把它记下来并停止继续往那个方向走遇到窄依赖就继续往上。遍历结束后得到的就是当前 Stage 的所有 shuffle 边界。这个方法的返回值是一个集合所以同一个 shuffle 依赖在多条路径上被引用也不会重复创建 Stage。3.2 从 runJob 到 submitMissingTasks 的完整调用链把链路串起来看SparkContext.runJob先做两件事——取调用点信息用于 UI 上显示是哪个代码行触发的和对传入的函数做闭包清理然后交给DAGScheduler.runJob。DAGScheduler 不是直接同步执行而是往自己的事件循环里投一个JobSubmitted事件由事件处理线程异步处理。这点值得强调Driver 端的调度是事件驱动的所有 Stage 的创建、提交、失败重试都在同一个线程里串行处理所以 UI 上看到的调度延迟有时跟业务代码无关是事件队列排队的表现。事件被处理后先创建 ResultStage包含最终 RDD 和要在每个分区上执行的函数然后递归地往上找父 Stage。父 Stage 的创建过程中如果是第一次遇到这个 shuffle还要额外做一件事把产生该 shuffle 的 RDD 数据先算出来——这就是submitMissingTasks里缺分区概念的来源。已经算完的分区不会重算只提交还没有结果的分区这是 Stage 重试能续跑而不是全量重来的基础。submitMissingTasks里做了几件关键的事确定要提交哪些分区、把 RDD 和函数或 RDD 和 shuffle 依赖序列化成taskBinary、按数据本地性给每个 Task 算一个偏好位置、把 Task 集合交给TaskSchedulerImpl。序列化这一步用的是 closureSerializer也就是 Java 序列化这意味着 taskBinary 的体积跟你的闭包大小正相关。闭包里有大对象时光是序列化和网络传输就能拖垮调度——这是广播变量存在的真实原因不是为了语法好看。3.3 ShuffleMapTask 与 ResultTask 在 runTask 里的分歧点Task 只有两种分歧点在runTask的第一行反序列化。ShuffleMapTask反序列化出的是(RDD, ShuffleDependency)二元组然后向 shuffle 管理器要一个 writer把rdd.iterator(partition, context)的结果写进去最后返回一个MapStatus记录每个 reduce 分区数据的大小和位置。ResultTask反序列化出的是(RDD, 函数)直接把迭代器喂给函数返回值就是要交回 Driver 的结果。这里有个细节很能说明设计意图两者调用的都是rdd.iterator(partition, context)而不是compute。也就是说缓存和检查点的逻辑对两类 Task 是一视同仁的都从同一个入口进。另一个细节是ShuffleMapTask返回的MapStatus里存的是每个 reduce 分区的数据大小而不是记录数——因为下游拉取时需要的正是从哪个位置读多少字节这个信息直接对应到 shuffle 数据文件和索引文件。写 shuffle 文件的过程也值得一提。排序型 shuffle 管理器会为每个 map 分区写一个数据文件和一个索引文件数据文件里按 reduce 分区连续存放索引文件记录每个 reduce 分区的偏移和长度。所以shuffle 文件数约等于 map 分区数乘以 2而不是 map 分区数乘以 reduce 分区数。这个换算关系在排查小文件过多的问题时非常有用一个 2000 分区的 map 阶段光 shuffle 文件就是 4000 个再乘以重试次数和 Stage 数量很容易把存储系统的元数据服务压垮。3.4 AQE 对 RDD 作业无效这件事很多文档没说清自适应查询执行在数据框和 SQL 层面做了很多优化动态合并分区、动态切换连接策略、动态处理倾斜。但如果你用的是 RDD API这些优化一个都不会生效。原因就在上面这条链路上RDD 的 Stage 切分是静态的在作业提交时就把依赖图定死了shuffle 分区数在构造ShuffleDependency时由分区器确定运行时不会再改。这个结论影响很大。我见过不少团队为了用上自适应优化把 RDD 代码硬包一层转成数据框结果因为 schema 推导和序列化开销性能反而下降。判断标准其实很简单如果你的作业主体是 RDD 算子就老老实实按静态切分的思路去调分区数和并行度别指望运行时帮你兜底如果你能做到底层用数据框表达那才谈得上自适应优化。两条路线不要混着用混用只会让问题定位变得极其困难。4. cache 与 checkpoint 在源码里是两条完全不同的路径这两个 API 在文档里经常被放在一起讲实际源码里它们几乎不共享代码。缓存走的是 BlockManager 的存取路径检查点走的是重算 落盘 改依赖的路径。搞混这两个机制就会写出既 cache 又 checkpoint 但都不生效的代码。4.1 iterator 入口处的第一道分叉getOrComputeiterator里的判断只有一行storageLevel ! StorageLevel.NONE就走getOrCompute。getOrCompute的逻辑是用RDDBlockId(rddId, partitionIndex)构造块 ID去 BlockManager 问有没有这个块有就直接返回没有就调computeOrReadCheckpoint算出迭代器然后把迭代器交给 BlockManager 存起来。这里的关键是存这个动作的实现方式。BlockManager 拿到的不是一个已经物化的集合而是一个迭代器它会根据存储级别决定策略纯内存级别会尝试把迭代器展开成内存中的块展开过程中申请不到足够内存就退化成边展开边让调用方消费也就是不缓存内存加磁盘级别在内存不够时会把数据写到磁盘纯磁盘级别直接写盘。所以cache是否真的生效取决于运行时内存是否够用而不是取决于你有没有调用它。还有一个容易忽略的点persist在被调用时会把 RDD 注册到 SparkContext 的持久化 RDD 表里同时注册给清理器用于后续的引用计数回收。这意味着persist之后忘了unpersist这块内存会一直占着直到整个上下文结束或者被清理器判定可以回收。长驻的流式作业里这个问题尤其明显——每批次都 persist 一个新的 RDD 却不释放内存曲线会一路向上。4.2 BlockManager 落盘与 unroll 内存阈值的相互作用展开迭代器的过程有个专门的阈值参数默认是 1MB。含义是先按这个量级逐步申请内存边申请边展开如果展开到一半发现还需要更多继续申请如果总计需要的量超出可用内存就放弃内存缓存。这个渐进式展开的设计是为了避免申请一大块内存结果用不完的浪费但也带来一个现象小分区会被缓存大分区可能一次都缓存不上。我遇到过一个很典型的案例某作业的每个分区数据量差异很大大部分分区几百 KB 都缓存住了少部分几 GB 的分区全部没缓存上导致下游每次都要重算那几个大分区。查监控发现缓存命中率一直是 90% 多看着很健康实际性能被那几个大分区拖死。解决办法不是调大阈值而是先用repartition把数据打散让分区大小均匀再缓存。这个顺序很重要——缓存均匀的小分区才有意义缓存不均匀的大分区基本是白费内存。另外要注意存储级别里堆外和副本这两个维度。副本数大于 1 时缓存会占用成倍的内存或磁盘并且写入要走网络堆外内存则绕开 JVM 堆减少 GC 压力但对内存管理器的配置更敏感。选级别的时候我一般按能不能承受重算来决定重算代价极高的用内存加磁盘双保险重算代价一般的就纯内存数据量特别大且重算便宜的干脆不缓存靠上游的分区裁剪来省成本。4.3 checkpoint 的递归触发时机为什么它总在作业结束之后调用checkpoint()的时候源码里做的事情少得让人意外只是给当前 RDD 打上一个标记创建一个检查点数据对象挂上去其他什么都没干。真正的落盘发生在作业执行过程中——在SparkContext.runJob的最后一行会调用最终 RDD 的doCheckpoint。doCheckpoint的逻辑是递归的如果当前 RDD 自己带了检查点标记就先递归让所有父 RDD 处理然后执行自己的检查点动作如果自己没标记就直接递归父 RDD。执行检查点动作时会把当前 RDD 算一遍把结果写到一个可靠存储上生成一个新的检查点 RDD然后把原 RDD 的依赖清空、分区信息重置、指向新的检查点 RDD。这套机制解释了两个常见困惑。第一为什么checkpoint之后第一次执行没有变快——因为这次执行正是用来生成检查点数据的快的是下一次。第二为什么必须配合缓存使用——如果不缓存生成检查点的这次计算会把整条血缘重算一遍长血缘的作业会付出巨大代价加上缓存之后重算变成读缓存成本就降下来了。我在实际项目里见过血缘长达几十个算子的作业只加了检查点没加缓存结果第一次执行时间翻了三倍就是因为整条链被重算了一遍。还有一个细节依赖被清空之后toDebugString输出里那一支会显示成检查点标记血缘就断了。有人以为这是血缘丢失的 bug其实是设计如此——断掉血缘正是检查点的目的让后续的容错恢复从这个点开始而不是从最初的数据源重算。4.4 cache 明明调了却不生效的几种典型情况第一种是缓存级别设成了仅内存但数据量超出存储内存配额展开失败缓存被静默放弃。判断方法是看 BlockManager 的日志里有没有展开失败的记录或者直接看存储页面上这个 RDD 的缓存比例。第二种是先调了persist但 RDD 在被真正计算之前就被unpersist了或者作业在某次失败重试时把 RDD 对象换掉了某些算子在特定配置下会重建 RDD 对象。第三种最隐蔽缓存的 RDD 被后续算子改写了分区器导致下游读的其实是另一个 RDD。比如缓存了rdd然后执行rdd.reduceByKey(...)下游读的是ShuffledRDD而ShuffledRDD并没有被缓存所以该 shuffle 还是每次都重算。正确的做法是把reduceByKey之后的结果也缓存或者缓存之后的所有关键中间结果而不是只缓存源头。第四种跟闭包有关缓存的是 RDD 数据但如果你的compute里每次都重新拉一遍外部数据比如每次都查一次维表接口缓存住的是最终结果这部分外部调用只省了第一次之后的重算——听上去没问题但如果这个 RDD 因为内存不足被淘汰下次重算又会去拉外部数据。这种情况下正确的做法是把维表做成广播变量常驻内存而不是靠 RDD 缓存来兜。5. 摸清调用链之后我在实际项目里改掉的几个习惯读源码最大的回报不是能背出类名而是有了预判的能力写完一段代码能大致说出它会生成几个 Stage、起多少 Task、写多少文件、占多少内存。以下几点是我在实际项目里改得最彻底的。5.1 分区数量、并行度与 shuffle 文件数之间的关系以前调并行度基本靠拍脑袋看作业慢就加加了没用就再加。搞清链路之后有了明确的推导路径分区数决定 Task 数Task 数决定并发度而 shuffle 阶段的并发度还受 reduce 分区数约束。假设一个 map 阶段 2000 分区、reduce 阶段 200 分区那任何一个瞬间最多只有 200 个 reduce Task 在跑多出来的资源都是闲着的。这时候要调的是 reduce 分区数而不是总并行度。另一个换算关系是文件数。map 分区 2000每个 map Task 写一个数据文件加一个索引文件那就是 4000 个文件。如果这个 Stage 因为失败重试跑了三遍每遍都写新文件、成功后删旧的但删除可能延迟瞬时文件数能到上万。存储系统的元数据服务对文件数很敏感超过某个量级就会出现明显的创建和列目录延迟进而拖慢整个 Stage。我现在的习惯是分区数超过某个量级时先评估一下这个 Stage 到底需不需要这么多分区多数情况下把上游做一次合并比硬扛更划算。5.2 闭包瘦身把 Driver 端的大对象挡在序列化之外闭包被序列化进 taskBinary体积直接影响调度开销。我见过把整个配置对象、整个维表、甚至整个 Spring 容器塞进闭包的代码单个 Task 序列化出来几 MB几千个 Task 光是分发就要好几分钟。这个问题在源码层面的根因很清楚clean会尝试剔除闭包里用不到的字段但只要有一个字段被函数体引用到整个引用链都会被带上。改法有两条。第一是能广播的一律广播广播变量走的是独立的传输通道每个 Executor 只存一份Task 里带的是引用而不是数据本身。第二是给闭包里的类加transient或者把用不到的字段抽出去别让一个大对象因为一个字段的引用被整体拖进去。这两条做下来我负责的一个作业的 taskBinary 平均体积从几 MB 降到几十 KB调度延迟肉眼可见地下降。顺带说一个判断技巧如果作业在提交阶段就卡住UI 上 Task 迟迟不出现或者 Driver 日志里出现序列化超时优先怀疑闭包体积而不是集群资源。这个判断方向搞反了会浪费很多时间在扩容上。5.3 用 toDebugString 与 Stage 详情页反查自己的假设toDebugString打印的是血缘树带缩进和依赖类型标记能直接看出哪一步产生了 shuffle。我现在的习惯是任何作业上线前先把它的输出贴到笔记里作为一份预期结构。上线后如果 Stage 数量、Task 数量和预期对不上说明代码路径或者配置跟我以为的不一样立刻去查。Stage 详情页里我重点看三个数字每个 Task 的输入大小分布、shuffle 写入大小分布、GC 时间占比。输入大小分布能立刻暴露倾斜个别 Task 是别人的几十倍shuffle 写入分布能暴露分区器选错GC 占比高则指向内存问题而不是计算问题。这三个指标配合血缘树看绝大多数性能问题能在第一次排查就定位到具体算子而不是靠反复改参数试。6. 一份可以直接照抄的自定义 RDD 骨架理论讲完落到能跑的代码上。下面这个骨架解决的是从一批分片文件里读数据并解析成对象这种最常见的自定义场景把前面提到的分区索引校验、迭代器生命周期、资源释放都处理掉了。6.1 从分片文件源构造 RDD 的完整代码import org.apache.spark._ import org.apache.spark.rdd.RDD import java.io.{BufferedReader, File, FileReader} // 自定义分区除了 index还要带上下手要读的文件 private[example] class FileSplitPartition(val index: Int, val file: File) extends Partition class TextSegmentsRDD( sc: SparkContext, inputDir: String, minPartitions: Int ) extends RDD[String](sc, Nil) { // 没有父 RDD依赖传 Nil // 1. 分区信息Driver 端算好索引必须从 0 严格递增 override protected def getPartitions: Array[Partition] { val files new File(inputDir).listFiles() .filter(f f.isFile !f.getName.startsWith(.)) .sortBy(_.getName) require(files.nonEmpty, s目录下没有可用文件: $inputDir) // 文件数少于期望分区数时按文件数来不要硬凑 val target math.min(files.length, math.max(minPartitions, 1)) // 每个分区按文件编号取模分配保证每个分区都有活干 val grouped Array.fill(target)(new scala.collection.mutable.ArrayBuffer[File]()) files.zipWithIndex.foreach { case (f, i) grouped(i % target) f } grouped.zipWithIndex.map { case (fs, idx) new MultiFilePartition(idx, fs.toArray) }.asInstanceOf[Array[Partition]] } // 2. 计算逻辑返回迭代器绝不物化 override def compute(split: Partition, context: TaskContext): Iterator[String] { val part split.asInstanceOf[MultiFilePartition] val readers scala.collection.mutable.ArrayBuffer[BufferedReader]() // 注册清理回调Task 成功或失败都会执行 context.addTaskCompletionListener[Unit] { _ readers.foreach(r try r.close() catch { case _: Throwable }) readers.clear() } // 惰性打开文件逐行吐出 part.files.iterator.flatMap { file val reader new BufferedReader(new FileReader(file)) readers reader Iterator.continually(reader.readLine()).takeWhile(_ ! null) } } // 3. 本地性告诉调度器这个分区最好在哪些节点上跑 override protected def getPreferredLocations(split: Partition): Seq[String] split.asInstanceOf[MultiFilePartition].hosts.toSeq } private[example] class MultiFilePartition( val index: Int, val files: Array[File], val hosts: Array[String] Array.empty ) extends Partition这份代码里有几个决定成败的细节。getPartitions里做了require校验和文件数少于期望分区数时按文件数来的处理避免生成一批空分区——空分区不会报错但会白白占用 Task 槽位还会让分区大小的统计失真。compute里用flatMap把多个文件串成一个迭代器全程没有toList内存占用与文件大小无关。最关键是清理回调。用Iterator.continually加takeWhile的写法在正常消费完全部内容时文件句柄不会自动关闭——因为迭代器不知道流结束了意味着该关文件。只有靠addTaskCompletionListener才能保证无论正常结束还是被 kill 都能释放。我第一次写这类 RDD 时没加这个回调跑了几百个 Task 之后 Executor 上报打开文件数过多排查了半天才发现是句柄泄漏。6.2 上线前的自查清单写完自定义 RDD我一般会过一遍下面这张表。这些问题不一定会立刻暴露但压测或者长跑之后大概率冒出来。检查项为什么重要典型故障表现分区索引是否从 0 连续递增partitions方法内部有强校验提交作业即抛异常提示分区索引不匹配compute是否返回惰性迭代器全量物化会让单分区数据全进堆内存大分区时 Executor 内存溢出迭代器是否只被消费一次下游可能拿到已耗尽的迭代器结果是空的但没有任何报错hasNext是否无副作用内部代码会连续多次调用数据间断性丢失极难复现外部资源是否有清理回调Task 被 kill 时正常路径不执行文件句柄或连接数持续增长分区是否可能为空空分区占槽位不出活并行度看着够实际有效并发很低是否覆写了本地性方法默认Nil等于放弃数据本地性网络传输量异常偏高是否声明了分区器不声明会让下游多一次重分区多出一次不必要的 shuffle这张表里我踩得最深的是迭代器只被消费一次那条。当时写了个 RDD在compute里把迭代器赋给一个成员变量方便调试打印结果线上偶发空结果本地怎么都复现不了——因为本地数据量小调试打印和正式消费的时序恰好错开了。这类问题的教训是compute里除了构造并返回迭代器其他什么都别做任何顺手存一下的念头都要掐掉。还有一条经验值得单独说自定义 RDD 写完先用极小数据集跑通再用单分区超大文件和多分区小文件两种极端分布各跑一遍。这两种分布能分别暴露内存问题和分区不均问题比等生产环境出问题再回头查要省事得多。我自己现在的固定流程是写完 → 三种数据分布各跑一遍 → 看 Stage 页面的分区大小分布是否均匀 → 再看 GC 占比 → 通过之后才敢往生产推。最后分享一个调试时很管用的小技巧在compute开头打一行日志把分区索引和 Task 的尝试次数一起打出来。这样从 Executor 日志里能直接看出哪些分区在被反复重算、哪些分区一次都没跑过。这个方法比看 UI 上的聚合指标更直接尤其是在排查某些分区反复失败这类问题时一眼就能定位到具体是哪个分区、第几次尝试。