Implement Python drop_duplicates with cudf::stable_distinct.

cpp/include/cudf/detail/stream_compaction.hpp

@@ -91,24 +91,9 @@ std::unique_ptr<table> distinct(
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
 /**
- * @brief Create a new table without duplicate rows.
+ * @copydoc cudf::stable_distinct
  *
- * Given an `input` table_view, each row is copied to the output table to create a set of distinct
- * rows. The row order is guaranteed to be preserved as in the input.
- *
- * If there are duplicate rows, which row to be copied depends on the specified value of the `keep`
- * parameter.
- *
- * This API produces exactly the same set of output rows as `cudf::distinct`.
- *
- * @param input The input table
- * @param keys Vector of indices indicating key columns in the `input` table
- * @param keep Copy any, first, last, or none of the found duplicates
- * @param nulls_equal Flag to specify whether null elements should be considered as equal
- * @param nans_equal Flag to specify whether NaN elements should be considered as equal
- * @param stream CUDA stream used for device memory operations and kernel launches
- * @param mr Device memory resource used to allocate the returned table
- * @return A table containing the resulting distinct rows
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches.
  */
 std::unique_ptr<table> stable_distinct(
   table_view const& input,

cpp/include/cudf/stream_compaction.hpp

@@ -280,6 +280,33 @@ std::unique_ptr<table> distinct(
   nan_equality nans_equal             = nan_equality::ALL_EQUAL,
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
+/**
+ * @brief Create a new table without duplicate rows.
+ *
+ * Given an `input` table_view, each row is copied to the output table to create a set of distinct
+ * rows. The row order is guaranteed to be preserved as in the input.
+ *
+ * If there are duplicate rows, which row to be copied depends on the specified value of the `keep`
+ * parameter.
+ *
+ * This API produces exactly the same set of output rows as `cudf::distinct`.
+ *
+ * @param input The input table
+ * @param keys Vector of indices indicating key columns in the `input` table
+ * @param keep Copy any, first, last, or none of the found duplicates
+ * @param nulls_equal Flag to specify whether null elements should be considered as equal
+ * @param nans_equal Flag to specify whether NaN elements should be considered as equal
+ * @param mr Device memory resource used to allocate the returned table
+ * @return A table containing the resulting distinct rows
+ */
+std::unique_ptr<table> stable_distinct(
+  table_view const& input,
+  std::vector<size_type> const& keys,
+  duplicate_keep_option keep          = duplicate_keep_option::KEEP_ANY,
+  null_equality nulls_equal           = null_equality::EQUAL,
+  nan_equality nans_equal             = nan_equality::ALL_EQUAL,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
 /**
  * @brief Count the number of consecutive groups of equivalent rows in a column.
  *

cpp/src/stream_compaction/distinct.cu

@@ -162,4 +162,16 @@ std::unique_ptr<table> distinct(table_view const& input,
     input, keys, keep, nulls_equal, nans_equal, cudf::default_stream_value, mr);
 }
 
+std::unique_ptr<table> stable_distinct(table_view const& input,
+                                       std::vector<size_type> const& keys,
+                                       duplicate_keep_option keep,
+                                       null_equality nulls_equal,
+                                       nan_equality nans_equal,
+                                       rmm::mr::device_memory_resource* mr)
+{
+  CUDF_FUNC_RANGE();
+  return detail::stable_distinct(
+    input, keys, keep, nulls_equal, nans_equal, cudf::default_stream_value, mr);
+}
+
 }  // namespace cudf

python/cudf/cudf/_lib/cpp/stream_compaction.pxd

@@ -43,3 +43,10 @@ cdef extern from "cudf/stream_compaction.hpp" namespace "cudf" \
         column_view source_table,
         null_policy null_handling,
         nan_policy nan_handling) except +
+
+    cdef unique_ptr[table] stable_distinct(
+        table_view input,
+        vector[size_type] keys,
+        duplicate_keep_option keep,
+        null_equality nulls_equal,
+    ) except +

python/cudf/cudf/_lib/stream_compaction.pyx

@@ -15,6 +15,7 @@ from cudf._lib.cpp.stream_compaction cimport (
     distinct_count as cpp_distinct_count,
     drop_nulls as cpp_drop_nulls,
     duplicate_keep_option,
+    stable_distinct as cpp_stable_distinct,
     unique as cpp_unique,
 )
 from cudf._lib.cpp.table.table cimport table
@@ -107,6 +108,66 @@ def apply_boolean_mask(list columns, Column boolean_mask):
     return columns_from_unique_ptr(move(c_result))
 
 
+def stable_distinct(list columns,
+                    object keys,
+                    object keep='first',
+                    bool nulls_are_equal=True):
+
+    """
+    Drop duplicate rows in provided columns. Retain original order.
+
+    Parameters
+    ----------
+    columns : List of columns
+    keys : List of column indices. If set, then these columns are checked for
+           duplicates rather than all of columns (optional)
+    keep : keep 'first' or 'last' or none of the duplicate rows
+    nulls_are_equal : if True, nulls are treated equal else not.
+
+    Returns
+    -------
+    List of columns with duplicates dropped
+
+    """
+    cdef vector[size_type] cpp_keys = (
+        keys if keys is not None else range(len(columns))
+    )
+    cdef duplicate_keep_option cpp_keep_option
+
+    if keep == 'first':
+        cpp_keep_option = duplicate_keep_option.KEEP_FIRST
+    elif keep == 'last':
+        cpp_keep_option = duplicate_keep_option.KEEP_LAST
+    elif keep is False:
+        cpp_keep_option = duplicate_keep_option.KEEP_NONE
+    else:
+        raise ValueError('keep must be either "first", "last" or False')
+
+    # shifting the index number by number of index columns
+    cdef null_equality cpp_nulls_equal = (
+        null_equality.EQUAL
+        if nulls_are_equal
+        else null_equality.UNEQUAL
+    )
+
+    cdef table_view source_table_view = table_view_from_columns(columns)
+    cdef table_view keys_view = source_table_view.select(cpp_keys)
+
+    cdef unique_ptr[table] c_result
+
+    with nogil:
+        c_result = move(
+            cpp_stable_distinct(
+                source_table_view,
+                cpp_keys,
+                cpp_keep_option,
+                cpp_nulls_equal
+            )
+        )
+
+    return columns_from_unique_ptr(move(c_result))
+
+
 def drop_duplicates(list columns,
                     object keys=None,
                     object keep='first',

python/cudf/cudf/core/dataframe.py

@@ -2946,7 +2946,12 @@ def diff(self, periods=1, axis=0):
 
     @_cudf_nvtx_annotate
     def drop_duplicates(
-        self, subset=None, keep="first", inplace=False, ignore_index=False
+        self,
+        subset=None,
+        keep="first",
+        inplace=False,
+        ignore_index=False,
+        preserve_order=False,
     ):
         """
         Return DataFrame with duplicate rows removed, optionally only
@@ -3017,7 +3022,10 @@ def drop_duplicates(
         1  Yum Yum   cup     4.0
         """  # noqa: E501
         outdf = super().drop_duplicates(
-            subset=subset, keep=keep, ignore_index=ignore_index
+            subset=subset,
+            keep=keep,
+            ignore_index=ignore_index,
+            preserve_order=preserve_order,
         )
 
         return self._mimic_inplace(outdf, inplace=inplace)

python/cudf/cudf/core/indexed_frame.py

@@ -1536,6 +1536,7 @@ def drop_duplicates(
         keep="first",
         nulls_are_equal=True,
         ignore_index=False,
+        preserve_order=False,
     ):
         """
         Drop duplicate rows in frame.
@@ -1569,18 +1570,32 @@ def drop_duplicates(
         keys = self._positions_from_column_names(
             subset, offset_by_index_columns=not ignore_index
         )
-        return self._from_columns_like_self(
-            libcudf.stream_compaction.drop_duplicates(
-                list(self._columns)
-                if ignore_index
-                else list(self._index._columns + self._columns),
-                keys=keys,
-                keep=keep,
-                nulls_are_equal=nulls_are_equal,
-            ),
-            self._column_names,
-            self._index.names if not ignore_index else None,
-        )
+        if preserve_order:
+            return self._from_columns_like_self(
+                libcudf.stream_compaction.stable_distinct(
+                    list(self._columns)
+                    if ignore_index
+                    else list(self._index._columns + self._columns),
+                    keys=keys,
+                    keep=keep,
+                    nulls_are_equal=nulls_are_equal,
+                ),
+                self._column_names,
+                self._index.names if not ignore_index else None,
+            )
+        else:
+            return self._from_columns_like_self(
+                libcudf.stream_compaction.drop_duplicates(
+                    list(self._columns)
+                    if ignore_index
+                    else list(self._index._columns + self._columns),
+                    keys=keys,
+                    keep=keep,
+                    nulls_are_equal=nulls_are_equal,
+                ),
+                self._column_names,
+                self._index.names if not ignore_index else None,
+            )
 
     @_cudf_nvtx_annotate
     def _empty_like(self, keep_index=True):