ENH: #59237

Author: kirill-bash

pandas/core/frame.py

@@ -20,6 +20,7 @@
     Iterator,
     Mapping,
     Sequence,
+    Set as AbstractSet,
 )
 import functools
 from io import StringIO
@@ -6534,7 +6535,7 @@ def dropna(
     @overload
     def drop_duplicates(
         self,
-        subset: Hashable | Sequence[Hashable] | None = ...,
+        subset: Hashable | Sequence[Hashable] | AbstractSet | None = ...,
         *,
         keep: DropKeep = ...,
         inplace: Literal[True],
@@ -6544,7 +6545,7 @@ def drop_duplicates(
     @overload
     def drop_duplicates(
         self,
-        subset: Hashable | Sequence[Hashable] | None = ...,
+        subset: Hashable | Sequence[Hashable] | AbstractSet | None = ...,
         *,
         keep: DropKeep = ...,
         inplace: Literal[False] = ...,
@@ -6554,7 +6555,7 @@ def drop_duplicates(
     @overload
     def drop_duplicates(
         self,
-        subset: Hashable | Sequence[Hashable] | None = ...,
+        subset: Hashable | Sequence[Hashable] | AbstractSet | None = ...,
         *,
         keep: DropKeep = ...,
         inplace: bool = ...,
@@ -6563,7 +6564,7 @@ def drop_duplicates(
 
     def drop_duplicates(
         self,
-        subset: Hashable | Sequence[Hashable] | None = None,
+        subset: Hashable | Sequence[Hashable] | AbstractSet | None = None,
         *,
         keep: DropKeep = "first",
         inplace: bool = False,
@@ -6667,7 +6668,7 @@ def drop_duplicates(
 
     def duplicated(
         self,
-        subset: Hashable | Sequence[Hashable] | None = None,
+        subset: Hashable | Sequence[Hashable] | AbstractSet | None = None,
         keep: DropKeep = "first",
     ) -> Series:
         """
@@ -6792,8 +6793,13 @@ def f(vals) -> tuple[np.ndarray, int]:
             raise KeyError(Index(diff))
 
         if len(subset) == 1 and self.columns.is_unique:
+            # GH#59237 adding support for single element sets
+            if isinstance(subset, set):
+                elem = subset.pop()
+            else:
+                elem = subset[0]
             # GH#45236 This is faster than get_group_index below
-            result = self[subset[0]].duplicated(keep)
+            result = self[elem].duplicated(keep)
             result.name = None
         else:
             vals = (col.values for name, col in self.items() if name in subset)

pandas/tests/frame/methods/test_drop_duplicates.py

@@ -476,3 +476,41 @@ def test_drop_duplicates_non_boolean_ignore_index(arg):
     msg = '^For argument "ignore_index" expected type bool, received type .*.$'
     with pytest.raises(ValueError, match=msg):
         df.drop_duplicates(ignore_index=arg)
+
+
+def test_drop_duplicates_set():
+    # GH#59237
+    df = DataFrame(
+        {
+            "AAA": ["foo", "bar", "foo", "bar", "foo", "bar", "bar", "foo"],
+            "B": ["one", "one", "two", "two", "two", "two", "one", "two"],
+            "C": [1, 1, 2, 2, 2, 2, 1, 2],
+            "D": range(8),
+        }
+    )
+    # single column
+    result = df.drop_duplicates({"AAA"})
+    expected = df[:2]
+    tm.assert_frame_equal(result, expected)
+
+    result = df.drop_duplicates({"AAA"}, keep="last")
+    expected = df.loc[[6, 7]]
+    tm.assert_frame_equal(result, expected)
+
+    result = df.drop_duplicates({"AAA"}, keep=False)
+    expected = df.loc[[]]
+    tm.assert_frame_equal(result, expected)
+    assert len(result) == 0
+
+    # multi column
+    expected = df.loc[[0, 1, 2, 3]]
+    result = df.drop_duplicates({"AAA", "B"})
+    tm.assert_frame_equal(result, expected)
+
+    result = df.drop_duplicates({"AAA", "B"}, keep="last")
+    expected = df.loc[[0, 5, 6, 7]]
+    tm.assert_frame_equal(result, expected)
+
+    result = df.drop_duplicates({"AAA", "B"}, keep=False)
+    expected = df.loc[[0]]
+    tm.assert_frame_equal(result, expected)