ENH: add functional (but pretty slow) json_normalize function for flattening nested records per #1067

Author: wesm

pandas/io/json.py

@@ -1,6 +1,7 @@
-
 # pylint: disable-msg=E1101,W0613,W0603
-from StringIO import StringIO
+
+from collections import defaultdict
+
 import os
 
 from pandas import Series, DataFrame, to_datetime
@@ -11,17 +12,19 @@
 
 import numpy as np
 from pandas.tslib import iNaT
-import pandas.lib as lib
 
 ### interface to/from ###
 
-def to_json(path_or_buf, obj, orient=None, date_format='epoch', double_precision=10, force_ascii=True):
-        
+def to_json(path_or_buf, obj, orient=None, date_format='epoch',
+            double_precision=10, force_ascii=True):
+
     if isinstance(obj, Series):
-        s = SeriesWriter(obj, orient=orient, date_format=date_format, double_precision=double_precision, 
+        s = SeriesWriter(obj, orient=orient, date_format=date_format,
+                         double_precision=double_precision,
                          ensure_ascii=force_ascii).write()
     elif isinstance(obj, DataFrame):
-        s = FrameWriter(obj, orient=orient, date_format=date_format, double_precision=double_precision,
+        s = FrameWriter(obj, orient=orient, date_format=date_format,
+                        double_precision=double_precision,
                         ensure_ascii=force_ascii).write()
     else:
         raise NotImplementedError
@@ -36,12 +39,13 @@ def to_json(path_or_buf, obj, orient=None, date_format='epoch', double_precision
 
 class Writer(object):
 
-    def __init__(self, obj, orient, date_format, double_precision, ensure_ascii):
+    def __init__(self, obj, orient, date_format, double_precision,
+                 ensure_ascii):
         self.obj = obj
 
         if orient is None:
             orient = self._default_orient
-            
+
         self.orient = orient
         self.date_format = date_format
         self.double_precision = double_precision
@@ -64,15 +68,18 @@ def _format_to_date(self, data):
         if self._needs_to_date(data):
             return data.apply(lambda x: x.isoformat())
         return data
-    
+
     def copy_if_needed(self):
         """ copy myself if necessary """
         if not self.is_copy:
             self.obj = self.obj.copy()
             self.is_copy = True
 
     def write(self):
-        return dumps(self.obj, orient=self.orient, double_precision=self.double_precision, ensure_ascii=self.ensure_ascii)
+        return dumps(self.obj, orient=self.orient,
+                     double_precision=self.double_precision,
+                     ensure_ascii=self.ensure_ascii)
+
 
 class SeriesWriter(Writer):
     _default_orient = 'index'
@@ -186,13 +193,13 @@ def read_json(path_or_buf=None, orient=None, typ='frame', dtype=True,
     return obj
 
 class Parser(object):
-    
+
     def __init__(self, json, orient, dtype=True, convert_axes=True, convert_dates=True, keep_default_dates=False, numpy=False):
         self.json = json
 
         if orient is None:
             orient = self._default_orient
-            
+
         self.orient = orient
         self.dtype = dtype
 
@@ -207,7 +214,7 @@ def __init__(self, json, orient, dtype=True, convert_axes=True, convert_dates=Tr
 
     def parse(self):
 
-        # try numpy 
+        # try numpy
         numpy = self.numpy
         if numpy:
             self._parse_numpy()
@@ -269,7 +276,7 @@ def _try_convert_data(self, name, data, use_dtypes=True, convert_dates=True):
                 pass
 
         if data.dtype == 'float':
-            
+
             # coerce floats to 64
             try:
                 data = data.astype('float64')
@@ -291,7 +298,7 @@ def _try_convert_data(self, name, data, use_dtypes=True, convert_dates=True):
 
         # coerce ints to 64
         if data.dtype == 'int':
-            
+
             # coerce floats to 64
             try:
                 data = data.astype('int64')
@@ -322,7 +329,7 @@ def _try_convert_to_date(self, data):
         if issubclass(new_data.dtype.type,np.number):
             if not ((new_data == iNaT) | (new_data > 31536000000000000L)).all():
                 return data, False
-                
+
         try:
             new_data = to_datetime(new_data)
         except:
@@ -342,7 +349,7 @@ class SeriesParser(Parser):
     _default_orient = 'index'
 
     def _parse_no_numpy(self):
-    
+
         json = self.json
         orient = self.orient
         if orient == "split":
@@ -446,3 +453,197 @@ def is_ok(col):
                 new_data, result = self._try_convert_to_date(self.obj[col])
                 if result:
                     self.obj[col] = new_data
+
+
+#----------------------------------------------------------------------
+# JSON normalization routines
+
+def nested_to_record(ds,prefix="",level=0):
+    """a simplified json_normalize
+
+    converts a nested dict into a flat dict ("record"), unlike json_normalize,
+    it does not attempt to extract a subset of the data.
+
+    Parameters
+    ----------
+    ds : dict or list of dicts
+
+    Returns
+    -------
+    d - dict or list of dicts, matching `ds`
+
+    Example:
+    IN[52]: nested_to_record(dict(flat1=1,dict1=dict(c=1,d=2),nested=dict(e=dict(c=1,d=2),d=2)))
+    Out[52]:
+    {'dict1.c': 1,
+     'dict1.d': 2,
+     'flat1': 1,
+     'nested.d': 2,
+     'nested.e.c': 1,
+     'nested.e.d': 2}
+    """
+    singleton = False
+    if isinstance(ds,dict):
+        ds = [ds]
+        singleton = True
+
+    for d in ds:
+        for k,v in d.items(): # modifying keys inside loop, not lazy
+            # each key gets renamed with prefix
+            if level == 0:
+                newkey = str(k)
+            else:
+                newkey = prefix+'.'+ str(k)
+
+            # only dicts gets recurse-flattend
+            # only at level>1 do we rename the rest of the keys
+            if not isinstance(v,dict):
+                if level!=0: # so we skip copying for top level, common case
+                    v = d.pop(k)
+                    d[newkey]= v
+                continue
+            else:
+                v = d.pop(k)
+                d.update(nested_to_record(v,newkey,level+1))
+
+    if singleton:
+        return ds[0]
+    return ds
+
+
+def json_normalize(data, record_path=None, meta=None,
+                   meta_prefix=None,
+                   record_prefix=None):
+    """
+    "Normalize" semi-structured JSON data into a flat table
+
+    Parameters
+    ----------
+    data : dict or list of dicts
+        Unserialized JSON objects
+    record_path : string or list of strings, default None
+        Path in each object to list of records. If not passed, data will be
+        assumed to be an array of records
+    meta : list of paths (string or list of strings)
+        Fields to use as metadata for each record in resulting table
+    record_prefix : string, default None
+        If True, prefix records with dotted (?) path, e.g. foo.bar.field if
+        path to records is ['foo', 'bar']
+    meta_prefix : string, default None
+
+    Examples
+    --------
+    data = [{'state': 'Florida',
+             'shortname': 'FL',
+             'info': {
+                  'governor': 'Rick Scott'
+             },
+             'counties': [{'name': 'Dade', 'population': 12345},
+                         {'name': 'Broward', 'population': 40000},
+                         {'name': 'Palm Beach', 'population': 60000}]},
+            {'state': 'Ohio',
+             'shortname': 'OH',
+             'info': {
+                  'governor': 'John Kasich'
+             },
+             'counties': [{'name': 'Summit', 'population': 1234},
+                          {'name': 'Cuyahoga', 'population': 1337}]}]
+
+    result = json_normalize(data, 'counties', ['state', 'shortname',
+                                              ['info', 'governor']])
+
+      state    governor
+    Florida  Rick Scott
+
+
+    Returns
+    -------
+    frame : DataFrame
+    """
+    def _pull_field(js, spec):
+        result = js
+        if isinstance(spec, list):
+            for field in spec:
+                result = result[field]
+        else:
+            result = result[spec]
+
+        return result
+
+    # A bit of a hackjob
+    if isinstance(data, dict):
+        data = [data]
+
+    if record_path is None:
+        if any([isinstance(x,dict) for x in data[0].itervalues()]):
+            # naive normalization, this is idempotent for flat records
+            # and potentially will inflate the data considerably for
+            # deeply nested structures:
+            #  {VeryLong: { b: 1,c:2}} -> {VeryLong.b:1 ,VeryLong.c:@}
+            #
+            # TODO: handle record value which are lists, at least error reasonabley
+            data = nested_to_record(data)
+        return DataFrame(data)
+    elif not isinstance(record_path, list):
+        record_path = [record_path]
+
+    if meta is None:
+        meta = []
+    elif not isinstance(meta, list):
+        meta = [meta]
+
+    for i, x in enumerate(meta):
+        if not isinstance(x, list):
+            meta[i] = [x]
+
+    # Disastrously inefficient for now
+    records = []
+    lengths = []
+
+    meta_vals = defaultdict(list)
+    meta_keys = ['.'.join(val) for val in meta]
+
+    def _recursive_extract(data, path, seen_meta, level=0):
+        if len(path) > 1:
+            for obj in data:
+                for val, key in zip(meta, meta_keys):
+                    if level + 1 == len(val):
+                        seen_meta[key] = _pull_field(obj, val[-1])
+
+                _recursive_extract(obj[path[0]], path[1:],
+                                   seen_meta, level=level+1)
+        else:
+            for obj in data:
+                recs = _pull_field(obj, path[0])
+
+                # For repeating the metadata later
+                lengths.append(len(recs))
+
+                for val, key in zip(meta, meta_keys):
+                    if level + 1 > len(val):
+                        meta_val = seen_meta[key]
+                    else:
+                        meta_val = _pull_field(obj, val[level:])
+                    meta_vals[key].append(meta_val)
+
+                records.extend(recs)
+
+    _recursive_extract(data, record_path, {}, level=0)
+
+    result = DataFrame(records)
+
+    if record_prefix is not None:
+        result.rename(columns=lambda x: record_prefix + x, inplace=True)
+
+    # Data types, a problem
+    for k, v in meta_vals.iteritems():
+        if meta_prefix is not None:
+            k = meta_prefix + k
+
+        if k in result:
+            raise ValueError('Conflicting metadata name %s, '
+                             'need distinguishing prefix ' % k)
+
+        result[k] = np.array(v).repeat(lengths)
+
+    return result