Repository: databricks/koalas

Issue:
Implement 'keep' parameter for ``drop_duplicates``
Possible implementation #1303 


Discussion:

Resolving patch:
diff --git a/databricks/koalas/frame.py b/databricks/koalas/frame.py
--- a/databricks/koalas/frame.py
+++ b/databricks/koalas/frame.py
@@ -3567,6 +3567,53 @@ def op(kser):
 
         return self._apply_series_op(op)
 
+    def _mark_duplicates(self, subset=None, keep="first"):
+        if len(self._internal.index_names) > 1:
+            raise ValueError("We don't support multi-index yet.")
+
+        if subset is None:
+            subset = self._internal.column_labels
+        else:
+            if isinstance(subset, str):
+                subset = [(subset,)]
+            elif isinstance(subset, tuple):
+                subset = [subset]
+            else:
+                subset = [sub if isinstance(sub, tuple) else (sub,) for sub in subset]
+            diff = set(subset).difference(set(self._internal.column_labels))
+            if len(diff) > 0:
+                raise KeyError(", ".join([str(d) if len(d) > 1 else d[0] for d in diff]))
+        group_cols = [self._internal.column_name_for(label) for label in subset]
+
+        index_column = self._internal.index_columns[0]
+        if self._internal.index_names[0] is not None:
+            name = self._internal.index_names[0]
+        else:
+            name = ("0",)
+        column = "__duplicated__"
+
+        sdf = self._sdf
+
+        if keep == "first" or keep == "last":
+            if keep == "first":
+                ord_func = spark.functions.asc
+            else:
+                ord_func = spark.functions.desc
+            window = (
+                Window.partitionBy(group_cols)
+                .orderBy(ord_func(NATURAL_ORDER_COLUMN_NAME))
+                .rowsBetween(Window.unboundedPreceding, Window.currentRow)
+            )
+            sdf = sdf.withColumn(column, F.row_number().over(window) > 1)
+        elif not keep:
+            window = Window.partitionBy(group_cols).rowsBetween(
+                Window.unboundedPreceding, Window.unboundedFollowing
+            )
+            sdf = sdf.withColumn(column, F.count("*").over(window) > 1)
+        else:
+            raise ValueError("'keep' only supports 'first', 'last' and False")
+        return name, name_like_string(name), sdf
+
     def duplicated(self, subset=None, keep="first"):
         """
         Return boolean Series denoting duplicate rows, optionally only considering certain columns.
@@ -3623,55 +3670,11 @@ def duplicated(self, subset=None, keep="first"):
         """
         from databricks.koalas.series import _col
 
-        if len(self._internal.index_names) > 1:
-            raise ValueError("Now we don't support multi-index Now.")
-
-        if subset is None:
-            subset = self._internal.column_labels
-        else:
-            if isinstance(subset, str):
-                subset = [(subset,)]
-            elif isinstance(subset, tuple):
-                subset = [subset]
-            else:
-                subset = [sub if isinstance(sub, tuple) else (sub,) for sub in subset]
-            diff = set(subset).difference(set(self._internal.column_labels))
-            if len(diff) > 0:
-                raise KeyError(", ".join([str(d) if len(d) > 1 else d[0] for d in diff]))
-        group_cols = [self._internal.column_name_for(label) for label in subset]
-
+        name, column, sdf = self._mark_duplicates(subset, keep)
         index_column = self._internal.index_columns[0]
-        if self._internal.index_names[0] is not None:
-            name = self._internal.index_names[0]
-        else:
-            name = ("0",)
-        column = name_like_string(name)
 
-        sdf = self._sdf
-        if column == index_column:
-            index_column = SPARK_DEFAULT_INDEX_NAME
-            sdf = sdf.select(
-                [self._internal.index_scols[0].alias(index_column)] + self._internal.data_scols
-            )
+        sdf = sdf.withColumn(column, F.col("__duplicated__"))
 
-        if keep == "first" or keep == "last":
-            if keep == "first":
-                ord_func = spark.functions.asc
-            else:
-                ord_func = spark.functions.desc
-            window = (
-                Window.partitionBy(group_cols)
-                .orderBy(ord_func(NATURAL_ORDER_COLUMN_NAME))
-                .rowsBetween(Window.unboundedPreceding, Window.currentRow)
-            )
-            sdf = sdf.withColumn(column, F.row_number().over(window) > 1)
-        elif not keep:
-            window = Window.partitionBy(group_cols).rowsBetween(
-                Window.unboundedPreceding, Window.unboundedFollowing
-            )
-            sdf = sdf.withColumn(column, F.count("*").over(window) > 1)
-        else:
-            raise ValueError("'keep' only support 'first', 'last' and False")
         sdf = sdf.select(scol_for(sdf, index_column), scol_for(sdf, column))
         return _col(
             DataFrame(
@@ -7378,8 +7381,7 @@ def describe(self, percentiles: Optional[List[float]] = None) -> "DataFrame":
         )
         return DataFrame(internal).astype("float64")
 
-    # TODO: implements 'keep' parameters
-    def drop_duplicates(self, subset=None, inplace=False):
+    def drop_duplicates(self, subset=None, keep="first", inplace=False):
         """
         Return DataFrame with duplicate rows removed, optionally only
         considering certain columns.
@@ -7388,13 +7390,19 @@ def drop_duplicates(self, subset=None, inplace=False):
         ----------
         subset : column label or sequence of labels, optional
             Only consider certain columns for identifying duplicates, by
-            default use all of the columns
+            default use all of the columns.
+        keep : {'first', 'last', False}, default 'first'
+            Determines which duplicates (if any) to keep.
+            - ``first`` : Drop duplicates except for the first occurrence.
+            - ``last`` : Drop duplicates except for the last occurrence.
+            - False : Drop all duplicates.
         inplace : boolean, default False
-            Whether to drop duplicates in place or to return a copy
+            Whether to drop duplicates in place or to return a copy.
 
         Returns
         -------
         DataFrame
+            DataFrame with duplicates removed or None if ``inplace=True``.
 
         >>> df = ks.DataFrame(
         ...     {'a': [1, 2, 2, 2, 3], 'b': ['a', 'a', 'a', 'c', 'd']}, columns = ['a', 'b'])
@@ -7406,39 +7414,44 @@ def drop_duplicates(self, subset=None, inplace=False):
         3  2  c
         4  3  d
 
-        >>> df.drop_duplicates().sort_values(['a', 'b'])
+        >>> df.drop_duplicates().sort_index()
            a  b
         0  1  a
         1  2  a
         3  2  c
         4  3  d
 
-        >>> df.drop_duplicates('a').sort_values(['a', 'b'])
+        >>> df.drop_duplicates('a').sort_index()
            a  b
         0  1  a
         1  2  a
         4  3  d
 
-        >>> df.drop_duplicates(['a', 'b']).sort_values(['a', 'b'])
+        >>> df.drop_duplicates(['a', 'b']).sort_index()
            a  b
         0  1  a
         1  2  a
         3  2  c
         4  3  d
+
+        >>> df.drop_duplicates(keep='last').sort_index()
+           a  b
+        0  1  a
+        2  2  a
+        3  2  c
+        4  3  d
+
+        >>> df.drop_duplicates(keep=False).sort_index()
+           a  b
+        0  1  a
+        3  2  c
+        4  3  d
         """
         inplace = validate_bool_kwarg(inplace, "inplace")
-        if subset is None:
-            subset = self._internal.column_labels
-        elif isinstance(subset, str):
-            subset = [(subset,)]
-        elif isinstance(subset, tuple):
-            subset = [subset]
-        else:
-            subset = [sub if isinstance(sub, tuple) else (sub,) for sub in subset]
 
-        sdf = self._sdf.drop(*HIDDEN_COLUMNS).drop_duplicates(
-            subset=[self._internal.column_name_for(label) for label in subset]
-        )
+        sdf = self._mark_duplicates(subset, keep)[-1]
+
+        sdf = sdf.where(~F.col("__duplicated__")).drop("__duplicated__")
         internal = self._internal.with_new_sdf(sdf)
         if inplace:
             self._internal = internal
diff --git a/databricks/koalas/series.py b/databricks/koalas/series.py
--- a/databricks/koalas/series.py
+++ b/databricks/koalas/series.py
@@ -1378,23 +1378,29 @@ def to_list(self):
 
     tolist = to_list
 
-    def drop_duplicates(self, inplace=False):
+    def drop_duplicates(self, keep="first", inplace=False):
         """
-        Return koalas Series with duplicate values removed.
+        Return Series with duplicate values removed.
 
         Parameters
         ----------
-        inplace: bool, default False
-            If True, performs operation inpalce and returns None.
+        keep : {'first', 'last', ``False``}, default 'first'
+            Method to handle dropping duplicates:
+            - 'first' : Drop duplicates except for the first occurrence.
+            - 'last' : Drop duplicates except for the last occurrence.
+            - ``False`` : Drop all duplicates.
+        inplace : bool, default ``False``
+            If ``True``, performs operation inplace and returns None.
 
         Returns
         -------
         Series
-            Series with deplicates dropped.
+            Series with duplicates dropped.
 
         Examples
         --------
         Generate a Series with duplicated entries.
+
         >>> s = ks.Series(['lama', 'cow', 'lama', 'beetle', 'lama', 'hippo'],
         ...               name='animal')
         >>> s.sort_index()
@@ -1406,15 +1412,40 @@ def drop_duplicates(self, inplace=False):
         5     hippo
         Name: animal, dtype: object
 
+        With the 'keep' parameter, the selection behaviour of duplicated values
+        can be changed. The value 'first' keeps the first occurrence for each
+        set of duplicated entries. The default value of keep is 'first'.
+
         >>> s.drop_duplicates().sort_index()
         0      lama
         1       cow
         3    beetle
         5     hippo
         Name: animal, dtype: object
+
+        The value 'last' for parameter 'keep' keeps the last occurrence for
+        each set of duplicated entries.
+
+        >>> s.drop_duplicates(keep='last').sort_index()
+        1       cow
+        3    beetle
+        4      lama
+        5     hippo
+        Name: animal, dtype: object
+
+        The value ``False`` for parameter 'keep' discards all sets of
+        duplicated entries. Setting the value of 'inplace' to ``True`` performs
+        the operation inplace and returns ``None``.
+
+        >>> s.drop_duplicates(keep=False, inplace=True)
+        >>> s.sort_index()
+        1       cow
+        3    beetle
+        5     hippo
+        Name: animal, dtype: object
         """
         inplace = validate_bool_kwarg(inplace, "inplace")
-        kseries = _col(self.to_frame().drop_duplicates())
+        kseries = _col(self.to_frame().drop_duplicates(keep=keep))
 
         if inplace:
             self._internal = kseries._internal


Test patch:
diff --git a/databricks/koalas/tests/test_dataframe.py b/databricks/koalas/tests/test_dataframe.py
--- a/databricks/koalas/tests/test_dataframe.py
+++ b/databricks/koalas/tests/test_dataframe.py
@@ -2071,27 +2071,40 @@ def test_drop_duplicates(self):
         kdf = ks.from_pandas(pdf)
 
         # inplace is False
-        self.assert_eq(pdf.drop_duplicates().sort_index(), kdf.drop_duplicates().sort_index())
-        self.assert_eq(pdf.drop_duplicates("a").sort_index(), kdf.drop_duplicates("a").sort_index())
-        self.assert_eq(
-            pdf.drop_duplicates(["a", "b"]).sort_index(),
-            kdf.drop_duplicates(["a", "b"]).sort_index(),
-        )
+        for keep in ["first", "last", False]:
+            with self.subTest(keep=keep):
+                self.assert_eq(
+                    pdf.drop_duplicates(keep=keep).sort_index(),
+                    kdf.drop_duplicates(keep=keep).sort_index(),
+                )
+                self.assert_eq(
+                    pdf.drop_duplicates("a", keep=keep).sort_index(),
+                    kdf.drop_duplicates("a", keep=keep).sort_index(),
+                )
+                self.assert_eq(
+                    pdf.drop_duplicates(["a", "b"], keep=keep).sort_index(),
+                    kdf.drop_duplicates(["a", "b"], keep=keep).sort_index(),
+                )
 
-        # multi-index columns, inplace is False
         columns = pd.MultiIndex.from_tuples([("x", "a"), ("y", "b")])
         pdf.columns = columns
         kdf.columns = columns
 
-        self.assert_eq(pdf.drop_duplicates().sort_index(), kdf.drop_duplicates().sort_index())
-        self.assert_eq(
-            pdf.drop_duplicates(("x", "a")).sort_index(),
-            kdf.drop_duplicates(("x", "a")).sort_index(),
-        )
-        self.assert_eq(
-            pdf.drop_duplicates([("x", "a"), ("y", "b")]).sort_index(),
-            kdf.drop_duplicates([("x", "a"), ("y", "b")]).sort_index(),
-        )
+        # inplace is False
+        for keep in ["first", "last", False]:
+            with self.subTest("multi-index columns", keep=keep):
+                self.assert_eq(
+                    pdf.drop_duplicates(keep=keep).sort_index(),
+                    kdf.drop_duplicates(keep=keep).sort_index(),
+                )
+                self.assert_eq(
+                    pdf.drop_duplicates(("x", "a"), keep=keep).sort_index(),
+                    kdf.drop_duplicates(("x", "a"), keep=keep).sort_index(),
+                )
+                self.assert_eq(
+                    pdf.drop_duplicates([("x", "a"), ("y", "b")], keep=keep).sort_index(),
+                    kdf.drop_duplicates([("x", "a"), ("y", "b")], keep=keep).sort_index(),
+                )
 
         # inplace is True
         subset_list = [None, "a", ["a", "b"]]
@@ -2494,7 +2507,7 @@ def test_duplicated(self):
             pd.Series(pdf.duplicated(subset=["b"]), name="0").sort_index(),
             kdf.duplicated(subset=["b"]).sort_index(),
         )
-        with self.assertRaisesRegex(ValueError, "'keep' only support 'first', 'last' and False"):
+        with self.assertRaisesRegex(ValueError, "'keep' only supports 'first', 'last' and False"):
             kdf.duplicated(keep="false")
         with self.assertRaisesRegex(KeyError, "'d'"):
             kdf.duplicated(subset=["d"])
diff --git a/databricks/koalas/tests/test_series.py b/databricks/koalas/tests/test_series.py
--- a/databricks/koalas/tests/test_series.py
+++ b/databricks/koalas/tests/test_series.py
@@ -19,6 +19,7 @@
 from distutils.version import LooseVersion
 import inspect
 from io import BytesIO
+from itertools import product
 from datetime import datetime, timedelta
 
 import matplotlib
@@ -1138,17 +1139,22 @@ def test_xs(self):
         self.assert_eq(kser.xs(("a", "lama", "speed")), pser.xs(("a", "lama", "speed")))
 
     def test_duplicates(self):
-        # test on texts
-        pser = pd.Series(["lama", "cow", "lama", "beetle", "lama", "hippo"], name="animal")
-        kser = ks.Series(pser)
-
-        self.assert_eq(pser.drop_duplicates().sort_values(), kser.drop_duplicates().sort_values())
-
-        # test on numbers
-        pser = pd.Series([1, 1, 2, 4, 3])
-        kser = ks.Series(pser)
-
-        self.assert_eq(pser.drop_duplicates().sort_values(), kser.drop_duplicates().sort_values())
+        psers = {
+            "test on texts": pd.Series(
+                ["lama", "cow", "lama", "beetle", "lama", "hippo"], name="animal"
+            ),
+            "test on numbers": pd.Series([1, 1, 2, 4, 3]),
+        }
+        keeps = ["first", "last", False]
+
+        for (msg, pser), keep in product(psers.items(), keeps):
+            with self.subTest(msg, keep=keep):
+                kser = ks.Series(pser)
+
+                self.assert_eq(
+                    pser.drop_duplicates(keep=keep).sort_values(),
+                    kser.drop_duplicates(keep=keep).sort_values(),
+                )
 
     def test_update(self):
         pser = pd.Series([10, 20, 15, 30, 45], name="x")