diff --git a/bind/minipandas.py b/bind/minipandas.py index f1a5f5f..6f2ba18 100644 --- a/bind/minipandas.py +++ b/bind/minipandas.py @@ -144,19 +144,23 @@ class MiniFrame: def copy(self) -> "MiniFrame": return MiniFrame({k: v.copy() for k, v in self._data.items()}) - def to_pandas(self): + def to_pandas(self, indiv_cols: bool = True): import pandas as pd - df = pd.DataFrame() + out = {} for k, v in self._data.items(): if v.ndim == 1: - df[k] = v - elif v.ndim == 2: - # expand fixed-size vectors into wide columns + out[k] = v # 1D numeric column + elif v.ndim == 2 and indiv_cols: + # expand into wide columns for i in range(v.shape[1]): - df[f"{k}_{i}"] = v[:, i] + out[f"{k}_{i}"] = v[:, i] + elif v.ndim == 2 and not indiv_cols: + # legacy nested: one Python list per row (object dtype) + # NOTE: v.tolist() -> List[List[...]] which most legacy code expects + out[k] = pd.Series(v.tolist(), dtype=object) else: raise ValueError(f"Column '{k}' has ndim={v.ndim}, not supported.") - return df + return pd.DataFrame(out) # convenience constructor @staticmethod diff --git a/bind/test_minipandas.py b/bind/test_minipandas.py index 95657fd..e2f9ef3 100644 --- a/bind/test_minipandas.py +++ b/bind/test_minipandas.py @@ -49,4 +49,44 @@ print("\n=== Describe ===") desc = mf_all.describe() print(desc) -print("\nAll basic minipandas features appear to work correctly.") \ No newline at end of file +print("\nAll basic minipandas features appear to work correctly.") + + + +# --- Create a simple MiniFrame --- +N, M = 3, 4 +mf = mpd.MiniFrame({ + "run": 1, # scalar -> broadcast + "event": np.arange(N), # 1D scalar column + "energy": np.arange(N * M).reshape(N, M), # 2D vector column +}) + +# --- Case 1: indiv_cols=True (default) --- +df_wide = mf.to_pandas(indiv_cols=True) +print("=== Expanded columns ===") +print(df_wide) + +# Expected shape: N rows, 1 scalar column ('event'), 1 broadcasted column ('run'), and M expanded columns +assert list(df_wide.columns) == ["run", "event"] + [f"energy_{i}" for i in range(M)] +assert df_wide.shape == (N, 2 + M) +assert np.allclose(df_wide["energy_0"], [0, 4, 8]) + +# --- Case 2: indiv_cols=False (legacy nested format) --- +df_nested = mf.to_pandas(indiv_cols=False) +print("\n=== Nested legacy format ===") +print(df_nested) + +# Expected columns: run, event, energy +assert list(df_nested.columns) == ["run", "event", "energy"] +assert isinstance(df_nested.loc[0, "energy"], list) +assert df_nested.loc[0, "energy"] == [0, 1, 2, 3] + +# --- Round-trip compatibility check --- +mf2 = mpd.MiniFrame({ + "run": df_nested["run"].to_numpy(), + "event": df_nested["event"].to_numpy(), + "energy": np.stack(df_nested["energy"].to_numpy()), +}) + +assert np.allclose(mf["energy"], mf2["energy"]) +print("\n✅ All checks passed successfully!") \ No newline at end of file