将大型 DataFrame 拆分为包含列中唯一值记录的 DataFrame

Question

A csv file has 90 million rows. One of the Columns in named "State". It has 12 unique values at present. (The count of unique values in the "State" column is dynamic and can change with each csv file.)

I want to split the DataFrame into smaller chunks and then save State-wise files. The code below is not working.

source_path = "DataJune.txt"
for i,chunk in enumerate(pd.read_csv(source_path, sep = '|',chunksize=1000000)):
    dfs = dict(tuple(chunk.groupby('State')))
    for i, df in dfs.items():
        df = df.append(df)
        df.to_csv("tempcsv/" + i +".csv",sep=",", index = False)

Answer 1

IIUC, Try:

source_path = "DataJune.txt"

from collections import defaultdict

def def_value():
    return pd.DataFrame()
      
# Defining the dict
d = defaultdict(def_value)

for i,chunk in enumerate(pd.read_csv(source_path, sep = '|',chunksize=2)):
    chunk_states = chunk['State'].unique()
    for state in chunk_states:
        d[state]=d[state].append(chunk[chunk['State']==state])
for i, df in d.items():
    df.to_csv("tempcsv/" + str(i) +".csv",sep=",", index = False)

Another version, based on the @Corralien comment:

source_path = "DataJune.txt"

for i,chunk in enumerate(pd.read_csv(source_path, sep = '|',chunksize=2)):
    chunk_states = chunk['State'].unique()
    
    for state in chunk_states:
        with open("tempcsv/" + str(state) +".csv",mode='a+') as file:
            for i, row in chunk[chunk['State']==state].iterrows():
                file.write(','.join([str(x) for x in row]))
                file.write('\n')

Another version:

source_path = "DataJune.txt"
from os.path import exists
import csv

for i,chunk in enumerate(pd.read_csv(source_path, sep = '|',chunksize=2)):
    chunk_states = chunk['State'].unique()
    
    for state in chunk_states:
        path = "tempcsv/" + str(state) +".csv"
        file_exists = exists(path)
        if not file_exists:
            with open(path,newline='',mode='a+') as file:
                writer = csv.writer(file)
                writer.writerow(chunk.columns)
                print(chunk.columns)
        with open(path,newline='',mode='a+') as file:
            writer = csv.writer(file)
            writer.writerows(chunk[chunk['State']==state].values)

Answer 2

You can use:

import pandas as pd
import os

source_path = 'DataJune.txt'
fps = {}

for chunk in pd.read_csv(source_path, sep='|', chunksize=1000000, dtype=object):
    for state, df in chunk.groupby('State'):
        # New state, create a new file and write headers
        if state not in fps:
            fps[state] = open(f'tempcsv/{state}.csv', 'w')
            fps[state].write(f"{','.join(df.columns)}{os.linesep}")

        # Write data without headers
        df.to_csv(fps[state], index=False, header=False)

# Close files properly
for fp in fps.values():
    fp.close()
del fps

Update

Try to replace:

# Write data without headers
df.to_csv(fps[state], index=False, header=False)

By

# Write data without headers
g = (row.strip() for row in df.to_csv(index=False, header=None, sep=',').split(os.linesep) if row)
print(*g, sep=os.linesep, file=fps[state])