mirror of
https://github.com/OpenBB-finance/OpenBB.git
synced 2026-10-27 04:31:13 +00:00
[BugFix] #7672-Preserve unnamed MultiIndex levels in DataFrame round-trip. (#7762)
* #7672-Fixed DataFrame round-trip issue * Address review feedback on MultiIndex round trip (#7672) --------- Co-authored-by: dlee-bq <dlee@brightquery.com>
This commit is contained in:
@@ -17,6 +17,13 @@ if TYPE_CHECKING:
|
|||||||
from numpy import ndarray
|
from numpy import ndarray
|
||||||
from pandas import DataFrame, Series
|
from pandas import DataFrame, Series
|
||||||
|
|
||||||
|
_MULTIINDEX_META = ("is_multiindex", "multiindex_names", "multiindex_columns")
|
||||||
|
|
||||||
|
|
||||||
|
def _multiindex_columns(names: list) -> list:
|
||||||
|
"""Get the column names that reset_index gives the levels of a MultiIndex."""
|
||||||
|
return [n if n is not None else f"level_{i}" for i, n in enumerate(names)]
|
||||||
|
|
||||||
|
|
||||||
def basemodel_to_df(
|
def basemodel_to_df(
|
||||||
data: list[Data] | Data,
|
data: list[Data] | Data,
|
||||||
@@ -25,6 +32,7 @@ def basemodel_to_df(
|
|||||||
"""Convert list of BaseModel to a Pandas DataFrame."""
|
"""Convert list of BaseModel to a Pandas DataFrame."""
|
||||||
pd = require_optional("pandas")
|
pd = require_optional("pandas")
|
||||||
DataFrame, to_datetime = pd.DataFrame, pd.to_datetime # type: ignore[union-attr]
|
DataFrame, to_datetime = pd.DataFrame, pd.to_datetime # type: ignore[union-attr]
|
||||||
|
is_numeric_dtype = pd.api.types.is_numeric_dtype # type: ignore[union-attr]
|
||||||
|
|
||||||
if isinstance(data, list):
|
if isinstance(data, list):
|
||||||
df = DataFrame(
|
df = DataFrame(
|
||||||
@@ -40,12 +48,20 @@ def basemodel_to_df(
|
|||||||
)
|
)
|
||||||
|
|
||||||
if "is_multiindex" in df.columns:
|
if "is_multiindex" in df.columns:
|
||||||
col_names = ast.literal_eval(df.multiindex_names.unique()[0])
|
names = ast.literal_eval(df.multiindex_names.unique()[0])
|
||||||
df = df.set_index(col_names)
|
# The level columns follow the reset_index naming unless df_to_basemodel
|
||||||
df = df.drop(["is_multiindex", "multiindex_names"], axis=1)
|
# had to rename them to avoid a clash, in which case it stores them.
|
||||||
|
if "multiindex_columns" in df.columns:
|
||||||
|
columns = ast.literal_eval(df.multiindex_columns.unique()[0])
|
||||||
|
else:
|
||||||
|
columns = _multiindex_columns(names)
|
||||||
|
df = df.set_index(columns).rename_axis(names)
|
||||||
|
df = df.drop(columns=[c for c in _MULTIINDEX_META if c in df.columns])
|
||||||
|
|
||||||
# If the date column contains dates only, convert them to a date to avoid encoding time data.
|
# If the date column contains dates only, convert them to a date to avoid encoding time data.
|
||||||
if "date" in df.columns:
|
# Numeric values would be read as epoch offsets, so only text and datetimes are parsed.
|
||||||
|
parsed = None
|
||||||
|
if "date" in df.columns and not is_numeric_dtype(df["date"]):
|
||||||
try:
|
try:
|
||||||
# Parse the column in a single call instead of once per row.
|
# Parse the column in a single call instead of once per row.
|
||||||
parsed = to_datetime(df["date"])
|
parsed = to_datetime(df["date"])
|
||||||
@@ -53,8 +69,13 @@ def basemodel_to_df(
|
|||||||
# A column that mixes UTC offsets - a series crossing a daylight
|
# A column that mixes UTC offsets - a series crossing a daylight
|
||||||
# savings boundary, for instance - has no single datetime64
|
# savings boundary, for instance - has no single datetime64
|
||||||
# representation, so fall back to parsing each value on its own.
|
# representation, so fall back to parsing each value on its own.
|
||||||
|
try:
|
||||||
parsed = df["date"].apply(to_datetime)
|
parsed = df["date"].apply(to_datetime)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
# Not a date column after all, so leave it as it is.
|
||||||
|
parsed = None
|
||||||
|
|
||||||
|
if parsed is not None:
|
||||||
df["date"] = parsed
|
df["date"] = parsed
|
||||||
|
|
||||||
if parsed.dtype == object:
|
if parsed.dtype == object:
|
||||||
@@ -79,6 +100,7 @@ def df_to_basemodel(
|
|||||||
"""Convert from a Pandas DataFrame to list of BaseModel."""
|
"""Convert from a Pandas DataFrame to list of BaseModel."""
|
||||||
pd = require_optional("pandas")
|
pd = require_optional("pandas")
|
||||||
MultiIndex, Series, to_datetime = pd.MultiIndex, pd.Series, pd.to_datetime # type: ignore[union-attr]
|
MultiIndex, Series, to_datetime = pd.MultiIndex, pd.Series, pd.to_datetime # type: ignore[union-attr]
|
||||||
|
is_numeric_dtype = pd.api.types.is_numeric_dtype # type: ignore[union-attr]
|
||||||
|
|
||||||
is_multiindex = isinstance(df.index, MultiIndex)
|
is_multiindex = isinstance(df.index, MultiIndex)
|
||||||
|
|
||||||
@@ -89,13 +111,34 @@ def df_to_basemodel(
|
|||||||
|
|
||||||
# Check if df has multiindex. If so, add the index names to the df and a boolean column
|
# Check if df has multiindex. If so, add the index names to the df and a boolean column
|
||||||
if isinstance(df.index, MultiIndex):
|
if isinstance(df.index, MultiIndex):
|
||||||
|
names = list(df.index.names)
|
||||||
|
default_columns = _multiindex_columns(names)
|
||||||
|
# A level whose column name is already taken, by a column or another level,
|
||||||
|
# gets a free name instead so that reset_index does not fail.
|
||||||
|
taken = set(df.columns) | set(_MULTIINDEX_META)
|
||||||
|
columns = []
|
||||||
|
for default in default_columns:
|
||||||
|
column = default
|
||||||
|
while column in taken:
|
||||||
|
column = f"_{column}"
|
||||||
|
taken.add(column)
|
||||||
|
columns.append(column)
|
||||||
|
df = df.rename_axis(columns).reset_index()
|
||||||
df["is_multiindex"] = True
|
df["is_multiindex"] = True
|
||||||
df["multiindex_names"] = str(df.index.names)
|
df["multiindex_names"] = str(names)
|
||||||
df = df.reset_index()
|
if columns != default_columns:
|
||||||
|
df["multiindex_columns"] = str(columns)
|
||||||
|
|
||||||
# Converting to JSON will add T00:00:00.000 to all dates with no time element unless we format it as a string first.
|
# Converting to JSON will add T00:00:00.000 to all dates with no time element unless we format it as a string first.
|
||||||
if "date" in df.columns:
|
# Numeric values would be read as epoch offsets, so only text and datetimes are parsed.
|
||||||
df["date"] = df["date"].apply(to_datetime)
|
if "date" in df.columns and not is_numeric_dtype(df["date"]):
|
||||||
|
try:
|
||||||
|
dates = df["date"].apply(to_datetime)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
# Not a date column after all, so leave it as it is.
|
||||||
|
dates = None
|
||||||
|
if dates is not None:
|
||||||
|
df["date"] = dates
|
||||||
if all(t.time() == time(0, 0) for t in df["date"]):
|
if all(t.time() == time(0, 0) for t in df["date"]):
|
||||||
df["date"] = df["date"].apply(lambda x: x.date().strftime("%Y-%m-%d"))
|
df["date"] = df["date"].apply(lambda x: x.date().strftime("%Y-%m-%d"))
|
||||||
|
|
||||||
|
|||||||
@@ -75,6 +75,59 @@ def test_basemodel_to_multiindex_df():
|
|||||||
assert isinstance(df.index, pd.MultiIndex)
|
assert isinstance(df.index, pd.MultiIndex)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"names",
|
||||||
|
[["x", "y"], [None, None], ["x", None], [None, "y"]],
|
||||||
|
)
|
||||||
|
def test_multiindex_round_trip(names):
|
||||||
|
"""Test that a MultiIndex survives a df -> basemodel -> df round trip."""
|
||||||
|
original = df_multiindex.copy()
|
||||||
|
original.index.names = names
|
||||||
|
|
||||||
|
result = basemodel_to_df(df_to_basemodel(original))
|
||||||
|
|
||||||
|
pd.testing.assert_frame_equal(result, original)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"names, columns",
|
||||||
|
[
|
||||||
|
(["z", "y"], ["z"]),
|
||||||
|
([None, "y"], ["level_0"]),
|
||||||
|
(["x", "x"], ["z"]),
|
||||||
|
(["is_multiindex", "y"], ["z"]),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_multiindex_round_trip_name_clash(names, columns):
|
||||||
|
"""Test that index levels named after a column or each other survive a round trip."""
|
||||||
|
original = df_multiindex.copy()
|
||||||
|
original.index.names = names
|
||||||
|
original.columns = columns
|
||||||
|
|
||||||
|
result = basemodel_to_df(df_to_basemodel(original))
|
||||||
|
|
||||||
|
pd.testing.assert_frame_equal(result, original)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("values", [list("abcdefghij"), list(range(10))])
|
||||||
|
def test_multiindex_round_trip_non_date_date_level(values):
|
||||||
|
"""Test that an index level named date that holds non-date values survives a round trip."""
|
||||||
|
original = df.assign(date=values).set_index(["date", "y"])
|
||||||
|
|
||||||
|
result = basemodel_to_df(df_to_basemodel(original))
|
||||||
|
|
||||||
|
pd.testing.assert_frame_equal(result, original)
|
||||||
|
|
||||||
|
|
||||||
|
def test_df_to_basemodel_does_not_modify_multiindex_input():
|
||||||
|
"""Test that converting a MultiIndex DataFrame leaves the input untouched."""
|
||||||
|
original = df.set_index(["x", "y"])
|
||||||
|
|
||||||
|
df_to_basemodel(original)
|
||||||
|
|
||||||
|
pd.testing.assert_frame_equal(original, df.set_index(["x", "y"]))
|
||||||
|
|
||||||
|
|
||||||
def test_get_target_column():
|
def test_get_target_column():
|
||||||
"""Test the get_target_column helper."""
|
"""Test the get_target_column helper."""
|
||||||
target = get_target_column(df, "x")
|
target = get_target_column(df, "x")
|
||||||
|
|||||||
Reference in New Issue
Block a user