如何通过Python将IMDB多值字段转为数组存入MongoDB？

使用Pandas的read_csv读取IMDB的TSV格式数据集（以name.basics.tsv为例），前期将处理后的数据存入MongoDB。临近收尾时，因字段类型问题无法用Metabase制作仪表盘，需要在读取文件或处理DataFrame阶段，把含逗号分隔信息的字段（如knownForTitles、primaryProfession）转为数组类型后存入MongoDB。已知可手动修改MongoDB字段类型，但不想手动操作，也不想切换到MySQL，求可行的Python处理方案。

当前读取代码 #

import pandas as pd
from pymongo import MongoClient
import numpy as np
from ast import literal_eval

mc = MongoClient("mongodb://localhost:27017/")
db = mc._database_default_options(name='imdbdash')

file = "../IMDBDATASET/name.basics.tsv"
dfname = pd.read_csv(filepath_or_buffer=file,
                      sep = '\t',
                      #dtype={'primaryProfession': np.array , 'knownForTitles':'np.array'},
                      #converters={'primaryProfession': literal_eval,'knownForTitles': literal_eval,},
                      low_memory=False,
                      header=0,
                      na_values=["\\N","nan","NaN"],
                      keep_default_na=False,
                      na_filter=True,
                      verbose=False)

已尝试的无效方法 #

• 使用converters={'primaryProfession': ast.literal_eval,'knownForTitles': ast.literal_eval}
• 将列转为numpy数组：test = np.array(dfname['knownForTitles'].tolist())
• 字符串拼接操作："',''.join(map(str,dfname['knownForTitles']))
• 指定dtype={'primaryProfession': np.array , 'knownForTitles':'np.array'}

可行解决方案 #

方案1：读取文件时直接转换（推荐） #

自定义转换器函数，将逗号分隔的字符串转为列表，同时处理空值：

import pandas as pd
from pymongo import MongoClient

def str_to_list(s):
    # 处理空值或空字符串，返回空列表
    if pd.isna(s) or s.strip() == '':
        return []
    # 按逗号分割并去除每个元素的前后空格
    return [item.strip() for item in s.split(',')]

# 修正数据库连接方式
mc = MongoClient("mongodb://localhost:27017/")
db = mc['imdbdash']

file = "../IMDBDATASET/name.basics.tsv"
dfname = pd.read_csv(filepath_or_buffer=file,
                      sep='\t',
                      converters={
                          'primaryProfession': str_to_list,
                          'knownForTitles': str_to_list
                      },
                      low_memory=False,
                      header=0,
                      na_values=["\\N", "nan", "NaN"],
                      keep_default_na=False,
                      na_filter=True,
                      verbose=False)

# 存入MongoDB，列表会自动转为MongoDB数组类型
db['name_basics'].insert_many(dfname.to_dict('records'))

方案2：读取数据后再处理 #

如果已经读取了原始数据，可对DataFrame的目标列单独处理：

# 对已读取的dfname进行字段转换
dfname['knownForTitles'] = dfname['knownForTitles'].apply(str_to_list)
dfname['primaryProfession'] = dfname['primaryProfession'].apply(str_to_list)

# 存入MongoDB
db['name_basics'].insert_many(dfname.to_dict('records'))

关键说明 #

• 自定义的str_to_list函数会处理空值、空字符串，避免存入MongoDB时出现null，同时去除元素前后空格保证数据整洁。
• 原代码中获取数据库的方式错误，应改为mc['imdbdash']，否则无法正确连接目标数据库。
• 使用df.to_dict('records')将DataFrame转为字典列表，存入MongoDB时，Python列表会自动映射为MongoDB的数组类型，Metabase可正常识别。

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