SNOW-2439665: Add support for drop_duplicates

Author: sfc-gh-helmeleegy

CHANGELOG.md

@@ -143,6 +143,7 @@
   - `groupby.median`
   - `groupby.std`
   - `groupby.var`
+  - `drop_duplicates`
 - Reuse row count from the relaxed query compiler in `get_axis_len`.
 
 #### Bug Fixes

src/snowflake/snowpark/modin/plugin/_internal/indexing_utils.py

@@ -1130,7 +1130,9 @@ def get_frame_by_row_label(
 
     # boolean indexer
     if isinstance(key_datatype, BooleanType):
-        return _get_frame_by_row_label_boolean_frame(internal_frame, key)
+        return _get_frame_by_row_label_boolean_frame(
+            internal_frame, key, dummy_row_pos_mode
+        )
 
     return _get_frame_by_row_label_non_boolean_frame(
         internal_frame, key, dummy_row_pos_mode
@@ -1470,6 +1472,7 @@ def generate_bound_column(
 def _get_frame_by_row_label_boolean_frame(
     internal_frame: InternalFrame,
     key: InternalFrame,
+    dummy_row_pos_mode: bool = False,
 ) -> InternalFrame:
     """
     Select rows with boolean frame key. Here, if the frame and key's index are aligned, then the join is on their row
@@ -1488,6 +1491,7 @@ def _get_frame_by_row_label_boolean_frame(
         internal_frame,
         key,
         "coalesce",
+        dummy_row_pos_mode,
     )
 
     key_bool_val_col = col(

src/snowflake/snowpark/modin/plugin/_internal/ordered_dataframe.py

@@ -1257,17 +1257,31 @@ def join(
 
             left = left.ensure_row_position_column(dummy_row_pos_mode=True)
             if len(left_on_cols) == 1 and ROW_POSITION_COLUMN_LABEL in left_on_cols[0]:
+                left_row_position_snowflake_quoted_identifier = (
+                    left.row_position_snowflake_quoted_identifier
+                )
                 left.row_position_snowflake_quoted_identifier = None
-                left = left.ensure_row_position_column(dummy_row_pos_mode=False)
+                new_left = left.ensure_row_position_column(dummy_row_pos_mode=False)
+                left.row_position_snowflake_quoted_identifier = (
+                    left_row_position_snowflake_quoted_identifier
+                )
+                left = new_left
                 assert left.row_position_snowflake_quoted_identifier is not None
                 left_on_cols = [left.row_position_snowflake_quoted_identifier]
             right = right.ensure_row_position_column(dummy_row_pos_mode=True)
             if (
                 len(right_on_cols) == 1
                 and ROW_POSITION_COLUMN_LABEL in right_on_cols[0]
             ):
+                right_row_position_snowflake_quoted_identifier = (
+                    right.row_position_snowflake_quoted_identifier
+                )
                 right.row_position_snowflake_quoted_identifier = None
-                right = right.ensure_row_position_column(dummy_row_pos_mode=False)
+                new_right = right.ensure_row_position_column(dummy_row_pos_mode=False)
+                right.row_position_snowflake_quoted_identifier = (
+                    right_row_position_snowflake_quoted_identifier
+                )
+                right = new_right
                 assert right.row_position_snowflake_quoted_identifier is not None
                 right_on_cols = [right.row_position_snowflake_quoted_identifier]
 

tests/integ/modin/test_faster_pandas.py

@@ -279,6 +279,35 @@ def test_drop(session):
     assert_frame_equal(snow_result, native_result)
 
 
+@sql_count_checker(query_count=3, join_count=2)
+def test_drop_duplicates(session):
+    # create tables
+    table_name = Utils.random_name_for_temp_object(TempObjectType.TABLE)
+    session.create_dataframe(
+        native_pd.DataFrame([[2, 12], [2, 12], [3, 13]], columns=["A", "B"])
+    ).write.save_as_table(table_name, table_type="temp")
+
+    # create snow dataframes
+    df = pd.read_snowflake(table_name)
+    snow_result = df.drop_duplicates()
+
+    # verify that the input dataframe has a populated relaxed query compiler
+    assert df._query_compiler._relaxed_query_compiler is not None
+    assert df._query_compiler._relaxed_query_compiler._dummy_row_pos_mode is True
+    # verify that the output dataframe also has a populated relaxed query compiler
+    assert snow_result._query_compiler._relaxed_query_compiler is not None
+    assert (
+        snow_result._query_compiler._relaxed_query_compiler._dummy_row_pos_mode is True
+    )
+
+    # create pandas dataframes
+    native_df = df.to_pandas()
+    native_result = native_df.drop_duplicates()
+
+    # compare results
+    assert_frame_equal(snow_result, native_result)
+
+
 @sql_count_checker(query_count=3, join_count=1)
 def test_duplicated(session):
     # create tables