无法导出含美国城市州信息的Spark大数据框，求解决方案

现有数据处理流程 #

1. 读取数据并添加索引 #

import pyspark.sql.functions as F
from pyspark.sql.window import Window

sourcedata = spark.read.parquet(
    "...",
    header="true",
    inferSchema="true",
)

sourcedata = sourcedata.select("*").withColumn(
    "index", F.row_number().over(Window.orderBy(F.monotonically_increasing_id()))
)  # 创建索引用于后续关联

处理后DataFrame的列结构：

+-----------+--------------------+---------------------+---------------+------------------+------------------+---------+---------+-----------------+------------------+---------+------------+-----------------+---------+-------+-------+---------+---------+-----+
|vendor_name|Trip_Pickup_DateTime|Trip_Dropoff_DateTime|Passenger_Count|Trip_Distance     |Start_Lon         |Start_Lat|Rate_Code|store_and_forward|End_Lon           |End_Lat  |Payment_Type|Fare_Amt         |surcharge|mta_tax|Tip_Amt|Tolls_Amt|Total_Amt|index|
+-----------+--------------------+---------------------+---------------+------------------+------------------+---------+---------+-----------------+------------------+---------+------------+-----------------+---------+-------+-------+---------+---------+-----+
|VTS        |2009-01-04 02:52:00 |2009-01-04 03:02:00  |1              |2.63              |-73.991957        |40.721567|null     |null             |-73.993803        |40.695922|CASH        |8.9              |0.5      |null   |0.0    |0.0      |9.4      |1    |
|VTS        |2009-01-04 03:31:00 |2009-01-04 03:38:00  |3              |4.55              |-73.982102        |40.73629 |null     |null             |-73.95585         |40.76803 |Credit      |12.1             |0.5      |null   |2.0    |0.0      |14.6     |2    |
|VTS        |2009-01-03 15:43:00 |2009-01-03 15:57:00  |5              |10.35             |-74.002587        |40.739748|null     |null             |-73.869983        |40.770225|Credit      |23.7             |0.0      |null   |4.74   |0.0      |28.44    |3    |

数据类型：

root
|-- vendor_name: string (nullable = true)
|-- Trip_Pickup_DateTime: string (nullable = true)
|-- Trip_Dropoff_DateTime: string (nullable = true)
|-- Passenger_Count: long (nullable = true)
|-- Trip_Distance: double (nullable = true)
|-- Start_Lon: double (nullable = true)
|-- Start_Lat: double (nullable = true)
|-- Rate_Code: double (nullable = true)
|-- store_and_forward: double (nullable = true)
|-- End_Lon: double (nullable = true)
|-- End_Lat: double (nullable = true)
|-- Payment_Type: string (nullable = true)
|-- Fare_Amt: double (nullable = true)
|-- surcharge: double (nullable = true)
|-- mta_tax: double (nullable = true)
|-- Tip_Amt: double (nullable = true)
|-- Tolls_Amt: double (nullable = true)
|-- Total_Amt: double (nullable = true)
|-- index: integer (nullable = false)

2. 反向地理编码获取地址信息 #

from pyspark.sql.functions import col, split, size
from pyspark.sql.types import StringType
from geopy.geocoders import Nominatim

df = sourcedata.select(
    "index", col("Start_Lat").alias("latitude"), col("Start_Lon").alias("longitude")
)

# 定义反向地理编码UDF
@udf(StringType())
def reverse_geocode(latitude, longitude):
    latitude = float(latitude)
    longitude = float(longitude)
    geolocator = Nominatim(user_agent="reverse_geocode")
    location = geolocator.reverse((latitude, longitude), language="en")
    if location:
        return location.address
    else:
        return None

df_with_location = df.withColumn(
    "location", reverse_geocode(df["latitude"], df["longitude"])
)

处理后结果：

+-----+---------+------------------+--------------------------------------------------------------------------------------------------------------------------+
|index|latitude |longitude         |location                                                                                                                  |
+-----+---------+------------------+--------------------------------------------------------------------------------------------------------------------------+
|1    |40.721567|-73.991957        |183, Chrystie Street, Manhattan Community Board 3, Manhattan, New York County, New York, 10002, United States             |
|2    |40.73629 |-73.982102        |359, 2nd Avenue, Manhattan Community Board 6, Manhattan, New York County, New York, 10010, United States                  |

3. 提取城市和州信息 #

df_with_city_state = (
    df_with_location.withColumn("location", split("location", ","))
    .withColumn("city", col("location")[size("location") - 4])
    .withColumn("state", col("location")[size("location") - 5])
)

df_with_city_state = df_with_city_state.select("index", "city", "state")

4. 导出超时问题 #

尝试导出时出现超时：

df_with_location.write.csv('test', mode='overwrite')

相同行数的其他DataFrame可正常导出，推测问题出在地理编码后的DataFrame处理环节。

解决方法 #

1. 优化反向地理编码UDF #

• 复用Geolocator实例：不在UDF内部重复创建Nominatim实例，减少资源开销：

geolocator = Nominatim(user_agent="reverse_geocode")

@udf(StringType())
def reverse_geocode(latitude, longitude):
    try:
        lat = float(latitude)
        lon = float(longitude)
        location = geolocator.reverse((lat, lon), language="en", timeout=10)
        return location.address if location else None
    except Exception as e:
        return None

• 添加超时参数：在reverse方法中设置timeout，避免单个请求长时间阻塞。
• 异常捕获：捕获网络请求异常，防止单条数据处理失败导致整个任务中断。

2. 调整Spark配置 #

增加Spark的超时和资源配置，适配IO密集型任务：

# 调整网络超时和心跳间隔
spark.conf.set("spark.network.timeout", "300s")
spark.conf.set("spark.executor.heartbeatInterval", "60s")

# 根据集群资源调整执行器内存和核心数
spark.conf.set("spark.executor.memory", "8g")
spark.conf.set("spark.executor.cores", "4")

3. 分批次导出 #

对DataFrame进行分区，控制单次导出的数据量：

# 按index范围分区，拆分数据
df_with_city_state = df_with_city_state.repartitionByRange(10, "index")

# 导出时指定表头，控制输出文件数量
df_with_city_state.write.option("header", "true").csv('test', mode='overwrite')

4. 预处理数据减少开销 #

• 过滤无效数据：提前过滤latitude/longitude为null或异常值的行，减少无效网络请求。
• 缓存中间结果：对地理编码后的DataFrame缓存，避免重复计算：

df_with_location.cache()
df_with_location.count()  # 触发缓存执行

5. 改用批量地理编码方案 #

如果数据量较大，单条请求的UDF效率极低，建议：

• 使用支持批量请求的地理编码服务
• 下载离线地理数据（如GeoJSON）进行本地匹配，完全避免网络请求开销

内容的提问来源于stack exchange，提问作者Quang Phạm