Permute levels of nested mappings

[thorwhalen]

The problem of needing to flip the levels of nested mappings popped up several times.

If we were working with tables we would have identified the problem earlier on: It's related the problem of groupby, also related to sort by, where we want to pick one or several keys to group (or sort) the rows by. Rather, our problem here would be one of first ungrouping some grouped data, then regrouping according to different grouping rules.

Nested mappings with consistent depth (i.e. the depth of the leafs are all at the same level) could be seen as a representation of a (sparse) tensor (i.e. a multi-dimensional matrix).

Let's get to the concrete stuff.

The simple_flip_first_two_levels function below flips the first two levels of a nested mapping.

def flip_first_two_levels(mapping):
    """Flips the keys and values of a nested dictionary."""
    result = {}
    for key1, sub_mapping in mapping.items():
        for key2, value in sub_mapping.items():
            if key2 not in result:
                result[key2] = {}
            result[key2][key1] = value
    return result

See it in action:

        
depth_2_mapping = {
    'a': {'x': 1, 'y': 2},
    'b': {'x': 3, 'z': 4},
    'c': {'y': 5, 'z': 6}
}

assert flip_first_two_levels(depth_2_mapping) == {
    'x': {'a': 1, 'b': 3}, 
    'y': {'a': 2, 'c': 5}, 
    'z': {'b': 4, 'c': 6}
}

But the function is limited in many ways:

• It only works with two levels, we would like to be able to have more control over the restructuring
• It works on any input typing.Mapping, but returns a dict
• It only works if the nested depth is consistent

Note what happens when we use the function on a dictionary of depth 3:

depth_3_mapping = {
    'a': {'x': {'foo': 1, 'bar': 2}, 'y': {'foo': 3, 'baz': 4}},
    'b': {'x': {'bar': 5, 'baz': 6}, 'z': {'foo': 7, 'bar': 8}}
}

assert flip_first_two_levels(depth_3_mapping) == {
    'x': {
        'a': {'foo': 1, 'bar': 2}, 
        'b': {'bar': 5, 'baz': 6}
    },
    'y': {'a': {'foo': 3, 'baz': 4}},
    'z': {'b': {'foo': 7, 'bar': 8}}
}

Example of permuting levels

But what if we wanted to restructure it as, say:

{
    "x": {
        "foo": {"a": 1}, 
        "bar": {"a": 2, "b": 5}, 
        "baz": {"b": 6}
    },
    "y": {
        "foo": {"a": 3}
    },
    "z": {
        "foo": {"b": 7}, 
        "bar": {"b": 8}
    },
}

And then, there's the fact that no matter what kind of mapping we input, we'll always get a dict as our output, which may not be what we want -- sure, we could always cast it to the mapping we want from after getting the dict, but that could be undesirable if our mapping is an interface to significant amounts of data.

We propose to write a flip_levels (permute_levels? renest_mapping? any name suggestions?) that would look something like

Proposal

from typing import Mapping, Sequence, Callable

MappingCast = Callable[[Mapping], Mapping]

def flip_levels(
    mapping: Mapping, 
    permutation: Sequence[int],  # one-line representation of a permutation (with zero-based indexing)
    casts: Sequence[MappingCast]  # specifies how to cast each level of the output mapping
):
    """Permutes the nested levels of the input mapping"""

See One line representation of a permutation, but of course we could handle other representations.

Examples

For example, our flip_first_two_levels would basically be functools.partial(flip_levels, permutation=(1, 0), casts=[dict] * 2.

assert (
    flip_first_two_levels(mapping) 
    == 
    flip_levels(mapping, permutation=(1, 0), casts=[dict] * 2
)

Code ideas we can use

In the following code, we use FlatReader to flatten the input mapping and cast_mapping to carry out the mapping casting.

Note though that FlatReader only flattens two levels.
There's a discussion about that.

It seems like a good idea to implicitly or explicitly restructure through flattening.

In fact, we should think about wording and design. Might it make more sense to make group and ungroup functions (restructuring/permutation being carried out by ungrouping and grouping differently)?

from collections import defaultdict
from dol.sources import FlatReader

tree = lambda: defaultdict(tree)

# TODO: Generalize to use permutation of levels
def flip_first_two_levels(mapping):
    d = tree()
    for (k1, k2), v in FlatReader(mapping).items():
        d[k2][k1] = v
    return cast_mapping(d, [dict] * 2)
                
    
def cast_mapping(mapping, casts):
    """Casts a mapping recursively, with the ability to use different casting functions
    at different levels"""
    cast, *_remaining_casts = casts
    if not _remaining_casts:
        return cast(mapping)
    else:
        def _gen_items():
            for k, v in mapping.items():
                yield k, cast_mapping(v, _remaining_casts)
        return cast(_gen_items())