Efficient Way of subset a dataframe into small ones based on unique values and simultaneously write out to csv file

Viewed 364

What is the most efficient way to subset a large dataframe df into small subsets based on a unique/ filter condition? For example, I have a dataset with a dimension of 22050 rows with 5 columns, something like this

id, nationality, age, gender, income 
10001, France, 20, M, 45007
13328, UK, 52, F, 72308
11654, USA, 57, F, 95645
11765, UK, 39, M, 77343
10081, UAE, 41,M, 83117
10503, France, 22, F, 25665

There are 15 unique nationalities in the entire dataset, I want to subset the dataset into 15 dataframes based on the 15 unique countries and simultaneously write out the 15 dataframes in 15 csv output files.

Desired output should look like this

dataframe-one in a csv file

id, nationality, age, gender, income
10001, France, 20, M, 45007
10503, France, 22, F, 25665

dataframe-two in a csv file

13328, UK, 52, F, 72308
11765, UK, 39, M, 77343

likewise for dataframes 3 to 15

Here is my attempt:

fran = df[df.nationality == 'France']
fran.to_csv(file_name, sep=',')
uk = df[df.nationality =='UK']
uk.to_csv(file_name, sep=',')
USA = df[df.nationality == 'USA']
usa.to_csv(file_name, sep=',')

I want a more efficient way, apply | lambda or a loop approach

2 Answers

In base R, we can split the data by the 'nationality' column into a list of data.frame

lst1 <- split(df, df$nationality)

Then loop over the list and write it to different files

lapply(names(lst1), function(nm) write.csv(lst[[nm]], 
          paste0(nm, ".csv"), row.names = FALSE, quote = FALSE))

NOTE: The split method would be much faster than the == based subsetting

Assuming you want to use pandas you can loop over a grouped dataframe

for i, part_df in df.groupby('nationality'):
    part_df.to_csv(f'{i}.csv', index=False)

Your example is not clear if you want comma or tab delimiter. For a tab delimited file add sep='\t' as an argument to to_csv.

Related