diff --git a/leanframe/core/frame.py b/leanframe/core/frame.py index 08393e9..063d9a7 100644 --- a/leanframe/core/frame.py +++ b/leanframe/core/frame.py @@ -222,6 +222,79 @@ def to_ibis(self) -> ibis_types.Table: """Return the underlying Ibis expression.""" return self._data + def drop( + self, + labels=None, + *, + axis=0, + index=None, + columns=None, + level=None, + inplace: bool = False, + errors: str = "ignore", + ) -> DataFrame: + """Drop specified labels from columns or index. + + Unlike pandas, the default for errors is "ignore". + Using errors="raise" is unsupported in leanframe. + """ + if inplace: + raise NotImplementedError("inplace=True is not supported in leanframe.") + + if level is not None: + raise NotImplementedError("level is not supported in leanframe.") + + if errors == "raise": + raise ValueError("errors='raise' is not supported in leanframe.") + + if labels is not None: + if index is not None or columns is not None: + raise ValueError("Cannot specify both 'labels' and 'index'/'columns'") + + if axis in (0, "index"): + index = labels + elif axis in (1, "columns"): + columns = labels + else: + raise ValueError(f"No axis named {axis} for object type DataFrame") + + if labels is None and columns is None and index is None: + raise ValueError( + "Need to specify at least one of 'labels', 'index' or 'columns'" + ) + + new_data = self._data + + if index is not None: + if not hasattr(self, "_index") or self._index is None: + raise ValueError( + "Cannot drop rows by index labels without an index set. " + "Use .set_index('column_name') first." + ) + + if isinstance(index, str) or not hasattr(index, "__iter__"): + idx_to_drop = [index] + else: + idx_to_drop = list(index) + + if len(idx_to_drop) > 0: + index_col = new_data[self._index.columns[0]] + new_data = new_data.filter(~index_col.isin(idx_to_drop)) + + if columns is not None: + if isinstance(columns, str) or not hasattr(columns, "__iter__"): + cols_to_drop = [columns] + else: + cols_to_drop = list(columns) + + existing_cols = new_data.columns + cols_to_keep = [col for col in existing_cols if col not in cols_to_drop] + new_data = new_data.select(*cols_to_keep) + + res = DataFrame(new_data) + res._index = self._index + return res + def set_index( self, columns: str | list[str], diff --git a/tasks/001-dataframe.md b/tasks/001-dataframe.md index f052ea0..b5edec6 100644 --- a/tasks/001-dataframe.md +++ b/tasks/001-dataframe.md @@ -55,7 +55,7 @@ Implement all methods and properties for the pandas DataFrame class. - [ ] `div` - [ ] `divide` - [ ] `dot` -- [ ] `drop` +- [x] `drop` - [ ] `drop_duplicates` - [x] `droplevel` - [ ] `dropna` diff --git a/tests/unit/test_frame.py b/tests/unit/test_frame.py index cdaf4fb..98cc7ef 100644 --- a/tests/unit/test_frame.py +++ b/tests/unit/test_frame.py @@ -17,6 +17,7 @@ import pandas as pd import pandas.testing as tm import pyarrow as pa +import pytest import leanframe @@ -185,3 +186,91 @@ def test_dataframe_assign_overwrite(session: leanframe.Session): result_lf = df_lf.assign(col1=session.col("col1") * 2) expected_pd = df_pd.assign(col1=df_pd["col1"] * 2) tm.assert_frame_equal(result_lf.to_pandas(), expected_pd) + + +def test_dataframe_drop_columns(session: leanframe.Session): + df_pd = pd.DataFrame({ + "col1": [1, 2, 3], + "col2": ["a", "b", "c"], + "col3": [1.1, 2.2, 3.3], + }).astype({ + "col1": pd.ArrowDtype(pa.int64()), + "col2": pd.ArrowDtype(pa.string()), + "col3": pd.ArrowDtype(pa.float64()), + }) + df_lf = session.DataFrame(df_pd) + + # Drop single column + result1 = df_lf.drop(columns="col2") + expected1 = df_pd.drop(columns="col2") + tm.assert_frame_equal(result1.to_pandas(), expected1) + + # Drop multiple columns + result2 = df_lf.drop(columns=["col1", "col3"]) + expected2 = df_pd.drop(columns=["col1", "col3"]) + tm.assert_frame_equal(result2.to_pandas(), expected2) + + # Drop using labels and axis=1 + result3 = df_lf.drop(["col2"], axis=1) + expected3 = df_pd.drop(["col2"], axis=1) + tm.assert_frame_equal(result3.to_pandas(), expected3) + +def test_dataframe_drop_errors(session: leanframe.Session): + df_pd = pd.DataFrame({ + "col1": [1, 2, 3], + "col2": ["a", "b", "c"], + }) + df_lf = session.DataFrame(df_pd) + + # errors="raise" is unsupported + with pytest.raises(ValueError, match="errors='raise' is not supported"): + df_lf.drop(columns="missing", errors="raise") + + # Test errors='ignore' (default behavior now) + result = df_lf.drop(columns=["col1", "missing"]) + # Need to handle pandas type conversion for assert + tm.assert_frame_equal(result.to_pandas(), df_lf.to_pandas().drop(columns=["col1", "missing"], errors="ignore")) + +def test_dataframe_drop_unsupported(session: leanframe.Session): + df_pd = pd.DataFrame({ + "col1": [1, 2, 3], + }) + df_lf = session.DataFrame(df_pd) + + with pytest.raises(NotImplementedError, match="inplace=True is not supported"): + df_lf.drop(columns="col1", inplace=True) + + with pytest.raises(NotImplementedError, match="level is not supported"): + df_lf.drop(columns="col1", level=1) + + with pytest.raises(ValueError, match="Cannot drop rows by index labels without an index set"): + df_lf.drop(index=[0]) + + with pytest.raises(ValueError, match="Cannot drop rows by index labels without an index set"): + df_lf.drop(labels=[0], axis=0) + + with pytest.raises(ValueError, match="Need to specify at least one of 'labels', 'index' or 'columns'"): + df_lf.drop() + +def test_dataframe_drop_index(session: leanframe.Session): + df_pd = pd.DataFrame({ + "id": [1, 2, 3, 4], + "col2": ["a", "b", "c", "d"], + }) + df_lf = session.DataFrame(df_pd) + + # Must set index first + df_lf = df_lf.set_index("id") + df_pd = df_pd.set_index("id") + + # Drop single index value + result1 = df_lf.drop(index=2) + expected1 = df_pd.drop(index=2) + + # We use check_dtype=False because Ibis output column order or exact type might vary slightly + tm.assert_frame_equal(result1.to_pandas().set_index("id"), expected1, check_dtype=False, check_index_type=False) + + # Drop multiple index values using labels & axis=0 + result2 = df_lf.drop([1, 4], axis=0) + expected2 = df_pd.drop([1, 4], axis=0) + tm.assert_frame_equal(result2.to_pandas().set_index("id"), expected2, check_dtype=False, check_index_type=False)