diff --git a/changelog_entry.yaml b/changelog_entry.yaml index e69de29..f7834a4 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -0,0 +1,4 @@ +- bump: patch + changes: + fixed: + - Allowed a MicroDataFrame to handle an empty index subset diff --git a/docs/gini.ipynb b/docs/gini.ipynb index 7b5bffd..9861c36 100644 --- a/docs/gini.ipynb +++ b/docs/gini.ipynb @@ -27,7 +27,7 @@ "# Create sample data\n", "x = [10, 20, 30, 40, 100]\n", "w = [1, 2, 3, 4, 5]\n", - "df = pd.DataFrame({'x': x, 'w': w})" + "df = pd.DataFrame({\"x\": x, \"w\": w})" ] }, { @@ -80,7 +80,7 @@ "outputs": [], "source": [ "# Create a MicroDataFrame\n", - "mdf_df = mdf.MicroDataFrame(df, weights='w')\n", + "mdf_df = mdf.MicroDataFrame(df, weights=\"w\")\n", "\n", "# Access column as MicroSeries and calculate gini\n", "print(f\"Gini from MicroDataFrame column: {mdf_df.x.gini():.4f}\")" diff --git a/microdf/microdataframe.py b/microdf/microdataframe.py index e749c8a..f43ee27 100644 --- a/microdf/microdataframe.py +++ b/microdf/microdataframe.py @@ -218,13 +218,15 @@ def set_weight_col( def __getitem__( self, key: Union[str, List] ) -> Union[pd.Series, pd.DataFrame]: + # Let pandas handle the initial slicing result = super().__getitem__(key) + + # If the result is a DataFrame, re-synchronize the weights if isinstance(result, pd.DataFrame): - try: - weights = self.weights[key] - except Exception: - weights = self.weights - return MicroDataFrame(result, weights=weights) + new_weights = self.weights.reindex(result.index) + return MicroDataFrame(result, weights=new_weights) + + # Otherwise, the result is a Series or a scalar, so just return it return result def catch_series_relapse(self) -> None: diff --git a/microdf/tests/test_microseries_dataframe.py b/microdf/tests/test_microseries_dataframe.py index 6a22d83..de75ab4 100644 --- a/microdf/tests/test_microseries_dataframe.py +++ b/microdf/tests/test_microseries_dataframe.py @@ -24,6 +24,15 @@ def test_df_init() -> None: assert df.a.mean() == np.average(arr, weights=w) +def test_handles_empty_index() -> None: + arr = np.array([0, 1, 1]) + w = np.array([3, 0, 9]) + df = mdf.MicroDataFrame({"a": arr}, weights=w) + + empty_index = pd.Index([]) + df[empty_index] # Implicit assert; checking for ValueError + + def test_series_getitem() -> None: arr = np.array([0, 1, 1]) w = np.array([3, 0, 9])