首页 > 解决方案 > 使用python读取数据集的特定列

问题描述

我一直在尝试获取一个非常大的 csv 文件并将其读入 python 并编写一个新的缩减 csv 文件。我创建了一个我想使用的列名列表。下面是我正在尝试使用的代码

redfile = open(file_path,'r')

import csv
reader=csv.reader(redfile)
names=next(reader)
for elem in names:
        if elem.startswith("W")==True:
            names.remove(elem)
for elem in names:
        if elem.startswith("P")==True:
            names.remove(elem)
for elem in names:
        if elem.startswith("X")==True:
            names.remove(elem)
names.remove("SCH_ID")
names.remove("STRAT_ID")
names.remove("STU_ID")

nameind = []
line0 = ''

wfile = open('reduced.csv','w')
for i, line in enumerate(redfile):
    redarray = [x for x in line.split(",")]
    line1 = ''
    if i == 0:
        for ii in range(0,len(redarray)):
            if redarray[ii] in names:
                nameind.append(ii)
                line0 = line0+redarray[ii]+','
        line0 = line0[:-1]
        print(line0)
        wfile.write(line0)
        wfile.write('\n')
        nameindarray = np.array(nameind)
    elif i < 25000:
        for ii in nameind:
            line1 = line1+redarray[ii]+','
        line1 = line1[:-1]
        wfile.write(line1)
        wfile.write('\n')
    else:
        break
redfile.close()
wfile.close()
print(i)

如您所见,redfile 由用户选择,而 names 是特定列名的数组。该程序仅在 2 小时左右后继续运行。作为参考,大约有 24,000 行数据和大约 5000 列。现在最后,如何通过不包括具有特定值(例如 -5)的列来减少列的数量?

标签: pythonpython-3.xnumpycsv

解决方案


我想,您只想将文件的内容复制file_pathreduced.csv删除所有列的文件中,这些列以字符之一开头X, PW并且没有列SCH_ID, STRAT_ID, STU_ID

如果是这样,您可以像这样使用 pandas 来执行此操作:

import pandas as pd

# read the first row only to get the column names
df= pd.read_csv(file_path, sep=',', dtype='str', nrows=1)
use_cols= [col for col in df.columns if col[:1] not in 'XPW' and col not in ['SCH_ID', 'STRAT_ID', 'STU_ID']]

df= pd.read_csv(file_path, sep=',', dtype='str', usecols=use_cols)
df.to_csv('reduced.csv', index=False, sep=',')

请将此视为伪代码,因为我无法在没有数据的情况下对其进行测试,但我非常有信心它可以工作。如果事实证明引用不是您喜欢的那样,您可以尝试将quotechar关键字添加到read_csvandto_csv中。

顺便提一句。如果您想简化代码并用于with确保文件在任何情况下都已关闭,则可以重写最后一个 while 循环,例如:

with open('reduced.csv','w') as wfile:
    for i, line in enumerate(redfile):
        redarray = list(line.split(','))
        line1 = ''
        if i == 0:
            for ii, token in enumerate(redarray):
                if token in names:
                    nameind.append(ii)
                    line0= line0 + token + ','
            line0 = line0[:-1]
            print(line0)
            wfile.write(line0)
            wfile.write('\n')
            nameindarray = np.array(nameind)
    elif i < 25000:
        line1= ','.join([redarray[i] for i in nameind])
        wfile.write(line1)
        wfile.write('\n')
    else:
        break

如果您想切换到第二个建议,您可能还想在一个with子句中打开您的输入文件。如果使用with,则无需显式关闭文件。with当块终止时,这会自动为您完成。


推荐阅读