如何在PySpark/Databricks中将DataFrame保存为shp/GeoJSON？

一、修复Geopandas在Databricks的路径问题 #

你遇到的核心问题是：Databricks的/mnt/挂载路径指向分布式存储（如ADLS、S3），而Geopandas的to_file是本地文件系统操作，无法直接写入分布式存储路径。本地运行时操作的是本地磁盘，所以可行，但Databricks中需要先写入节点本地临时目录，再复制到挂载路径。同时你的代码存在笔误，需修正。

修正后的代码 #

import geopandas as gpd
import os
from shutil import copyfile

# 转换PySpark DataFrame为Pandas DataFrame
pd_df = df.toPandas()
# 修正笔误：用pd_df而非pd，列名对应你的Wkt列
pd_df['geometry'] = pd_df['Wkt'].apply(gpd.GeoSeries.from_wkt)

# 创建GeoDataFrame
gdf = gpd.GeoDataFrame(pd_df, geometry='geometry')

# 1. 写入节点本地临时路径
local_temp_prefix = "/tmp/shapefile"
gdf.to_file(local_temp_prefix)

# 2. 复制所有Shapefile相关文件到挂载路径
target_dir = os.path.join(MOUNT_POINT, output_folder)
for ext in ['.shp', '.shx', '.dbf', '.prj']:
    src_file = f"{local_temp_prefix}{ext}"
    dst_file = os.path.join(target_dir, f"shapefile{ext}")
    if os.path.exists(src_file):
        copyfile(src_file, dst_file)

二、适合大数据量的Spark原生方案（推荐） #

10万行转Pandas虽然可行，但数据量进一步增大时会触发节点内存问题。推荐使用Spark原生地理空间库处理，无需转换为Pandas。

方法1：使用GeoSpark #

GeoSpark是开源的Spark地理空间库，支持分布式读写Shapefile/GeoJSON。

前置步骤

在Databricks集群的库管理中添加GeoSpark相关包（如org.datasyslab:geospark-sql_2.12:1.4.1）。

代码示例

from pyspark.sql import functions as F
from pyspark.sql import SparkSession
from geospark.register import upload_jars
from geospark.utils import GeoSparkKryoRegistrator
from geospark.sql.functions import ST_GeomFromWKT

# 初始化GeoSpark
upload_jars()
spark = SparkSession.builder \
    .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer") \
    .config("spark.kryo.registrator", GeoSparkKryoRegistrator.getName) \
    .getOrCreate()

# 将WKT列转换为GeoSpark几何类型
geo_df = df.withColumn("geometry", ST_GeomFromWKT(F.col("Wkt")))

# 保存为Shapefile
geo_df.write.format("geospark") \
    .option("shapefile", "true") \
    .save(os.path.join(MOUNT_POINT, output_folder, "shapefile_output"))

# 保存为GeoJSON
geo_df.write.format("geospark") \
    .option("geojson", "true") \
    .save(os.path.join(MOUNT_POINT, output_folder, "geojson_output"))

方法2：使用Databricks GeoAnalytics（需许可） #

如果你的Databricks集群有GeoAnalytics许可，可以直接用原生API处理：

from pyspark.sql import functions as F

# 转换WKT为几何类型
geo_df = df.withColumn("geometry", F.expr("ST_GeomFromWKT(Wkt)"))

# 保存为Shapefile
geo_df.write.format("shapefile") \
    .save(os.path.join(MOUNT_POINT, output_folder, "shapefile_output"))

# 保存为GeoJSON
geo_df.write.format("json") \
    .option("recordFormat", "geojson") \
    .save(os.path.join(MOUNT_POINT, output_folder, "geojson_output"))

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