基于Pandas DataFrame创建用户子目录及会话CSV文件

给定数据 #

我们有如下包含用户、会话ID、日志时间、经纬度的Pandas DataFrame：

import pandas as pd

data = {'user': [7, 7, 7, 7, 7, 7, 7, 11, 11, 11],
        'session_id': [15, 15, 15, 15, 31, 31, 31, 43, 43, 43],
        'logtime': ['2016-04-13 07:58:40','2016-04-13 07:58:41','2016-04-13 07:58:42',
                   '2016-04-13 07:58:43','2016-04-01 20:29:37','2016-04-01 20:29:42',
                   '2016-04-01 20:29:47','2016-03-30 06:21:59','2016-03-30 06:22:04',
                   '2016-03-30 06:22:09'],
        'lat': [41.1872084,41.1870716,41.1869719,41.1868664,41.1471521,
                41.1472466,41.1473038,41.2372125,41.2371444,41.2369725],
        'lon': [-8.6038931,-8.6037318,-8.6036908,-8.6036423,-8.5878757,
                -8.5874314,-8.586632,-8.6720773,-8.6721269,-8.6718833]}

d = pd.DataFrame(data)

需求说明 #

• 在当前工作目录下为每个用户创建子目录；
• 每个用户子目录下，为该用户的每个会话创建一个CSV文件；
• 每个CSV文件仅写入对应会话的logtime、lat、lon字段（不含session_id），文件名按file1.csv、file2.csv的格式命名；
• 遍历所有用户完成上述操作。

预期生成的目录结构及文件内容如下：

Data/
├── 11
│   └── file1.csv
|          logtime,lat,lon
|          2016-03-30 06:21:59,41.2372125,-8.6720773
|          2016-03-30 06:22:04,41.2371444,-8.6721269
|          2016-03-30 06:22:09,41.2369725,-8.6718833 
└── 7
    ├── file1.csv
    |      logtime,lat,lon
    |      2016-04-13 07:58:40,41.187208,-8.603893
    |      2016-04-13 07:58:41,41.187072,-8.603732
    |      2016-04-13 07:58:42,41.186972,-8.603691
    |      2016-04-13 07:58:43,41.186866,-8.603642
    └── file2.csv
           logtime,lat,lon
           2016-04-01 20:29:37,41.147152,-8.587876
           2016-04-01 20:29:42,41.147247,-8.587431
           2016-04-01 20:29:47,41.147304,-8.586632

解决方案代码 #

import os
import pandas as pd

# 创建根目录Data，已存在则不报错
root_dir = "Data"
os.makedirs(root_dir, exist_ok=True)

# 按用户、会话分组遍历数据
for (user_id, session_id), group in d.groupby(['user', 'session_id']):
    # 创建用户子目录
    user_dir = os.path.join(root_dir, str(user_id))
    os.makedirs(user_dir, exist_ok=True)
    
    # 确定当前会话对应的文件序号
    user_sessions = d[d['user'] == user_id]['session_id'].unique()
    file_num = list(user_sessions).index(session_id) + 1
    filename = f"file{file_num}.csv"
    
    # 提取指定字段并保存为CSV，不写入行索引
    group[['logtime', 'lat', 'lon']].to_csv(os.path.join(user_dir, filename), index=False)

代码说明 #

1. 先创建根目录Data，exist_ok=True参数避免目录已存在时抛出异常；
2. 通过groupby(['user', 'session_id'])对数据按用户和会话进行分组；
3. 为每个用户创建专属子目录，确保目录存在；
4. 提取当前用户的所有唯一会话，确定当前会话对应的文件序号（从1开始计数）；
5. 筛选分组数据中的logtime、lat、lon字段，保存为指定文件名的CSV文件，index=False避免将行索引写入文件。

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