SNOW-2467959: Add support for 9 date functions in faster pandas

Author: sfc-gh-helmeleegy

CHANGELOG.md

@@ -57,7 +57,15 @@
   - `str.lower`
   - `str.upper`
   - `str.title`
-  
+  - `dt.tz_localize`
+  - `dt.tz_convert`
+  - `dt.ceil`
+  - `dt.round`
+  - `dt.floor`
+  - `dt.normalize`
+  - `dt.month_name`
+  - `dt.day_name`
+  - `dt.strftime`  
 - Make faster pandas disabled by default (opt-in instead of opt-out).
 
 ## 1.41.0 (2025-10-23)

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

@@ -20987,6 +20987,34 @@ def dt_tz_localize(
         ambiguous: str = "raise",
         nonexistent: str = "raise",
         include_index: bool = False,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _dt_tz_localize_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._dt_tz_localize_internal(
+                    tz=tz,
+                    ambiguous=ambiguous,
+                    nonexistent=nonexistent,
+                    include_index=include_index,
+                )
+            )
+        qc = self._dt_tz_localize_internal(
+            tz=tz,
+            ambiguous=ambiguous,
+            nonexistent=nonexistent,
+            include_index=include_index,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _dt_tz_localize_internal(
+        self,
+        tz: Union[str, tzinfo],
+        ambiguous: str = "raise",
+        nonexistent: str = "raise",
+        include_index: bool = False,
     ) -> "SnowflakeQueryCompiler":
         """
         Localize tz-naive to tz-aware.
@@ -21027,6 +21055,28 @@ def dt_tz_convert(
         self,
         tz: Union[str, tzinfo],
         include_index: bool = False,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _dt_tz_convert_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._dt_tz_convert_internal(
+                    tz=tz,
+                    include_index=include_index,
+                )
+            )
+        qc = self._dt_tz_convert_internal(
+            tz=tz,
+            include_index=include_index,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _dt_tz_convert_internal(
+        self,
+        tz: Union[str, tzinfo],
+        include_index: bool = False,
     ) -> "SnowflakeQueryCompiler":
         """
         Convert time-series data to the specified time zone.
@@ -21060,6 +21110,32 @@ def dt_ceil(
         ambiguous: str = "raise",
         nonexistent: str = "raise",
         include_index: bool = False,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _dt_ceil_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._dt_ceil_internal(
+                freq=freq,
+                ambiguous=ambiguous,
+                nonexistent=nonexistent,
+                include_index=include_index,
+            )
+        qc = self._dt_ceil_internal(
+            freq=freq,
+            ambiguous=ambiguous,
+            nonexistent=nonexistent,
+            include_index=include_index,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _dt_ceil_internal(
+        self,
+        freq: Frequency,
+        ambiguous: str = "raise",
+        nonexistent: str = "raise",
+        include_index: bool = False,
     ) -> "SnowflakeQueryCompiler":
         """
         Args:
@@ -21143,6 +21219,32 @@ def dt_round(
         ambiguous: str = "raise",
         nonexistent: str = "raise",
         include_index: bool = False,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _dt_round_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._dt_round_internal(
+                freq=freq,
+                ambiguous=ambiguous,
+                nonexistent=nonexistent,
+                include_index=include_index,
+            )
+        qc = self._dt_round_internal(
+            freq=freq,
+            ambiguous=ambiguous,
+            nonexistent=nonexistent,
+            include_index=include_index,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _dt_round_internal(
+        self,
+        freq: Frequency,
+        ambiguous: str = "raise",
+        nonexistent: str = "raise",
+        include_index: bool = False,
     ) -> "SnowflakeQueryCompiler":
         """
         Args:
@@ -21304,6 +21406,32 @@ def dt_floor(
         ambiguous: str = "raise",
         nonexistent: str = "raise",
         include_index: bool = False,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _dt_floor_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._dt_floor_internal(
+                freq=freq,
+                ambiguous=ambiguous,
+                nonexistent=nonexistent,
+                include_index=include_index,
+            )
+        qc = self._dt_floor_internal(
+            freq=freq,
+            ambiguous=ambiguous,
+            nonexistent=nonexistent,
+            include_index=include_index,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _dt_floor_internal(
+        self,
+        freq: Frequency,
+        ambiguous: str = "raise",
+        nonexistent: str = "raise",
+        include_index: bool = False,
     ) -> "SnowflakeQueryCompiler":
         """
         Args:
@@ -21374,6 +21502,22 @@ def floor_func(column: SnowparkColumn) -> SnowparkColumn:
         )
 
     def dt_normalize(self, include_index: bool = False) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _dt_normalize_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._dt_normalize_internal(
+                    include_index=include_index
+                )
+            )
+        qc = self._dt_normalize_internal(include_index=include_index)
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _dt_normalize_internal(
+        self, include_index: bool = False
+    ) -> "SnowflakeQueryCompiler":
         """
         Set the time component of each date-time value to midnight.
 
@@ -21398,6 +21542,26 @@ def normalize_column(column: SnowparkColumn) -> SnowparkColumn:
 
     def dt_month_name(
         self, locale: Optional[str] = None, include_index: bool = False
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _dt_month_name_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._dt_month_name_internal(
+                    locale=locale,
+                    include_index=include_index,
+                )
+            )
+        qc = self._dt_month_name_internal(
+            locale=locale,
+            include_index=include_index,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _dt_month_name_internal(
+        self, locale: Optional[str] = None, include_index: bool = False
     ) -> "SnowflakeQueryCompiler":
         """
         Args:
@@ -21432,6 +21596,24 @@ def month_name_func(column: SnowparkColumn) -> SnowparkColumn:
 
     def dt_day_name(
         self, locale: Optional[str] = None, include_index: bool = False
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _dt_day_name_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._dt_day_name_internal(
+                locale=locale,
+                include_index=include_index,
+            )
+        qc = self._dt_day_name_internal(
+            locale=locale,
+            include_index=include_index,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _dt_day_name_internal(
+        self, locale: Optional[str] = None, include_index: bool = False
     ) -> "SnowflakeQueryCompiler":
         """
         Args:
@@ -21487,6 +21669,18 @@ def dt_total_seconds(self, include_index: bool = False) -> "SnowflakeQueryCompil
         )
 
     def dt_strftime(self, date_format: str) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _dt_strftime_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._dt_strftime_internal(
+                date_format=date_format
+            )
+        qc = self._dt_strftime_internal(date_format=date_format)
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _dt_strftime_internal(self, date_format: str) -> "SnowflakeQueryCompiler":
         """
         Format underlying date-time data using specified format.