如何用Python的Pandas正确加载标准格式的ARFF数据集？

报错原因 #

arff.load()返回的是生成器对象，生成器不支持下标（[]）操作，直接尝试访问生成器的['data']键就会抛出TypeError: 'generator' object is not subscriptable。

解决方案 #

步骤1：导入依赖库 #

import arff
import pandas as pd

步骤2：正确读取ARFF文件 #

标准ARFF文件通常只包含一个数据集，可通过next()提取生成器中的唯一数据集对象，或把生成器转为列表后访问：

方法一：用next()提取数据集

with open('ALOI.arff', 'r', encoding='utf-8') as f:
    dataset = next(arff.load(f))

方法二：转换为列表后访问

with open('ALOI.arff', 'r', encoding='utf-8') as f:
    dataset = list(arff.load(f))[0]

步骤3：转换为Pandas DataFrame #

从数据集的attributes字段提取列名，再将data字段转为DataFrame：

# attributes格式为[(属性名, 数据类型), ...]，提取属性名作为列名
columns = [attr_name for attr_name, _ in dataset['attributes']]
df = pd.DataFrame(dataset['data'], columns=columns)

可选：处理标称属性 #

如果ARFF文件包含标称类型属性，可在加载时添加encode_nominal=True参数自动将其转为数值：

with open('ALOI.arff', 'r', encoding='utf-8') as f:
    dataset = next(arff.load(f, encode_nominal=True))

完整示例代码 #

import arff
import pandas as pd

with open('ALOI.arff', 'r', encoding='utf-8') as f:
    dataset = next(arff.load(f, encode_nominal=True))

columns = [attr[0] for attr in dataset['attributes']]
df = pd.DataFrame(dataset['data'], columns=columns)

# 查看前5行验证结果
print(df.head())

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