Python+GIS — shapefile批量导入PostgreSQL (PostGIS)

背景

Postgresql+PostGIS 是存储、处理、分析空间数据的常用解决方案。相信多数GISer和我一样，安装好PostGIS的第一件事就是研究如何将Esri Shapefile导入数据库。常见方法是通过PostGIS提供的数据导入工具，或者通过shp2psql命令行工具进行导入。
最近在研究Python空间数据处理时，想要通过Python脚本将shp文件批量入库，查到了如下两篇文章：

• https://zhuanlan.zhihu.com/p/41249261
• https://blog.csdn.net/jin80506/article/details/91386307

尝试运行了一下代码，发现存在导入报错、中文乱码、无法导入MultiPolygon等问题。
于是基于二者的代码进行了改造，解决了这些问题。

步骤

1. 安装依赖包

需要的python包主要有：sqlalchemy, geoalchemy2, geopandas, psycopg2，这些包有的依赖其它包，安装起来比较麻烦，建议大家使用anaconda进行安装。

2. 准备工作

首先创建一个测试数据库gis_test。
然后加载PostGIS扩展:

CREATE EXTENSION postgis;

当然，你也可以使用python连接PostgreSQL进行数据库的创建和PostGIS扩展的加载，这里不再赘述。

3. 代码

本文代码删掉了上述参考文章中“解析shp的geometry code，获取geometry类型”的部分，直接将空间数据指定为GEOMETRY 类型。

"""
shp2pgsql.py
import Shapefile to PostgreSQL.
author: mango
version: 0.1.0
Compatible with Python versions 3.x
"""
import geopandas as gpd
from sqlalchemy import create_engine
from geoalchemy2 import Geometry, WKTElement
import os

def shp2pgsql(file, engine):
    """单个shp文件入库"""
    file_name = os.path.split(file)[1]
    print('正在写入:'+file)
    tbl_name = file_name.split('.')[0]  # 表名
    map_data = gpd.GeoDataFrame.from_file(file)
    spatial_ref = map_data.crs.srs.split(':')[-1]  # 读取shp的空间参考
    map_data['geometry'] = map_data['geometry'].apply(
        lambda x: WKTElement(x.wkt, spatial_ref))
    # geopandas 的to_sql()方法继承自pandas, 将GeoDataFrame中的数据写入数据库
    map_data.to_sql(
        name=tbl_name,
        con=engine,
        if_exists='replace', # 如果表存在，则替换原有表
        chunksize=1000,  # 设置一次入库大小，防止数据量太大卡顿
        # 指定geometry的类型,这里直接指定geometry_type='GEOMETRY'，防止MultiPolygon无法写入
        dtype={'geometry': Geometry(
            geometry_type='GEOMETRY', srid=spatial_ref)},
        method='multi'
    )
    return None


def shp2pgsql_batch(dir_name, username, password, host, port, dbname):
    """创建批量任务"""
    os.chdir(dir_name)  # 改变当前工作目录到指定路径
    engine = create_engine(username, password, host, port, dbname)
    file_list = os.listdir(dir_name)
    for file in file_list:
        if file.split('.')[-1] == 'shp':
            file = os.path.abspath(file)
            shp2pgsql(file, engine)
    return None


# 执行任务计划
if __name__ == '__main__':
    file_path = r'D:\Data\mango'
    username = 'postgres'
    password = '123ewq'
    host = '127.0.0.1'
    port = '5432'
    dbname = 'gis_test'
    shp2pgsql_batch(file_path, username, password, host, port, dbname)

用改进后的代码测试了多种点、线、面数据，均能成功写入，且无中文乱码问题。大功告成！