Чтение в определенных столбцах набора данных с помощью Python

Question

Я пытался взять очень большой 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?

Answer 1

Полагаю, вы просто хотите скопировать содержимое файла file_path в reduced.csv со всеми удаленными столбцами, которые начинаются с одного из символов X, P, W и без столбцов SCH_ID, STRAT_ID, STU_ID.

Если это так, вы можете сделать это с пандами следующим образом:

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_csv и to_csv.

Btw. если вы хотите упростить свой код и использовать 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 завершен.