Create new data frames from existing data frame based on unique column values

Question

I have a large data set (4.5 million rows, 35 columns). The columns of interest are company_id (string) and company_score (float). There are approximately 10,000 unique company_id's.

company_id    company_score    date_submitted    company_region
AA            .07              1/1/2017          NW
AB            .08              1/2/2017          NE
CD            .0003            1/18/2017         NW

My goal is to create approximately 10,000 new dataframes, by unique company_id, with only the relevant rows in that data frame.

The first idea I had was to create the collection of data frames shown below, then loop through the original data set and append in new values based on criteria.

company_dictionary = {}  
for company in df['company_id']:  
    company_dictionary[company_id] = pd.DataFrame([])

Is there a better way to do this by leveraging pandas? i.e., is there a way I can use a built-in pandas function to create new filtered dataframes with only the relevant rows?

Edit: I tried a new approach, but I'm now encountering an error message that I don't understanding.

    [In]  unique_company_id = np.unique(df[['ID_BB_GLOBAL']].values)
    [In]  unique_company_id
    [Out] array(['BBG000B9WMF7', 'BBG000B9XBP9', 'BBG000B9ZG58', ..., 'BBG00FWZQ3R9',
       'BBG00G4XRQN5', 'BBG00H2MZS56'], dtype=object)
    [In]  for id in unique_company_id:
    [In]      new_df = df[df['id'] == id]
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
C:\get_loc(self, key, method, tolerance)
   2133             try:
-> 2134                 return self._engine.get_loc(key)
   2135             except KeyError:

pandas\index.pyx in pandas.index.IndexEngine.get_loc (pandas\index.c:4433)()

pandas\index.pyx in pandas.index.IndexEngine.get_loc (pandas\index.c:4279)()

pandas\src\hashtable_class_helper.pxi in pandas.hashtable.PyObjectHashTable.get_item (pandas\hashtable.c:13742)()

pandas\src\hashtable_class_helper.pxi in pandas.hashtable.PyObjectHashTable.get_item (pandas\hashtable.c:13696)()

KeyError: 'id'

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
<ipython-input-50-dce34398f1e1> in <module>()
      1 for id in unique_bank_id:
----> 2     new_df = df[df['id'] == id]

C:\ in __getitem__(self, key)
   2057             return self._getitem_multilevel(key)
   2058         else:
-> 2059             return self._getitem_column(key)
   2060 
   2061     def _getitem_column(self, key):

C:\ in _getitem_column(self, key)
   2064         # get column
   2065         if self.columns.is_unique:
-> 2066             return self._get_item_cache(key)
   2067 
   2068         # duplicate columns & possible reduce dimensionality

C:\ in _get_item_cache(self, item)
   1384         res = cache.get(item)
   1385         if res is None:
-> 1386             values = self._data.get(item)
   1387             res = self._box_item_values(item, values)
   1388             cache[item] = res

C:\ in get(self, item, fastpath)
   3541 
   3542             if not isnull(item):
-> 3543                 loc = self.items.get_loc(item)
   3544             else:
   3545                 indexer = np.arange(len(self.items))[isnull(self.items)]

C:\ in get_loc(self, key, method, tolerance)
   2134                 return self._engine.get_loc(key)
   2135             except KeyError:
-> 2136                 return self._engine.get_loc(self._maybe_cast_indexer(key))
   2137 
   2138         indexer = self.get_indexer([key], method=method, tolerance=tolerance)

pandas\index.pyx in pandas.index.IndexEngine.get_loc (pandas\index.c:4433)()

pandas\index.pyx in pandas.index.IndexEngine.get_loc (pandas\index.c:4279)()

pandas\src\hashtable_class_helper.pxi in pandas.hashtable.PyObjectHashTable.get_item (pandas\hashtable.c:13742)()

pandas\src\hashtable_class_helper.pxi in pandas.hashtable.PyObjectHashTable.get_item (pandas\hashtable.c:13696)()

KeyError: 'id'

Answer 1

You can groupby company_id column and convert its result into a dictionary of DataFrames:

import pandas as pd

df = pd.DataFrame({
    "company_id": ["AA", "AB", "AA", "CD", "AB"],
    "company_score": [.07, .08, .06, .0003, .09],
    "company_region": ["NW", "NE", "NW", "NW", "NE"]})

# Approach 1
dict_of_companies = {k: v for k, v in df.groupby('company_id')}

# Approach 2
dict_of_companies = dict(tuple(df.groupby("company_id")))

import pprint
pprint.pprint(dict_of_companies)

Output:

{'AA':   company_id company_region  company_score
0         AA             NW           0.07
2         AA             NW           0.06,
 'AB':   company_id company_region  company_score
1         AB             NE           0.08
4         AB             NE           0.09,
 'CD':   company_id company_region  company_score
3         CD             NW         0.0003}