如何基于PostgreSQL关联表动态创建分层文件夹导出GPS轨迹CSV数据集

前置依赖 #

仅需要用到Python标准库（os、csv）和你已经配置好的psycopg2，不需要额外安装第三方包。

核心实现思路 #

• 直接在数据库层完成两张表的关联，利用数据库的计算能力完成匹配，避免在Python侧做数组拆分、关联匹配的重复工作
• 采用psycopg2流式游标分批读取数据，避免一次性拉取全量数据占满客户端内存，适配大规模数据集的导出需求
• 按user_id、trip_id、timestamp的顺序做结果排序，保证同一个行程的GPS点连续返回，减少文件句柄的频繁切换，提升写入效率
• 自动创建目录结构，首次生成CSV文件时自动写入表头，后续同行程的数据直接追加，不重复写表头

完整可运行脚本 #

import os
import csv
import psycopg2
from config import config

# 定义导出文件的根目录
DATA_ROOT = "Data"
# 提前创建根目录，重复运行不会报错
os.makedirs(DATA_ROOT, exist_ok=True)

conn = None
cursor = None
csv_handle = None
# 记录当前正在写入的用户、行程，用来判断是否需要切换文件
current_user = None
current_trip = None
csv_writer = None

try:
    # 读取数据库连接配置，和你原有测试逻辑保持一致
    params = config()
    print('Connecting to the PostgreSQL database...')
    conn = psycopg2.connect(**params)
    # 创建服务端流式游标，避免一次性拉取全量数据到内存
    cursor = conn.cursor(name="gps_export_stream_cursor")

    # 执行关联查询，直接拿到所有需要的字段，排序保证同行程数据连续
    query = """
        SELECT
            s.user_id,
            t.trip_id,
            s.session_id,
            s.timestamp,
            s.lat,
            s.lon
        FROM trips t
        JOIN sessions s 
            ON s.session_id = ANY(t.session_ids)
        ORDER BY s.user_id, t.trip_id, s.timestamp
    """
    cursor.execute(query)
    print("Start exporting GPS trajectory data...")

    # 每次批量拉取1000条记录，平衡IO效率和内存占用
    while True:
        rows = cursor.fetchmany(size=1000)
        if not rows:
            break
        for row in rows:
            user_id, trip_id, session_id, ts, lat, lon = row
            # 切换到新的用户/行程时，关闭旧文件句柄，打开新文件
            if user_id != current_user or trip_id != current_trip:
                if csv_handle:
                    csv_handle.close()
                current_user = user_id
                current_trip = trip_id
                # 创建对应用户的文件夹
                user_dir = os.path.join(DATA_ROOT, str(current_user))
                os.makedirs(user_dir, exist_ok=True)
                # 定位到行程对应的CSV文件路径
                trip_file = os.path.join(user_dir, f"{current_trip}.csv")
                # 判断文件是否已存在，不存在则需要写入表头
                need_write_header = not os.path.exists(trip_file)
                csv_handle = open(trip_file, "a", newline="", encoding="utf-8")
                csv_writer = csv.writer(csv_handle)
                if need_write_header:
                    csv_writer.writerow(["session_id", "timestamp", "lat", "lon"])
            # 写入当前GPS点记录
            csv_writer.writerow([session_id, ts, lat, lon])

    print("Data export completed.")

except (Exception, psycopg2.DatabaseError) as error:
    print(f"Error during export: {error}")
finally:
    # 确保所有资源被正确释放
    if csv_handle:
        csv_handle.close()
    if cursor:
        cursor.close()
    if conn:
        conn.close()
        print("Database connection closed.")

运行说明 #

• 直接运行脚本即可自动生成符合要求的目录结构，和给出的示例目录完全一致
• 脚本支持断点续跑：如果导出过程中中断，重新运行不会重复写入已生成的表头，会继续完成剩余数据的导出
• 如果需要全量重新导出，仅需删除已生成的Data根目录后重新运行即可
• CSV文件写入时加了newline=""参数，避免Windows下用Excel打开出现空行，同时用utf-8编码保证特殊字符不会乱码

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