Spark RDD中为每条记录匹配同Key最近经纬度记录的优化方案

问题背景 #

现有两个超大Spark RDD：rddData和rddMeta，需为rddData每条记录匹配同key下经纬度物理距离最近的rddMeta记录。但同key的记录数可达数亿级，全量Join完全不可行，必须将距离计算限定在局部空间范围内。当前采用3x3经纬度取整块复制rddMeta的方案虽能运行，但存在数据膨胀严重、候选集冗余、内存压力大、部署繁琐等问题，需更高效简洁的实现。

数据结构定义：

case class Meta(key: Any, lat: Double, lon: Double, meta: Any)

case class Data(key: Any, lat: Double, lon: Double, data: Any, meta: Option[Any] = None) { 
  def ingestMeta(m: Option[Meta]) = this.copy(meta = m) 
}

val rddMeta: RDD[Meta] = ???
val rddData: RDD[Data] = ???

当前方案的核心痛点 #

1. 数据膨胀严重：每条Meta记录被复制9份，直接导致rddMeta数据量放大9倍，增加IO和存储开销
2. 候选集冗余：1度网格对应的距离（纬度≈111公里）远大于需求的60公里，导致Join后候选Meta过多，距离计算冗余度高
3. 内存风险高：groupByKey会将同网格下的所有Meta加载到内存Iterable中，极易引发OOM

优化方案建议 #

1. 精准空间分块，减少数据复制 #

根据60公里的需求，计算更精准的网格大小，仅将Meta映射到自身网格及可能覆盖60公里范围的相邻网格，避免无意义的复制：

• 纬度方向：1度≈111公里，60公里对应≈0.54度，可采用0.5度为步长划分网格
• 经度方向：因经度距离随纬度变化，可按lon.toInt * 0.6（赤道附近≈66公里）或动态计算步长，确保网格覆盖范围不超过60公里
• 仅映射自身网格+上下左右4个相邻网格，复制份数从9份降至最多5份

修改后的Meta分块方法：

case class Meta(key: Any, lat: Double, lon: Double, meta: Any) {
  // 按0.5度步长划分纬度网格，0.6度步长划分经度网格
  private def getGrid(lat: Double, lon: Double): (Int, Int) = {
    val latGrid = (lat / 0.5).toInt
    val lonGrid = (lon / 0.6).toInt
    (latGrid, lonGrid)
  }

  def preciseBox: List[((Any, Int, Int), Meta)] = {
    val (latGrid, lonGrid) = getGrid(lat, lon)
    // 仅保留自身网格+上下左右相邻网格，确保覆盖60公里范围
    val adjacentGrids = List(
      (latGrid, lonGrid),
      (latGrid - 1, lonGrid),
      (latGrid + 1, lonGrid),
      (latGrid, lonGrid - 1),
      (latGrid, lonGrid + 1)
    )
    adjacentGrids.map(grid => ((key, grid._1, grid._2), this))
  }

  // Haversine距离计算实现
  def haversine(otherLat: Double, otherLon: Double): Double = {
    val R = 6371 // 地球半径（公里）
    val dLat = math.toRadians(otherLat - lat)
    val dLon = math.toRadians(otherLon - lon)
    val a = math.sin(dLat/2) * math.sin(dLat/2) +
            math.cos(math.toRadians(lat)) * math.cos(math.toRadians(otherLat)) *
            math.sin(dLon/2) * math.sin(dLon/2)
    val c = 2 * math.atan2(math.sqrt(a), math.sqrt(1-a))
    R * c
  }
}

2. 替换groupByKey为aggregateByKey，降低内存压力 #

避免直接用groupByKey加载全量Meta到内存，改用aggregateByKey在分区内提前聚合为数组，内存更可控：

import org.apache.spark.rdd.RDD
import scala.collection.mutable.ArrayBuffer

// 用aggregateByKey构建每个网格下的Meta数组（替代Iterable）
val rddMetaIndexed: RDD[((Any, Int, Int), Array[Meta])] =
  rddMeta
    .flatMap(_.preciseBox)
    .aggregateByKey(ArrayBuffer[Meta]())(
      (buf, meta) => buf += meta,
      (buf1, buf2) => buf1 ++= buf2
    )
    .mapValues(_.toArray)

3. 切换到Dataset API，利用Spark原生空间优化 #

Spark Dataset/SQL提供原生空间处理函数，结合分区策略可大幅简化实现并提升性能：

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._
import org.apache.spark.sql.expressions.Window

val spark = SparkSession.builder().getOrCreate()
import spark.implicits._

// 转换为Dataset并添加空间点列，按key+网格分区
val metaDs = rddMeta.toDF()
  .withColumn("point", expr("ST_Point(lon, lat)"))
  .repartition(col("key"), expr("ST_GridCode(point, 0.5, 0.6)"))

val dataDs = rddData.toDF()
  .withColumn("point", expr("ST_Point(lon, lat)"))
  .repartition(col("key"), expr("ST_GridCode(point, 0.5, 0.6)"))

// 按key分组，筛选60公里内的最近邻Meta
val windowSpec = Window.partitionBy("key", "lat", "lon").orderBy("distance")

val resultDs = dataDs.join(metaDs, Seq("key"), "left_outer")
  .withColumn("distance", expr("ST_Distance(data.point, meta.point)"))
  .filter(col("distance") <= 60 || col("distance").isNull)
  .withColumn("rank", row_number().over(windowSpec))
  .filter(col("rank") === 1)
  .select("data.*", "meta.meta")
  .as[Data]

4. 局部KD-Tree加速最近邻查询 #

为每个网格下的Meta集合构建KD-Tree，将线性遍历O(n)的查询复杂度降至O(log n)：

// 自定义简易2D KD-Tree实现
class KDTree(points: Array[(Double, Double, Meta)]) {
  private class Node(val point: (Double, Double, Meta), val axis: Int, val left: Option[Node], val right: Option[Node])

  private val root: Option[Node] = buildTree(points, 0)

  private def buildTree(points: Array[(Double, Double, Meta)], axis: Int): Option[Node] = {
    if (points.isEmpty) None
    else {
      val sorted = points.sortBy(_._1 + _._2) // 简化排序逻辑，可按当前轴排序优化
      val mid = sorted.length / 2
      val nextAxis = (axis + 1) % 2
      Some(Node(
        sorted(mid),
        axis,
        buildTree(sorted.take(mid), nextAxis),
        buildTree(sorted.drop(mid + 1), nextAxis)
      ))
    }
  }

  def nearest(targetLat: Double, targetLon: Double): Option[Meta] = {
    def search(node: Node, best: (Double, Meta)): (Double, Meta) = {
      val dist = math.hypot(node.point._1 - targetLat, node.point._2 - targetLon)
      val newBest = if (dist < best._1) (dist, node.point._3) else best
      val axis = node.axis
      val target = if (axis == 0) targetLat else targetLon
      val nodeVal = if (axis == 0) node.point._1 else node.point._2

      val (first, second) = if (target < nodeVal) (node.left, node.right) else (node.right, node.left)
      val updatedBest = first.map(search(_, newBest)).getOrElse(newBest)
      val planeDist = math.abs(target - nodeVal)
      if (planeDist < updatedBest._1) second.map(search(_, updatedBest)).getOrElse(updatedBest) else updatedBest
    }

    root.map(node => search(node, (Double.MaxValue, node.point._3))).map(_._2)
  }
}

// 构建KD-Tree索引
val rddMetaKDTree: RDD[((Any, Int, Int), KDTree)] =
  rddMetaIndexed.mapValues(arr => {
    val points = arr.map(m => (m.lat, m.lon, m))
    new KDTree(points)
  })

// 执行匹配查询
val result: RDD[Data] = 
  rddData
    .keyBy(d => {
      val (latGrid, lonGrid) = Meta("", d.lat, d.lon, "").getGrid(d.lat, d.lon)
      (d.key, latGrid, lonGrid)
    })
    .leftOuterJoin(rddMetaKDTree)
    .map { case ((key, lt, ln), (data, kdTreeOpt)) =>
      val closest = kdTreeOpt.flatMap(_.nearest(data.lat, data.lon))
      data.ingestMeta(closest.map(_.meta))
    }

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