Copyright (C) 2019-2022 HERE Global B.V. and its affiliate(s).

All rights reserved.

This software and other materials contain proprietary information

controlled by HERE and are protected by applicable copyright legislation.

Any use and utilization of this software and other materials and

disclosure to any third parties is conditional upon having a separate

agreement with HERE for the access, use, utilization or disclosure of this

software. In the absence of such agreement, the use of the software is not

allowed.

"""
Utilities to manipulate the structure of pandas DataFrame and Series
"""
from typing import List, Optional, Union, cast

import pandas as pd
from here.platform.utils.collection import flatten_iterator

The separator to use when prefixing, unprefixing, unpacking columns

_sep = "."

[docs]
def unpack(series: pd.Series, keep_prefix: bool = True, max_level: int = -1) -> pd.DataFrame:
"""
Unpack a Series that contains dictionaries into a DataFrame that has one
column for each of the field of the dictionaries found in the series.

In the process, nested dictionaries are also unpacked: if a field of a dictionary
contains a nested dictionary, its fields are added as columns as well, with a name composed
by both the names of the field and nested field, separated by ..

Unpacking is recursive until nested dictionaries are found, or a maximum level is reached.
For more details about the recursive unpacking process, please see
here.platform.utils.collection_utils.flatten_iterator.

The resulting columns, their order and their types are function of the data.
The index of the series is retained. None, pd.NA and other values that are not
dictionaries, including lists and scalar values, are discarded.

:param series: the series containing dictionaries to unpack
:param keep_prefix: keep the name of the series as prefix for the unpacked columns
:param max_level: the maximum level of the recursive unpacking. 0 performs no recursive
unpacking, 1 unpacks only the first nested dictionaries, 2 only the first and second
nested dictionaries, and so on. A negative number represents no limit.
:return: a DataFrame, with one column for each of the fields of the dictionaries
"""

Flatten nested dictionaries

prefix = cast(str, series.name) if keep_prefix else ""
flattened = series.apply(
lambda d: dict(
flatten_iterator(iter(d.items()), max_level=max_level, prefix=prefix, sep=_sep)
)
if isinstance(d, dict)
else None
)

All the flattened field names, without duplicates

dict maintains the element order, set doesn't, therefore the unused None

names = {n: None for _, d in flattened.items() for n in (list(d.keys()) if isinstance(d, dict) else [])}

Construct the resulting dataframe, one column per flattened field name

result = pd.DataFrame(
{n: flattened.apply(lambda d: d.get(n) if isinstance(d, dict) else None) for n in names},
index=flattened.index,
)
assert result.index is series.index
return result

[docs]
def replace_column(
dataframe: pd.DataFrame, column: str, new_columns: pd.DataFrame
) -> pd.DataFrame:
"""
Replace one column of a DataFrame with the all columns of another DataFrame.

The selected column is removed and the new columns are inserted in its place.
Indices are aligned, but only rows already present in the input dataframe
are retained: rows present in the new columns but not in the original
dataframe are discarded. The input dataframe is not altered.

:param dataframe: the input dataframe
:param column: the name of the column of the input dataframe to replace
:param new_columns: the DataFrame to replace the selected column with
:return: a new DataFrame with the column replaced
"""
idx = dataframe.columns.get_loc(column)
result = dataframe.iloc[:, :idx]
remaining = dataframe.iloc[:, (idx + 1) :]
for i, column in enumerate(new_columns.columns, start=len(result.columns)):
result.insert(i, column, new_columns[column])
for i, column in enumerate(remaining.columns, start=len(result.columns)):
result.insert(i, column, remaining[column])

insert always retains the index, and this is good, see however:

https://github.com/pandas-dev/pandas/issues/43718

assert result.index is dataframe.index
return result

[docs]
def unpack_columns(
dataframe: pd.DataFrame,
columns: Union[str, List[str]],
keep_columns: bool = False,
keep_prefix: bool = True,
max_level: int = -1,
) -> pd.DataFrame:
"""
Unpack one or more columns of a DataFrame, replacing them with columns
extracted from the fields of the dictionaries they contain.

Similarly to how pandas explode function can unroll a list to multiple
rows of a DataFrame, this function is useful to unpack columns containing dict
into constructs that are easier to manipulate with pandas.

The input dataframe is not altered and its index retained.

Unpacking is recursive, optionally down to a certain maximum level.
See unpack of a single Series for details of the unpacking algorithm.

:param dataframe: the DataFrame to unpack
:param columns: one or more names of columns to unpack
:param keep_columns: keep the original unpacked columns in the resulting dataframe,
retaining any single values but discarding dictionaries that have been unpacked
:param keep_prefix: keep the name of the columns as prefix for the unpacked columns
:param max_level: the maximum level of the recursive unpacking. 0 performs no recursive
unpacking, 1 unpacks only the first nested dictionaries, 2 only the first and second
nested dictionaries, and so on. A negative number represents no limit.
A single max_level is used for all the columns. To unpack columns, each with
a different maximum unpacking level, call this function more than once.
:return: the input dataframe with the selected columns unpacked.
"""
result = dataframe
for column in [columns] if isinstance(columns, str) else columns:
to_replace = unpack(result[column], keep_prefix=keep_prefix, max_level=max_level)
if keep_columns:
cleaned_up = result[column].apply(lambda x: pd.NA if isinstance(x, dict) else x)
to_replace.insert(0, column, cleaned_up)
result = replace_column(result, column, to_replace)

Index must be retained

assert dataframe.index is result.index
return result

[docs]
def prefix_columns(
data: pd.DataFrame, prefix: str, columns: Optional[List[str]] = None
) -> pd.DataFrame:
"""
Rename all or selected columns by adding a prefix to their name.
The prefix is prepended to the column names using . as separator.

For example, when applied with prefix my_attr, the function renames
columns a, b and c to my_attr.a, my_attr.b and my_attr.c.

:param data: the input dataframe
:param prefix: the prefix to prepend to the selected columns
:param columns: names of column to nest under the prefix. The operation
is applied to all the columns of the input dataframe if not specified.
:return: a dataframe with the selected columns nested under a prefix
:raises ValueError: in case the prefix is invalid
"""
if not prefix:
raise ValueError("Prefix is not specified")

to_rename = {c: (prefix + _sep + c) for c in data.columns if (not columns or c in columns)}

return data.rename(columns=to_rename)

[docs]
def unprefix_columns(data: pd.DataFrame, prefix: str) -> pd.DataFrame:
"""
Rename all the columns that start with a prefix by removing it.
The . separator after the prefix is checked and removed as well.

For example, when applied with prefix my_attr, the function renames
columns my_attr.a, my_attr.b and my_attr.c to a, b and c.

:param data: the input dataframe
:param prefix: the prefix to remove from the columns
:return: a dataframe with the given prefix removed from the column names
:raises ValueError: in case the prefix is invalid
"""
if not prefix:
raise ValueError("Prefix is not specified")
prefix = prefix + _sep
prefix_len = len(prefix)

to_rename = {c: c[prefix_len:] for c in data.columns if c.startswith(prefix) and len(c) > prefix_len}

return data.rename(columns=to_rename)