如何在Spark中对预排序DataFrame执行二分查找

Spark 没有提供开箱即用的DataFrame/Dataset级别的二分查找API，但只要你是提前对目标列做了全局排序（注意不是sortWithinPartitions的分区内排序），完全可以实现O(logN)复杂度的高效查找，比默认全表过滤的O(N)扫描性能高几个数量级，常用实现分两种场景：

分布式大数据场景：基于分区边界的二分剪枝法（无需拉全量数据到Driver） #

全局排序后的DataFrame本身满足两个有序特性：

• 每个分区内部的数据按目标列升序排列
• 后一个分区的所有数据值≥前一个分区的所有数据值
  利用这个特性可以只扫描单个分区就完成查找，步骤如下：
• 第一步：获取每个分区的首元素值，只需要遍历每个分区取第一条数据即可，开销极低，不会触发全表扫描
• 第二步：在Driver本地对分区首元素数组做二分查找，定位到目标值所在的分区ID
• 第三步：仅扫描定位到的单个分区，遍历找到第一个≥目标值的元素即可

参考Scala实现代码：

import org.apache.spark.sql.functions.col
// 提前对目标列做全局排序并缓存，避免后续重复计算
val sortedDf = sourceDf.orderBy(col("target_col")).cache()
// 收集各分区首元素
val partitionFirstValues = sortedDf.rdd.mapPartitions(iter => {
  if (iter.hasNext) Iterator(iter.next().getAs[Long]("target_col")) else Iterator.empty
}).collect()

val targetValue = 100L // 待查找的目标阈值
import scala.collection.Searching._
// 二分定位目标分区
val targetPartition = partitionFirstValues.search(targetValue) match {
  case Found(exactPartitionIdx) => exactPartitionIdx
  case InsertionPoint(insertIdx) => if (insertIdx == 0) 0 else insertIdx - 1
}
// 仅扫描目标分区取第一个符合条件的结果
val findResult = sortedDf.rdd.mapPartitionsWithIndex((partIdx, rowIter) => {
  if (partIdx != targetPartition) Iterator.empty
  else rowIter.filter(_.getAs[Long]("target_col") >= targetValue).take(1)
}).first()

注意：全局排序后不要执行会触发shuffle的操作（比如join、groupBy、重分区），否则会破坏分区有序性，导致查找逻辑失效。

小数据量场景：本地数组二分法 #

如果数据量不大，排序列可以直接放进Driver内存，最简便的方式是把排序后的目标列直接收集为本地数组，用语言内置的二分工具直接查找即可。
参考Python实现：

import bisect
# 收集已排序的目标列
sorted_col = [row["target_col"] for row in sorted_df.select("target_col").collect()]
target = 100
# bisect_left直接返回首个≥目标值的索引
insert_pos = bisect.bisect_left(sorted_col, target)
first_ge_value = sorted_col[insert_pos] if insert_pos < len(sorted_col) else None

常见避坑 #

不要直接写df.filter(col("target_col") >= target).first()实现需求，哪怕DataFrame已经排好序，Spark的优化器也不会自动做二分剪枝，这个语句会触发全表扫描，数据量大时性能极差。

内容的提问来源于stack exchange，提问作者user626528