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 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(
|
||||
data: list[Data] | Data,
|
||||
@@ -25,6 +32,7 @@ def basemodel_to_df(
|
||||
"""Convert list of BaseModel to a Pandas DataFrame."""
|
||||
pd = require_optional("pandas")
|
||||
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):
|
||||
df = DataFrame(
|
||||
@@ -40,12 +48,20 @@ def basemodel_to_df(
|
||||
)
|
||||
|
||||
if "is_multiindex" in df.columns:
|
||||
col_names = ast.literal_eval(df.multiindex_names.unique()[0])
|
||||
df = df.set_index(col_names)
|
||||
df = df.drop(["is_multiindex", "multiindex_names"], axis=1)
|
||||
names = ast.literal_eval(df.multiindex_names.unique()[0])
|
||||
# The level columns follow the reset_index naming unless df_to_basemodel
|
||||
# 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 "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:
|
||||
# Parse the column in a single call instead of once per row.
|
||||
parsed = to_datetime(df["date"])
|
||||
@@ -53,8 +69,13 @@ def basemodel_to_df(
|
||||
# A column that mixes UTC offsets - a series crossing a daylight
|
||||
# savings boundary, for instance - has no single datetime64
|
||||
# representation, so fall back to parsing each value on its own.
|
||||
try:
|
||||
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
|
||||
|
||||
if parsed.dtype == object:
|
||||
@@ -79,6 +100,7 @@ def df_to_basemodel(
|
||||
"""Convert from a Pandas DataFrame to list of BaseModel."""
|
||||
pd = require_optional("pandas")
|
||||
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)
|
||||
|
||||
@@ -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
|
||||
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["multiindex_names"] = str(df.index.names)
|
||||
df = df.reset_index()
|
||||
df["multiindex_names"] = str(names)
|
||||
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.
|
||||
if "date" in df.columns:
|
||||
df["date"] = df["date"].apply(to_datetime)
|
||||
# Numeric values would be read as epoch offsets, so only text and datetimes are parsed.
|
||||
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"]):
|
||||
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)
|
||||
|
||||
|
||||
@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():
|
||||
"""Test the get_target_column helper."""
|
||||
target = get_target_column(df, "x")
|
||||
|
||||
Reference in New Issue
Block a user