Repository: alteryx/woodwork (python)

Issue:
Add WoodworkTableAccessor.drop
- Creates a new DataFrame without the specified columns that has Woodwork initialized using the typing information from the original DataFrame


Pull request:
Add drop to WoodworkTableAccessor
- Allows for `df.ww.drop`, returning a new dataframe with Woodwork initialized
- Closes #615 


Patch:
diff --git a/docs/source/api_reference.rst b/docs/source/api_reference.rst
index d5b9119..dea1466 100644
--- a/docs/source/api_reference.rst
+++ b/docs/source/api_reference.rst
@@ -12,6 +12,7 @@ WoodworkTableAccessor
     WoodworkTableAccessor
     WoodworkTableAccessor.describe
     WoodworkTableAccessor.describe_dict
+    WoodworkTableAccessor.drop
     WoodworkTableAccessor.iloc
     WoodworkTableAccessor.init
     WoodworkTableAccessor.loc
diff --git a/docs/source/release_notes.rst b/docs/source/release_notes.rst
index eb2be94..ac2c7f9 100644
--- a/docs/source/release_notes.rst
+++ b/docs/source/release_notes.rst
@@ -34,6 +34,7 @@ Release Notes
         * Add ``value_counts`` to WoodworkTableAccessor (:pr:`632`)
         * Add KoalasColumnAccessor (:pr:`634`)
         * Add ``pop`` to WoodworkTableAccessor (:pr:`636`)
+        * Add ``drop`` to WoodworkTableAccessor (:pr:`640`)
     * Fixes
         * Create new Schema object when performing pandas operation on Accessors (:pr:`595`)
     * Changes
diff --git a/woodwork/table_accessor.py b/woodwork/table_accessor.py
index a5f9fe5..e34994b 100644
--- a/woodwork/table_accessor.py
+++ b/woodwork/table_accessor.py
@@ -457,6 +457,33 @@ class WoodworkTableAccessor:
 
         return series
 
+    def drop(self, columns):
+        """Drop specified columns from a DataFrame.
+
+        Args:
+            columns (str or list[str]): Column name or names to drop. Must be present in the DataTable.
+
+        Returns:
+            (pd.DataFrame): DataFrame with the specified columns removed, maintaining Woodwork typing information.
+
+        Note:
+            This method is used for removing columns only. To remove rows with ``drop``, go through the
+            DataFrame directly and then reinitialize Woodwork with ``DataFrame.ww.init``
+            instead of calling ``DataFrame.ww.drop``.
+        """
+        if not isinstance(columns, (list, set)):
+            columns = [columns]
+
+        not_present = [col for col in columns if col not in self._dataframe.columns]
+        if not_present:
+            raise ValueError(f'{not_present} not found in DataFrame')
+
+        new_schema = self._schema._get_subset_schema([col for col in self._dataframe.columns if col not in columns])
+        new_df = self._dataframe.drop(columns, axis=1)
+        new_df.ww.init(schema=new_schema)
+
+        return new_df
+
     def mutual_information_dict(self, num_bins=10, nrows=None):
         """
         Calculates mutual information between all pairs of columns in the DataFrame that


Test patch:
diff --git a/woodwork/tests/accessor/test_table_accessor.py b/woodwork/tests/accessor/test_table_accessor.py
index 7948061..f3b1df0 100644
--- a/woodwork/tests/accessor/test_table_accessor.py
+++ b/woodwork/tests/accessor/test_table_accessor.py
@@ -721,13 +721,6 @@ def test_underlying_index(sample_df):
     assert schema_df.index.name is None
     assert type(schema_df.index) == specified_index
 
-    # --> add back when df.ww.drop is implemented
-    # schema_dropped = schema.drop('made_index')
-    # assert 'made_index' not in schema_dropped.columns
-    # assert 'made_index' not in schema_dropped._dataframe.columns
-    # assert type(schema_dropped._dataframe.index) == unspecified_index
-    # assert type(schema_dropped.to_dataframe().index) == unspecified_index
-
 
 def test_accessor_already_sorted(sample_unsorted_df):
     if dd and isinstance(sample_unsorted_df, dd.DataFrame):
@@ -813,7 +806,7 @@ def test_ordinal_with_nan_values():
 
 def test_accessor_with_falsy_column_names(falsy_names_df):
     if dd and isinstance(falsy_names_df, dd.DataFrame):
-        pytest.xfail('Dask DataTables cannot handle integer column names')
+        pytest.xfail('Dask DataFrames cannot handle integer column names')
     xfail_dask_and_koalas(falsy_names_df)
 
     schema_df = falsy_names_df.copy()
@@ -1448,3 +1441,71 @@ def test_pop_error(sample_df):
 
     with pytest.raises(LookupError, match="Column with name missing not found in DataFrame"):
         sample_df.ww.pop("missing")
+
+
+def test_accessor_drop(sample_df):
+    xfail_dask_and_koalas(sample_df)
+
+    schema_df = sample_df.copy()
+    schema_df.ww.init()
+
+    single_input_df = schema_df.ww.drop('is_registered')
+    assert len(single_input_df.ww.columns) == (len(schema_df.columns) - 1)
+    assert 'is_registered' not in single_input_df.ww.columns
+    assert to_pandas(schema_df).drop('is_registered', axis='columns').equals(to_pandas(single_input_df))
+
+    list_input_df = schema_df.ww.drop(['is_registered'])
+    assert len(list_input_df.ww.columns) == (len(schema_df.columns) - 1)
+    assert 'is_registered' not in list_input_df.ww.columns
+    assert to_pandas(schema_df).drop('is_registered', axis='columns').equals(to_pandas(list_input_df))
+    # should be equal to the single input example above
+    assert single_input_df.ww.schema == list_input_df.ww.schema
+    assert to_pandas(single_input_df).equals(to_pandas(list_input_df))
+
+    multiple_list_df = schema_df.ww.drop(['age', 'full_name', 'is_registered'])
+    assert len(multiple_list_df.ww.columns) == (len(schema_df.columns) - 3)
+    assert 'is_registered' not in multiple_list_df.ww.columns
+    assert 'full_name' not in multiple_list_df.ww.columns
+    assert 'age' not in multiple_list_df.ww.columns
+
+    assert to_pandas(schema_df).drop(['is_registered', 'age', 'full_name'], axis='columns').equals(to_pandas(multiple_list_df))
+
+    # Drop the same columns in a different order and confirm resulting DataFrame column order doesn't change
+    different_order_df = schema_df.ww.drop(['is_registered', 'age', 'full_name'])
+    assert different_order_df.ww.schema == multiple_list_df.ww.schema
+    assert to_pandas(multiple_list_df).equals(to_pandas(different_order_df))
+
+
+def test_accessor_drop_indices(sample_df):
+    xfail_dask_and_koalas(sample_df)
+
+    sample_df.ww.init(index='id', time_index='signup_date')
+    assert sample_df.ww.index == 'id'
+    assert sample_df.ww.time_index == 'signup_date'
+
+    dropped_index_df = sample_df.ww.drop('id')
+    assert 'id' not in dropped_index_df.ww.columns
+    assert dropped_index_df.ww.index is None
+    assert dropped_index_df.ww.time_index == 'signup_date'
+
+    dropped_time_index_df = sample_df.ww.drop(['signup_date'])
+    assert 'signup_date' not in dropped_time_index_df.ww.columns
+    assert dropped_time_index_df.ww.time_index is None
+    assert dropped_time_index_df.ww.index == 'id'
+
+
+def test_accessor_drop_errors(sample_df):
+    xfail_dask_and_koalas(sample_df)
+
+    sample_df.ww.init()
+
+    error = re.escape("['not_present'] not found in DataFrame")
+    with pytest.raises(ValueError, match=error):
+        sample_df.ww.drop('not_present')
+
+    with pytest.raises(ValueError, match=error):
+        sample_df.ww.drop(['age', 'not_present'])
+
+    error = re.escape("['not_present1', 4] not found in DataFrame")
+    with pytest.raises(ValueError, match=error):
+        sample_df.ww.drop(['not_present1', 4])