在导入为 pandas dataframe 之前从压缩的 csv 文件获取标题

Question

I am trying to import csv files as a pandas dataframe, where the csv files are inside a zip file. For efficient importing, I'm trying to get the headers first before I load it into a pandas dataframe.

What I tried so far is this:

from zipfile import ZipFile
from io import TextIOWrapper
import pandas as pd

with ZipFile(zip_path, 'r') as zipfile:
    with zipfile.open(file_path, 'r') as file:
        reader = csv.reader(TextIOWrapper(file, 'utf-8', newline=''))
        headers = next(reader)

        df = pd.read_csv(file)

The problem is when I get the headers with next(reader) the underlying file is affected, and the file is imported as a pandas dataframe without headers.

Would really appreciate any fix.

Answer 1

You can reset CSV iterator using function seek():

with ZipFile('test.zip', 'r') as zipfile:
    with zipfile.open('test.csv', 'r') as file:
        reader = csv.reader(TextIOWrapper(file, 'utf-8', newline=''))
        headers = next(reader)
        # reset CSV iterator
        file.seek(0)
        df = pd.read_csv(file)