[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>

openbb_platform/core/openbb_core/app/utils/__init__.py

@@ -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])
+        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 "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.
-            parsed = df["date"].apply(to_datetime)
+            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,15 +111,36 @@ 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["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.
-    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"]):
-        if all(t.time() == time(0, 0) for t in df["date"]):
+        try:
-            df["date"] = df["date"].apply(lambda x: x.date().strftime("%Y-%m-%d"))
+            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"))
 
     return [
         Data(**d) for d in json.loads(df.to_json(orient="records", date_format="iso"))

openbb_platform/core/tests/app/test_utils.py

@@ -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")