SNOW-2429645: Add support for drop in faster pandas (#3892)

Author: sfc-gh-helmeleegy

CHANGELOG.md

@@ -110,6 +110,7 @@
   - `sort_values`
   - `loc` (setting columns)
   - `to_datetime`
+  - `drop`
 - Reuse row count from the relaxed query compiler in `get_axis_len`.
 
 #### Bug Fixes

src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py

@@ -13084,6 +13084,32 @@ def drop(
         columns: Optional[Sequence[Hashable]] = None,
         level: Optional[Level] = None,
         errors: Literal["raise", "ignore"] = "raise",
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _drop_internal to be supported in faster pandas.
+        """
+        relaxed_query_compiler = None
+        if self._relaxed_query_compiler is not None and index is None:
+            relaxed_query_compiler = self._relaxed_query_compiler._drop_internal(
+                index=index,
+                columns=columns,
+                level=level,
+                errors=errors,
+            )
+        qc = self._drop_internal(
+            index=index,
+            columns=columns,
+            level=level,
+            errors=errors,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _drop_internal(
+        self,
+        index: Optional[Sequence[Hashable]] = None,
+        columns: Optional[Sequence[Hashable]] = None,
+        level: Optional[Level] = None,
+        errors: Literal["raise", "ignore"] = "raise",
     ) -> "SnowflakeQueryCompiler":
         """
         Drop specified rows or columns.

tests/integ/modin/test_faster_pandas.py

@@ -194,6 +194,35 @@ def test_read_filter_join_flag_disabled(session):
         assert_frame_equal(snow_result, native_result)
 
 
+@sql_count_checker(query_count=3)
+def test_drop(session):
+    # create tables
+    table_name = Utils.random_name_for_temp_object(TempObjectType.TABLE)
+    session.create_dataframe(
+        native_pd.DataFrame([[2, True], [1, False], [3, False]], 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(columns=["B"])
+
+    # 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(columns=["B"])
+
+    # compare results
+    assert_frame_equal(snow_result, native_result)
+
+
 @pytest.mark.parametrize("func", ["isna", "isnull", "notna", "notnull"])
 @sql_count_checker(query_count=3)
 def test_isna_notna(session, func):