Repository: modin-project/modin

Issue:
Add range-partitioning implementation for `.unique()`/`.drop_duplicates()`



Discussion:

Resolving patch:
diff --git a/modin/core/storage_formats/base/query_compiler.py b/modin/core/storage_formats/base/query_compiler.py
--- a/modin/core/storage_formats/base/query_compiler.py
+++ b/modin/core/storage_formats/base/query_compiler.py
@@ -1768,24 +1768,34 @@ def to_timedelta(self, unit="ns", errors="raise"):  # noqa: PR02
             self, unit=unit, errors=errors
         )
 
-    # FIXME: get rid of `**kwargs` parameter (Modin issue #3108).
-    @doc_utils.add_one_column_warning
-    @doc_utils.add_refer_to("Series.unique")
-    def unique(self, **kwargs):
+    # 'qc.unique()' uses most of the arguments from 'df.drop_duplicates()', so refering to this method
+    @doc_utils.add_refer_to("DataFrame.drop_duplicates")
+    def unique(self, keep="first", ignore_index=True, subset=None):
         """
-        Get unique values of `self`.
+        Get unique rows of `self`.
 
         Parameters
         ----------
-        **kwargs : dict
-            Serves compatibility purpose. Does not affect the result.
+        keep : {"first", "last", False}, default: "first"
+            Which duplicates to keep.
+        ignore_index : bool, default: True
+            If ``True``, the resulting axis will be labeled ``0, 1, …, n - 1``.
+        subset : list, optional
+            Only consider certain columns for identifying duplicates, if `None`, use all of the columns.
 
         Returns
         -------
         BaseQueryCompiler
             New QueryCompiler with unique values.
         """
-        return SeriesDefault.register(pandas.Series.unique)(self, **kwargs)
+        if subset is not None:
+            mask = self.getitem_column_array(subset, ignore_order=True)
+        else:
+            mask = self
+        without_duplicates = self.getitem_array(mask.duplicated(keep=keep).invert())
+        if ignore_index:
+            without_duplicates = without_duplicates.reset_index(drop=True)
+        return without_duplicates
 
     @doc_utils.add_one_column_warning
     @doc_utils.add_refer_to("Series.searchsorted")
diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
--- a/modin/core/storage_formats/pandas/query_compiler.py
+++ b/modin/core/storage_formats/pandas/query_compiler.py
@@ -1933,13 +1933,36 @@ def str_split(self, pat=None, n=-1, expand=False, regex=None):
 
     # END String map partitions operations
 
-    def unique(self):
-        new_modin_frame = self._modin_frame.apply_full_axis(
-            0,
-            lambda x: x.squeeze(axis=1).unique(),
-            new_columns=self.columns,
+    def unique(self, keep="first", ignore_index=True, subset=None):
+        # kernels with 'pandas.Series.unique()' work faster
+        can_use_unique_kernel = (
+            subset is None and ignore_index and len(self.columns) == 1 and keep
         )
-        return self.__constructor__(new_modin_frame)
+
+        if not can_use_unique_kernel and not RangePartitioning.get():
+            return super().unique(keep=keep, ignore_index=ignore_index, subset=subset)
+
+        if RangePartitioning.get():
+            new_modin_frame = self._modin_frame._apply_func_to_range_partitioning(
+                key_columns=self.columns.tolist() if subset is None else subset,
+                func=(
+                    (lambda df: pandas.DataFrame(df.squeeze(axis=1).unique()))
+                    if can_use_unique_kernel
+                    else (
+                        lambda df: df.drop_duplicates(
+                            keep=keep, ignore_index=ignore_index, subset=subset
+                        )
+                    )
+                ),
+                preserve_columns=True,
+            )
+        else:
+            new_modin_frame = self._modin_frame.apply_full_axis(
+                0,
+                lambda x: x.squeeze(axis=1).unique(),
+                new_columns=self.columns,
+            )
+        return self.__constructor__(new_modin_frame, shape_hint=self._shape_hint)
 
     def searchsorted(self, **kwargs):
         def searchsorted(df):
diff --git a/modin/pandas/base.py b/modin/pandas/base.py
--- a/modin/pandas/base.py
+++ b/modin/pandas/base.py
@@ -1511,13 +1511,12 @@ def drop_duplicates(
                     subset = list(subset)
             else:
                 subset = [subset]
-            df = self[subset]
-        else:
-            df = self
-        duplicated = df.duplicated(keep=keep)
-        result = self[~duplicated]
-        if ignore_index:
-            result.index = pandas.RangeIndex(stop=len(result))
+            if len(diff := pandas.Index(subset).difference(self.columns)) > 0:
+                raise KeyError(diff)
+        result_qc = self._query_compiler.unique(
+            keep=keep, ignore_index=ignore_index, subset=subset
+        )
+        result = self.__constructor__(query_compiler=result_qc)
         if inplace:
             self._update_inplace(result._query_compiler)
         else:


Test patch:
diff --git a/modin/pandas/test/dataframe/test_join_sort.py b/modin/pandas/test/dataframe/test_join_sort.py
--- a/modin/pandas/test/dataframe/test_join_sort.py
+++ b/modin/pandas/test/dataframe/test_join_sort.py
@@ -34,6 +34,7 @@
     generate_multiindex,
     random_state,
     rotate_decimal_digits_or_symbols,
+    sort_if_range_partitioning,
     test_data,
     test_data_keys,
     test_data_values,
@@ -247,10 +248,6 @@ def test_join_6602():
     ],
 )
 def test_merge(test_data, test_data2):
-    # RangePartitioning merge always produces sorted result, so we have to sort
-    # pandas' result as well in order them to match
-    comparator = df_equals_and_sort if RangePartitioning.get() else df_equals
-
     modin_df = pd.DataFrame(
         test_data,
         columns=["col{}".format(i) for i in range(test_data.shape[1])],
@@ -283,7 +280,7 @@ def test_merge(test_data, test_data2):
             pandas_result = pandas_df.merge(
                 pandas_df2, how=hows[i], on=ons[j], sort=sorts[j]
             )
-            comparator(modin_result, pandas_result)
+            sort_if_range_partitioning(modin_result, pandas_result)
 
             modin_result = modin_df.merge(
                 modin_df2,
@@ -299,7 +296,7 @@ def test_merge(test_data, test_data2):
                 right_on="key",
                 sort=sorts[j],
             )
-            comparator(modin_result, pandas_result)
+            sort_if_range_partitioning(modin_result, pandas_result)
 
     # Test for issue #1771
     modin_df = pd.DataFrame({"name": np.arange(40)})
@@ -308,7 +305,7 @@ def test_merge(test_data, test_data2):
     pandas_df2 = pandas.DataFrame({"name": [39], "position": [0]})
     modin_result = modin_df.merge(modin_df2, on="name", how="inner")
     pandas_result = pandas_df.merge(pandas_df2, on="name", how="inner")
-    comparator(modin_result, pandas_result)
+    sort_if_range_partitioning(modin_result, pandas_result)
 
     frame_data = {
         "col1": [0, 1, 2, 3],
@@ -329,7 +326,7 @@ def test_merge(test_data, test_data2):
         # Defaults
         modin_result = modin_df.merge(modin_df2, how=how)
         pandas_result = pandas_df.merge(pandas_df2, how=how)
-        comparator(modin_result, pandas_result)
+        sort_if_range_partitioning(modin_result, pandas_result)
 
         # left_on and right_index
         modin_result = modin_df.merge(
@@ -338,7 +335,7 @@ def test_merge(test_data, test_data2):
         pandas_result = pandas_df.merge(
             pandas_df2, how=how, left_on="col1", right_index=True
         )
-        comparator(modin_result, pandas_result)
+        sort_if_range_partitioning(modin_result, pandas_result)
 
         # left_index and right_on
         modin_result = modin_df.merge(
@@ -347,7 +344,7 @@ def test_merge(test_data, test_data2):
         pandas_result = pandas_df.merge(
             pandas_df2, how=how, left_index=True, right_on="col1"
         )
-        comparator(modin_result, pandas_result)
+        sort_if_range_partitioning(modin_result, pandas_result)
 
         # left_on and right_on col1
         modin_result = modin_df.merge(
@@ -356,7 +353,7 @@ def test_merge(test_data, test_data2):
         pandas_result = pandas_df.merge(
             pandas_df2, how=how, left_on="col1", right_on="col1"
         )
-        comparator(modin_result, pandas_result)
+        sort_if_range_partitioning(modin_result, pandas_result)
 
         # left_on and right_on col2
         modin_result = modin_df.merge(
@@ -365,7 +362,7 @@ def test_merge(test_data, test_data2):
         pandas_result = pandas_df.merge(
             pandas_df2, how=how, left_on="col2", right_on="col2"
         )
-        comparator(modin_result, pandas_result)
+        sort_if_range_partitioning(modin_result, pandas_result)
 
         # left_index and right_index
         modin_result = modin_df.merge(
@@ -374,7 +371,7 @@ def test_merge(test_data, test_data2):
         pandas_result = pandas_df.merge(
             pandas_df2, how=how, left_index=True, right_index=True
         )
-        comparator(modin_result, pandas_result)
+        sort_if_range_partitioning(modin_result, pandas_result)
 
     # Cannot merge a Series without a name
     ps = pandas.Series(frame_data2.get("col1"))
@@ -383,7 +380,7 @@ def test_merge(test_data, test_data2):
         modin_df,
         pandas_df,
         lambda df: df.merge(ms if isinstance(df, pd.DataFrame) else ps),
-        comparator=comparator,
+        comparator=sort_if_range_partitioning,
         expected_exception=ValueError("Cannot merge a Series without a name"),
     )
 
@@ -394,7 +391,7 @@ def test_merge(test_data, test_data2):
         modin_df,
         pandas_df,
         lambda df: df.merge(ms if isinstance(df, pd.DataFrame) else ps),
-        comparator=comparator,
+        comparator=sort_if_range_partitioning,
     )
 
     with pytest.raises(TypeError):
diff --git a/modin/pandas/test/dataframe/test_map_metadata.py b/modin/pandas/test/dataframe/test_map_metadata.py
--- a/modin/pandas/test/dataframe/test_map_metadata.py
+++ b/modin/pandas/test/dataframe/test_map_metadata.py
@@ -38,6 +38,7 @@
     name_contains,
     numeric_dfs,
     random_state,
+    sort_if_range_partitioning,
     test_data,
     test_data_keys,
     test_data_values,
@@ -1068,7 +1069,7 @@ def test_drop_duplicates(data, keep, subset, ignore_index):
                 keep=keep, inplace=False, subset=subset, ignore_index=ignore_index
             )
     else:
-        df_equals(
+        sort_if_range_partitioning(
             pandas_df.drop_duplicates(
                 keep=keep, inplace=False, subset=subset, ignore_index=ignore_index
             ),
@@ -1078,7 +1079,7 @@ def test_drop_duplicates(data, keep, subset, ignore_index):
         )
 
     try:
-        pandas_results = pandas_df.drop_duplicates(
+        pandas_df.drop_duplicates(
             keep=keep, inplace=True, subset=subset, ignore_index=ignore_index
         )
     except Exception as err:
@@ -1087,10 +1088,10 @@ def test_drop_duplicates(data, keep, subset, ignore_index):
                 keep=keep, inplace=True, subset=subset, ignore_index=ignore_index
             )
     else:
-        modin_results = modin_df.drop_duplicates(
+        modin_df.drop_duplicates(
             keep=keep, inplace=True, subset=subset, ignore_index=ignore_index
         )
-        df_equals(modin_results, pandas_results)
+        sort_if_range_partitioning(modin_df, pandas_df)
 
 
 def test_drop_duplicates_with_missing_index_values():
@@ -1168,7 +1169,7 @@ def test_drop_duplicates_with_missing_index_values():
     modin_df = pd.DataFrame(data["data"], index=data["index"], columns=data["columns"])
     modin_result = modin_df.sort_values(["id", "time"]).drop_duplicates(["id"])
     pandas_result = pandas_df.sort_values(["id", "time"]).drop_duplicates(["id"])
-    df_equals(modin_result, pandas_result)
+    sort_if_range_partitioning(modin_result, pandas_result)
 
 
 def test_drop_duplicates_after_sort():
@@ -1183,7 +1184,7 @@ def test_drop_duplicates_after_sort():
 
     modin_result = modin_df.sort_values(["value", "time"]).drop_duplicates(["value"])
     pandas_result = pandas_df.sort_values(["value", "time"]).drop_duplicates(["value"])
-    df_equals(modin_result, pandas_result)
+    sort_if_range_partitioning(modin_result, pandas_result)
 
 
 def test_drop_duplicates_with_repeated_index_values():
@@ -1191,7 +1192,12 @@ def test_drop_duplicates_with_repeated_index_values():
     data = [[0], [1], [0]]
     index = [0, 0, 0]
     modin_df, pandas_df = create_test_dfs(data, index=index)
-    eval_general(modin_df, pandas_df, lambda df: df.drop_duplicates())
+    eval_general(
+        modin_df,
+        pandas_df,
+        lambda df: df.drop_duplicates(),
+        comparator=sort_if_range_partitioning,
+    )
 
 
 @pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)
diff --git a/modin/pandas/test/test_general.py b/modin/pandas/test/test_general.py
--- a/modin/pandas/test/test_general.py
+++ b/modin/pandas/test/test_general.py
@@ -32,6 +32,7 @@
     default_to_pandas_ignore_string,
     df_equals,
     eval_general,
+    sort_if_range_partitioning,
     sort_index_for_equal_values,
     test_data_keys,
     test_data_values,
@@ -566,14 +567,18 @@ def test_pivot_table():
 
 
 def test_unique():
+    comparator = lambda *args: sort_if_range_partitioning(  # noqa: E731
+        *args, comparator=assert_array_equal
+    )
+
     modin_result = pd.unique([2, 1, 3, 3])
     pandas_result = pandas.unique([2, 1, 3, 3])
-    assert_array_equal(modin_result, pandas_result)
+    comparator(modin_result, pandas_result)
     assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.unique(pd.Series([2] + [1] * 5))
     pandas_result = pandas.unique(pandas.Series([2] + [1] * 5))
-    assert_array_equal(modin_result, pandas_result)
+    comparator(modin_result, pandas_result)
     assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.unique(
@@ -582,7 +587,7 @@ def test_unique():
     pandas_result = pandas.unique(
         pandas.Series([pandas.Timestamp("20160101"), pandas.Timestamp("20160101")])
     )
-    assert_array_equal(modin_result, pandas_result)
+    comparator(modin_result, pandas_result)
     assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.unique(
@@ -601,7 +606,7 @@ def test_unique():
             ]
         )
     )
-    assert_array_equal(modin_result, pandas_result)
+    comparator(modin_result, pandas_result)
     assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.unique(
@@ -620,12 +625,12 @@ def test_unique():
             ]
         )
     )
-    assert_array_equal(modin_result, pandas_result)
+    comparator(modin_result, pandas_result)
     assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.unique(pd.Series(pd.Categorical(list("baabc"))))
     pandas_result = pandas.unique(pandas.Series(pandas.Categorical(list("baabc"))))
-    assert_array_equal(modin_result, pandas_result)
+    comparator(modin_result, pandas_result)
     assert modin_result.shape == pandas_result.shape
 
 
diff --git a/modin/pandas/test/test_series.py b/modin/pandas/test/test_series.py
--- a/modin/pandas/test/test_series.py
+++ b/modin/pandas/test/test_series.py
@@ -63,6 +63,7 @@
     quantiles_keys,
     quantiles_values,
     random_state,
+    sort_if_range_partitioning,
     string_na_rep_keys,
     string_na_rep_values,
     string_sep_keys,
@@ -1695,10 +1696,12 @@ def test_drop(data):
 @pytest.mark.parametrize("inplace", [True, False], ids=["True", "False"])
 def test_drop_duplicates(data, keep, inplace):
     modin_series, pandas_series = create_test_series(data)
-    df_equals(
-        modin_series.drop_duplicates(keep=keep, inplace=inplace),
-        pandas_series.drop_duplicates(keep=keep, inplace=inplace),
-    )
+    modin_res = modin_series.drop_duplicates(keep=keep, inplace=inplace)
+    pandas_res = pandas_series.drop_duplicates(keep=keep, inplace=inplace)
+    if inplace:
+        sort_if_range_partitioning(modin_series, pandas_series)
+    else:
+        sort_if_range_partitioning(modin_res, pandas_res)
 
 
 @pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)
@@ -3671,22 +3674,26 @@ def test_tz_localize():
 
 @pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)
 def test_unique(data):
+    comparator = lambda *args: sort_if_range_partitioning(  # noqa: E731
+        *args, comparator=assert_array_equal
+    )
+
     modin_series, pandas_series = create_test_series(data)
     modin_result = modin_series.unique()
     pandas_result = pandas_series.unique()
-    assert_array_equal(modin_result, pandas_result)
+    comparator(modin_result, pandas_result)
     assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.Series([2, 1, 3, 3], name="A").unique()
     pandas_result = pandas.Series([2, 1, 3, 3], name="A").unique()
-    assert_array_equal(modin_result, pandas_result)
+    comparator(modin_result, pandas_result)
     assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.Series([pd.Timestamp("2016-01-01") for _ in range(3)]).unique()
     pandas_result = pandas.Series(
         [pd.Timestamp("2016-01-01") for _ in range(3)]
     ).unique()
-    assert_array_equal(modin_result, pandas_result)
+    comparator(modin_result, pandas_result)
     assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.Series(
@@ -3695,12 +3702,12 @@ def test_unique(data):
     pandas_result = pandas.Series(
         [pd.Timestamp("2016-01-01", tz="US/Eastern") for _ in range(3)]
     ).unique()
-    assert_array_equal(modin_result, pandas_result)
+    comparator(modin_result, pandas_result)
     assert modin_result.shape == pandas_result.shape
 
     modin_result = pandas.Series(pd.Categorical(list("baabc"))).unique()
     pandas_result = pd.Series(pd.Categorical(list("baabc"))).unique()
-    assert_array_equal(modin_result, pandas_result)
+    comparator(modin_result, pandas_result)
     assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.Series(
@@ -3709,7 +3716,7 @@ def test_unique(data):
     pandas_result = pandas.Series(
         pd.Categorical(list("baabc"), categories=list("abc"), ordered=True)
     ).unique()
-    assert_array_equal(modin_result, pandas_result)
+    comparator(modin_result, pandas_result)
     assert modin_result.shape == pandas_result.shape
 
 
diff --git a/modin/pandas/test/utils.py b/modin/pandas/test/utils.py
--- a/modin/pandas/test/utils.py
+++ b/modin/pandas/test/utils.py
@@ -41,6 +41,8 @@
     Engine,
     MinPartitionSize,
     NPartitions,
+    RangePartitioning,
+    RangePartitioningGroupby,
     TestDatasetSize,
     TrackFileLeaks,
 )
@@ -663,6 +665,27 @@ def assert_dtypes_equal(df1, df2):
                 break
 
 
+def sort_data(data):
+    """Sort the passed sequence."""
+    if isinstance(data, (pandas.DataFrame, pd.DataFrame)):
+        return data.sort_values(data.columns.to_list(), ignore_index=True)
+    elif isinstance(data, (pandas.Series, pd.Series)):
+        return data.sort_values()
+    else:
+        return np.sort(data)
+
+
+def sort_if_range_partitioning(df1, df2, comparator=None):
+    """Sort the passed objects if 'RangePartitioning' is enabled and compare the sorted results."""
+    if comparator is None:
+        comparator = df_equals
+
+    if RangePartitioning.get() or RangePartitioningGroupby.get():
+        df1, df2 = sort_data(df1), sort_data(df2)
+
+    comparator(df1, df2)
+
+
 def df_equals(df1, df2, check_dtypes=True):
     """Tests if df1 and df2 are equal.