Merge branch 'main' into feature/aherrera/SNOW-2455523

Author: sfc-gh-aherreraaguilar

CHANGELOG.md

@@ -127,7 +127,11 @@
   - `expanding.std`
   - `expanding.var`
   - `expanding.sem`
+  - `cumsum`
+  - `cummin`
+  - `cummax`
 - Make faster pandas disabled by default (opt-in instead of opt-out).
+- Improve performance of `drop_duplicates` by avoiding joins when `keep!=False` in faster pandas.
 
 ## 1.42.0 (2025-10-28)
 

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

@@ -9105,6 +9105,26 @@ def _concat_internal(
     )
     def cumsum(
         self, axis: int = 0, skipna: bool = True, *args: Any, **kwargs: Any
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _cumsum_internal to be supported in faster pandas.
+        """
+        relaxed_query_compiler = None
+        if self._relaxed_query_compiler is not None:
+            relaxed_query_compiler = self._relaxed_query_compiler._cumsum_internal(
+                axis=axis,
+                skipna=skipna,
+                **kwargs,
+            )
+        qc = self._cumsum_internal(
+            axis=axis,
+            skipna=skipna,
+            **kwargs,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _cumsum_internal(
+        self, axis: int = 0, skipna: bool = True, *args: Any, **kwargs: Any
     ) -> "SnowflakeQueryCompiler":
         """
         Return cumulative sum over a DataFrame or Series axis.
@@ -9143,6 +9163,26 @@ def cumsum(
     )
     def cummin(
         self, axis: int = 0, skipna: bool = True, *args: Any, **kwargs: Any
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _cummin_internal to be supported in faster pandas.
+        """
+        relaxed_query_compiler = None
+        if self._relaxed_query_compiler is not None:
+            relaxed_query_compiler = self._relaxed_query_compiler._cummin_internal(
+                axis=axis,
+                skipna=skipna,
+                **kwargs,
+            )
+        qc = self._cummin_internal(
+            axis=axis,
+            skipna=skipna,
+            **kwargs,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _cummin_internal(
+        self, axis: int = 0, skipna: bool = True, *args: Any, **kwargs: Any
     ) -> "SnowflakeQueryCompiler":
         """
         Return cumulative min over a DataFrame or Series axis.
@@ -9181,6 +9221,26 @@ def cummin(
     )
     def cummax(
         self, axis: int = 0, skipna: bool = True, *args: Any, **kwargs: Any
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _cummax_internal to be supported in faster pandas.
+        """
+        relaxed_query_compiler = None
+        if self._relaxed_query_compiler is not None:
+            relaxed_query_compiler = self._relaxed_query_compiler._cummax_internal(
+                axis=axis,
+                skipna=skipna,
+                **kwargs,
+            )
+        qc = self._cummax_internal(
+            axis=axis,
+            skipna=skipna,
+            **kwargs,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _cummax_internal(
+        self, axis: int = 0, skipna: bool = True, *args: Any, **kwargs: Any
     ) -> "SnowflakeQueryCompiler":
         """
         Return cumulative max over a DataFrame or Series axis.
@@ -18622,6 +18682,31 @@ def add_substring(
         # Returning the query compiler with updated columns and index.
         return SnowflakeQueryCompiler(result_frame)
 
+    def drop_duplicates(self) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _drop_duplicates_internal to be supported in faster pandas.
+        """
+        relaxed_query_compiler = None
+        if self._relaxed_query_compiler is not None:
+            relaxed_query_compiler = (
+                self._relaxed_query_compiler._drop_duplicates_internal()
+            )
+        qc = self._drop_duplicates_internal()
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _drop_duplicates_internal(self) -> "SnowflakeQueryCompiler":
+        """
+        Return a DataFrame or Series after dropping the duplicate rows.
+        """
+        return self.groupby_agg(
+            by=self._modin_frame.data_column_pandas_labels,
+            agg_func={},
+            axis=0,
+            groupby_kwargs={"sort": False, "as_index": False, "dropna": False},
+            agg_args=[],
+            agg_kwargs={},
+        )
+
     def duplicated(
         self,
         subset: Union[Hashable, Sequence[Hashable]] = None,

src/snowflake/snowpark/modin/plugin/extensions/base_overrides.py

@@ -1529,8 +1529,12 @@ def drop_duplicates(
         df = self[subset]
     else:
         df = self
-    duplicated = df.duplicated(keep=keep)
-    result = self[~duplicated]
+    if pd.session.dummy_row_pos_optimization_enabled and keep in ["first", "last"]:
+        result_qc = df._query_compiler.drop_duplicates()
+        result = self.__constructor__(query_compiler=result_qc)
+    else:
+        duplicated = df.duplicated(keep=keep)
+        result = self[~duplicated]
     if ignore_index:
         result.index = pandas.RangeIndex(stop=len(result))
     if inplace:

tests/integ/modin/test_faster_pandas.py

@@ -340,6 +340,40 @@ def test_copy(session):
         assert_frame_equal(snow_result, native_result)
 
 
+@pytest.mark.parametrize("func", ["cumsum", "cummin", "cummax"])
+@sql_count_checker(query_count=3)
+def test_cumulative_functions(session, func):
+    with session_parameter_override(
+        session, "dummy_row_pos_optimization_enabled", True
+    ):
+        # create tables
+        table_name = Utils.random_name_for_temp_object(TempObjectType.TABLE)
+        session.create_dataframe(
+            native_pd.DataFrame([[1, 11], [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 = getattr(df["B"], func)()
+
+        # 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 = getattr(native_df["B"], func)()
+
+        # compare results
+        assert_series_equal(snow_result, native_result)
+
+
 @sql_count_checker(query_count=3)
 def test_drop(session):
     with session_parameter_override(
@@ -373,15 +407,15 @@ def test_drop(session):
         assert_frame_equal(snow_result, native_result)
 
 
-@sql_count_checker(query_count=3, join_count=2)
+@sql_count_checker(query_count=3)
 def test_drop_duplicates(session):
     with session_parameter_override(
         session, "dummy_row_pos_optimization_enabled", True
     ):
         # 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"])
+            native_pd.DataFrame([[2, 12], [3, 13], [2, 12]], columns=["A", "B"])
         ).write.save_as_table(table_name, table_type="temp")
 
         # create snow dataframes
@@ -403,7 +437,7 @@ def test_drop_duplicates(session):
         native_result = native_df.drop_duplicates()
 
         # compare results
-        assert_frame_equal(snow_result, native_result)
+        assert_frame_equal(snow_result, native_result, check_index_type=False)
 
 
 @sql_count_checker(query_count=3, join_count=1)