Spark Dataset按距离条件分组：合并5公里内城市人口

1. 把loc列转成整数元组 #

假设你的loc列是类似"(123,456)"的字符串格式，两种方法搞定：

方法1：用Spark内置函数拆分（推荐） #

直接拆成纬度、经度两列，后续计算距离更方便，要元组的话再组合回去：

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

// 先清理括号、拆分坐标转整数
val dfWithCoords = originalDF
  .withColumn("loc_clean", regexp_replace(col("loc"), "[()]", ""))
  .withColumn("lat", split(col("loc_clean"), ",")(0).cast("int"))
  .withColumn("lon", split(col("loc_clean"), ",")(1).cast("int"))
  .drop("loc_clean")

// 如果必须要整数元组格式，用struct组合
val dfWithTuple = dfWithCoords.withColumn("loc_tuple", struct(col("lat"), col("lon")))

方法2：自定义UDF（适合复杂格式） #

如果loc格式不规整，写个UDF解析更灵活：

import org.apache.spark.sql.functions.udf

val parseLocToTuple = udf((locStr: String) => {
  val coords = locStr.replaceAll("[()]", "").split(",")
  (coords(0).toInt, coords(1).toInt)
})

val dfWithTuple = originalDF.withColumn("loc_tuple", parseLocToTuple(col("loc")))

2. 5公里内城市分组汇总人口 #

直接用groupBy搞不定这个需求，因为得先算两两城市的距离，再把近的归为一组，用「自连接+Haversine距离+连通分量标记」的思路实现：

步骤1：写个Haversine距离计算函数 #

用来算地球上两点的球面距离，输出单位是公里：

import org.apache.spark.sql.functions.udf

val haversineDistance = udf((lat1: Int, lon1: Int, lat2: Int, lon2: Int) => {
  val earthRadius = 6371 // 地球半径（公里）
  val dLat = Math.toRadians(lat2 - lat1)
  val dLon = Math.toRadians(lon2 - lon1)
  val a = Math.sin(dLat/2) * Math.sin(dLat/2) +
          Math.cos(Math.toRadians(lat1)) * Math.cos(Math.toRadians(lat2)) *
          Math.sin(dLon/2) * Math.sin(dLon/2)
  val c = 2 * Math.atan2(Math.sqrt(a), Math.sqrt(1-a))
  earthRadius * c
})

步骤2：自连接筛选5公里内的城市对 #

用zipcode当唯一标识，自连接时避免重复配对：

val dfWithId = dfWithCoords.withColumn("city_id", col("zipcode"))

// 只连接ID更小的城市，避免(a,b)和(b,a)重复计算
val joinedDF = dfWithId.as("a")
  .join(dfWithId.as("b"), col("a.city_id") < col("b.city_id"))
  .withColumn("distance", haversineDistance(col("a.lat"), col("a.lon"), col("b.lat"), col("b.lon")))
  .filter(col("distance") <= 5.0)

步骤3：标记连通分量（把近的城市归为一组） #

用Spark GraphX计算连通分量，把城市当节点，近的城市对当边：

import org.apache.spark.graphx._
import org.apache.spark.rdd.RDD

// 构建节点RDD：(city_id, (城市名, 人口))
val vertices: RDD[(Long, (String, Int))] = dfWithId
  .select(col("city_id").cast("long"), col("city"), col("pop"))
  .rdd.map(row => (row.getLong(0), (row.getString(1), row.getInt(2))))

// 构建边RDD：(起点ID, 终点ID, 距离)
val edges: RDD[Edge[Double]] = joinedDF
  .select(col("a.city_id").cast("long"), col("b.city_id").cast("long"), col("distance"))
  .rdd.map(row => Edge(row.getLong(0), row.getLong(1), row.getDouble(2)))

// 计算连通分量，得到每个城市所属的组ID
val graph = Graph(vertices, edges)
val connectedComponents = graph.connectedComponents().vertices

// 关联回原数据集
val groupedDF = dfWithId
  .join(spark.createDataFrame(connectedComponents).toDF("city_id", "group_id"), "city_id")

步骤4：按组汇总人口 #

最后一步简单聚合：

val resultDF = groupedDF
  .groupBy("group_id")
  .agg(
    collect_list("city").alias("同组城市"),
    sum("pop").alias("组内总人口")
  )

额外提示 #

如果数据量很大，自连接会生成超多中间数据，建议先按经纬度分网格，只在同一或相邻网格内计算距离，能大幅提升性能。

内容的提问来源于stack exchange，提问作者Guillermo García