Expand Pandas DataFrame Column with JSON Object

Question

I'm looking for a clean, fast way to expand a pandas dataframe column which contains a json object (essentially a dict of nested dicts), so I could have one column for each element in the json column in json normalized form; however, this needs to retain all of the original dataframe columns as well. In some instances, this dict might have a common identifier I could use to merge with the original dataframe, but not always. For example:

import pandas as pd
import numpy as np
df = pd.DataFrame([
    {
        'col1': 'a',
        'col2': {'col2.1': 'a1', 'col2.2': {'col2.2.1': 'a2.1', 'col2.2.2': 'a2.2'}},
        'col3': '3a'
    },
    {
        'col1': 'b',
        'col2': np.nan,
        'col3': '3b'
    },
    {
        'col1': 'c',
        'col2': {'col2.1': 'c1', 'col2.2': {'col2.2.1': np.nan, 'col2.2.2': 'c2.2'}},
        'col3': '3c'
    }
])

Here is a sample dataframe. As you can see, col2 is a dict in all of these cases which has another nested dict inside of it, or could be a null value, containing nested elements I would like to be able to access. (For the nulls, I would want to be able to handle them at any level--entire elements in the dataframe, or just specific elements in the row.) In this case, they have no ID that could link up to the original dataframe. My end goal would be essentially to have this:

final = pd.DataFrame([
    {
        'col1': 'a',
        'col2.1': 'a1',
        'col2.2.col2.2.1': 'a2.1',
        'col2.2.col2.2.2': 'a2.2',
        'col3': '3a'
    },
    {
        'col1': 'b',
        'col2.1': np.nan,
        'col2.2.col2.2.1': np.nan,
        'col2.2.col2.2.2': np.nan,
        'col3': '3b'
    },
    {
        'col1': 'c',
        'col2.1': 'c1',
        'col2.2.col2.2.1': np.nan,
        'col2.2.col2.2.2': 'c2.2',
        'col3': '3c'
    }
])

In my instance, the dict could have up to 50 nested key-value pairs, and I might only need to access a few of them. Additionally, I have about 50 - 100 other columns of data I need to preserve with these new columns (so an end goal of around 100 - 150). So I suppose there might be two methods I'd be looking for--getting a column for each value in the dict, or getting a column for a select few. The former option I haven't yet found a great workaround for; I've looked at some prior answers but found them to be rather confusing, and most threw errors. This seems especially difficult when there are dicts nested inside of the column. To attempt the second solution, I tried the following code:

def get_val_from_dict(row, col, label):
    if pd.isnull(row[col]):
        return np.nan
    
    norm = pd.json_normalize(row[col])
    
    try:
        return norm[label]
    except:
        return np.nan


needed_cols = ['col2.1', 'col2.2.col2.2.1', 'col2.2.col2.2.2']


for label in needed_cols:
    df[label] = df.apply(get_val_from_dict, args = ('col2', label), axis = 1)

This seemed to work for this example, and I'm perfectly happy with the output, but for my actual dataframe which had substantially more data, this seemed a bit slow--and, I would imagine, is not a great or scalable solution. Would anyone be able to offer an alternative to this sluggish approach to resolving the issue I'm having?

(Also, apologies also about the massive amounts of nesting in my naming here. If helpful, I am adding in several images of the dataframes below--the original, then the target, and then the current output.)

user_2570261 · Answer

instead of using apply or pd.json_normalize on the column that has a dictionary, convert the whole data frame to dictionary & use pd.json_normalize on it, finally picking the fields you wish to keep. This works because while the individual column for any given row may be null, the entire row would not be.

example:

# note that this method also prefixes an extra `col2.` 
# at the start of the names of the denested data, 
# which is not present in the example output
# the column renaming conforms to your desired name.
import re
final_cols = ['col1', 'col2.col2.1', 'col2.col2.2.col2.2.1', 'col2.col2.2.col2.2.2', 'col3']
out = pd.json_normalize(df.to_dict(orient='records'))[final_cols]
out.rename(columns=lambda x: re.sub(r'^col2\.', '', x), inplace=True)
out
# out:
  col1 col2.1 col2.2.col2.2.1 col2.2.col2.2.2 col3
0    a     a1            a2.1            a2.2   3a
1    b    NaN             NaN             NaN   3b
2    c     c1             NaN            c2.2   3c

but for my actual dataframe which had substantially more data, this was quite slow

Right now I have 1000 rows of data, each row has about 100 columns, and then the column I want to expand has about 50 nested key/value pairs in it. I would expect that the data could scale up to 100k rows with the same number of columns over the next year or so, and so I'm hoping to have a scalable process ready to go at that point

pd.json_normalize should be faster than your attempt, but it is not faster than doing the flattening in pure python, so you might get more performance if you wrote a custom transform function & constructed the dataframe as below.

out = pd.DataFrame(transform(x) for x in df.to_dict(orient='records'))

Expand Pandas DataFrame Column with JSON Object

1 Answers