PySpark批量地址地理编码性能优化求助：75k行数据处理速度过慢

我现在正在写一个自动获取经纬度的函数，但我的数据文件有75k+行，测试处理1k行就花了大概24分钟，速度实在太慢了。我想在代码里加入批量处理的逻辑，但我Python不算精通，网上相关的资料也不多。目前我已经写出了基础代码，但不知道怎么优化才能提升处理速度，有没有大佬能给点建议呀？

我的代码如下：

from geopy.geocoders import Nominatim
from pyspark.sql.functions import udf, col
from pyspark.sql.types import StructType, StructField, StringType, DoubleType
from geopy.exc import GeocoderTimedOut, GeocoderUnavailable

# Read the CSV file
df = spark.read.format("csv").option("header","true").load("Files/locations.csv")

# Limit to 1000 rows
df_limited = df.limit(1000)

# Define a function to get coordinates
def get_coordinates(address):
    geolocator = Nominatim(user_agent="address_geocoder")
    try:
        location = geolocator.geocode(address)
        if location:
            return (location.latitude, location.longitude)
        else:
            return (None, None)
    except (GeocoderTimedOut, GeocoderUnavailable):
        return (None, None)

# Create a UDF (User Defined Function)
coordinates_udf = udf(get_coordinates, StructType([
    StructField("latitude", DoubleType(), True),
    StructField("longitude", DoubleType(), True)
]))

# Apply the UDF to create new columns
df_with_coordinates = df_limited.withColumn("coordinates", coordinates_udf(col("ADDRESS")))

# Split the coordinates into separate latitude and longitude columns
df_final = df_with_coordinates.withColumn("latitude", col("coordinates.latitude")) \
                              .withColumn("longitude", col("coordinates.longitude")) \
                              .drop("coordinates")

# Display the final DataFrame
display(df_final)

几个亲测有效的优化建议： #

• 不要在UDF里重复创建Geopy实例
  你现在每处理一个地址就新建一个Nominatim对象，这会带来大量重复的初始化开销，还容易触发服务的限流机制。把实例移到UDF外面，或者用PySpark的广播变量分发到各个节点，避免重复创建：

  # 提前初始化并广播Geolocator实例
  geolocator = Nominatim(user_agent="address_geocoder")
  geolocator_broadcast = spark.sparkContext.broadcast(geolocator)
  
  def get_coordinates(address):
      # 从广播变量获取实例
      locator = geolocator_broadcast.value
      try:
          location = locator.geocode(address)
          return (location.latitude, location.longitude) if location else (None, None)
      except (GeocoderTimedOut, GeocoderUnavailable):
          return (None, None)
  

• 添加请求延迟，规避Nominatim限流
  Nominatim免费版有请求频率限制（官方建议每秒不超过1次），你的代码大概率因为请求太密集被限流了，导致处理速度骤降。可以在函数里加个1秒左右的延迟，虽然单条处理慢了，但总时间会因为不被限流反而缩短：

  import time
  
  def get_coordinates(address):
      locator = geolocator_broadcast.value
      try:
          location = locator.geocode(address)
          time.sleep(1)  # 添加延迟
          return (location.latitude, location.longitude) if location else (None, None)
      except (GeocoderTimedOut, GeocoderUnavailable):
          time.sleep(1)
          return (None, None)
  

• 用Pandas UDF替代普通UDF
  普通UDF是逐行处理的，效率很低。Pandas UDF是向量化处理，能大幅提升速度。这里给你改个批量处理的版本：

  from pyspark.sql.functions import pandas_udf
  import pandas as pd
  
  @pandas_udf(StructType([
      StructField("latitude", DoubleType(), True),
      StructField("longitude", DoubleType(), True)
  ]))
  def batch_get_coordinates(addresses: pd.Series) -> pd.DataFrame:
      locator = geolocator_broadcast.value
      results = []
      for addr in addresses:
          try:
              loc = locator.geocode(addr)
              results.append((loc.latitude, loc.longitude) if loc else (None, None))
              time.sleep(1)
          except (GeocoderTimedOut, GeocoderUnavailable):
              results.append((None, None))
              time.sleep(1)
      return pd.DataFrame(results, columns=["latitude", "longitude"])
  
  # 替换原来的UDF调用
  df_with_coordinates = df_limited.withColumn("coordinates", batch_get_coordinates(col("ADDRESS")))
  

• 缓存重复地址，减少重复请求
  如果你的数据里有大量重复地址，先对地址去重，只处理一次，再关联回原数据，能省掉很多不必要的请求：

  # 提取唯一地址
  unique_addrs = df.select("ADDRESS").distinct()
  # 处理唯一地址
  unique_addrs_with_coords = unique_addrs.withColumn("coordinates", coordinates_udf(col("ADDRESS")))
  # 关联回原数据
  df_final = df.join(unique_addrs_with_coords, on="ADDRESS", how="left")
  

• 考虑切换到付费地理编码服务
  如果免费的Nominatim实在满足不了需求，Google Maps、Mapbox这类付费服务支持更高的请求频率和批量接口，速度和稳定性都会好很多，适合处理大体积数据。

备注：内容来源于stack exchange，提问作者German Munoz