SNOW-1632900, SNOW-1625232: Add support for Series.dt.normalize and D…

…atetimeIndex.normalize

CHANGELOG.md

@@ -35,6 +35,7 @@
 - Added support for `DatetimeIndex.month_name` and `DatetimeIndex.day_name`.
 - Added support for `Series.dt.weekday`, `Series.dt.time`, and `DatetimeIndex.time`.
 - Added support for subtracting two timestamps to get a Timedelta.
+- Added support for `Series.dt.normalize`, and `DatetimeIndex.normalize`.
 
 #### Bug Fixes
 

docs/source/modin/series.rst

@@ -270,6 +270,7 @@ Series
     Series.dt.floor
     Series.dt.ceil
     Series.dt.round
+    Series.dt.normalize
 
 
 .. rubric:: String accessor methods

docs/source/modin/supported/series_dt_supported.rst

@@ -84,7 +84,7 @@ the method in the left column.
 +-----------------------------+---------------------------------+----------------------------------------------------+
 | ``tz_convert``              | N                               |                                                    |
 +-----------------------------+---------------------------------+----------------------------------------------------+
-| ``normalize``               | N                               |                                                    |
+| ``normalize``               | Y                               |                                                    |
 +-----------------------------+---------------------------------+----------------------------------------------------+
 | ``strftime``                | N                               |                                                    |
 +-----------------------------+---------------------------------+----------------------------------------------------+

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

@@ -16233,17 +16233,32 @@ def dt_floor(
             )
         )
 
-    def dt_normalize(self) -> None:
+    def dt_normalize(self, include_index: bool = False) -> "SnowflakeQueryCompiler":
         """
         Set the time component of each date-time value to midnight.
 
+        Args:
+            include_index: Whether to include the index columns in the operation.
+
         Returns
         -------
         BaseQueryCompiler
             New QueryCompiler containing date-time values with midnight time.
         """
-        ErrorMessage.not_implemented(
-            "Snowpark pandas doesn't yet support the method 'Series.dt.normalize'"
+        internal_frame = self._modin_frame
+
+        def normalize_column(col_id: str) -> SnowparkColumn:
+            return builtin("date_trunc")("d", col(col_id))
+
+        snowflake_ids = internal_frame.data_column_snowflake_quoted_identifiers[0:1]
+        if include_index:
+            snowflake_ids.extend(
+                internal_frame.index_column_snowflake_quoted_identifiers
+            )
+        return SnowflakeQueryCompiler(
+            internal_frame.update_snowflake_quoted_identifiers_with_expressions(
+                {col_id: normalize_column(col_id) for col_id in snowflake_ids}
+            ).frame
         )
 
     def dt_month_name(

src/snowflake/snowpark/modin/plugin/docstrings/series_utils.py

@@ -957,7 +957,45 @@ def swapcase():
         pass
 
     def normalize():
-        pass
+        """
+        Convert times to midnight.
+
+
+        The time component of the date-time is converted to midnight i.e. 00:00:00. This is useful in cases, when the time does not matter. Length is unaltered. The timezones are unaffected.
+
+
+        This method is available on Series with datetime values under the .dt accessor, and directly on Datetime Array/Index.
+
+
+        Returns
+        -------
+        DatetimeArray, DatetimeIndex or Series
+            The same type as the original data. Series will have the same name and index. DatetimeIndex will have the same name.
+
+        See also
+        --------
+        floor
+            Floor the datetimes to the specified freq.
+        ceil
+            Ceil the datetimes to the specified freq.
+        round
+            Round the datetimes to the specified freq.
+
+        Examples
+        --------
+        >>> idx = pd.date_range(start='2014-08-01 10:00', freq='h',
+        ...                    periods=3, tz='Asia/Calcutta')
+        >>> idx
+        DatetimeIndex(['2014-08-01 10:00:00+05:30',
+                    '2014-08-01 11:00:00+05:30',
+                    '2014-08-01 12:00:00+05:30'],
+                        dtype='datetime64[ns, Asia/Calcutta]', freq=None)
+        >>> idx.normalize()
+        DatetimeIndex(['2014-08-01 00:00:00+05:30',
+                    '2014-08-01 00:00:00+05:30',
+                    '2014-08-01 00:00:00+05:30'],
+                    dtype='datetime64[ns, Asia/Calcutta]', freq=None)
+        """
 
     def translate():
         """

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

@@ -787,7 +787,6 @@ def indexer_between_time(
         array([0, 1])
         """
 
-    @datetime_index_not_implemented()
     def normalize(self) -> DatetimeIndex:
         """
         Convert times to midnight.
@@ -814,18 +813,21 @@ def normalize(self) -> DatetimeIndex:
         Examples
         --------
         >>> idx = pd.date_range(start='2014-08-01 10:00', freq='h',
-        ...                     periods=3, tz='Asia/Calcutta')  # doctest: +SKIP
-        >>> idx  # doctest: +SKIP
+        ...                     periods=3, tz='Asia/Calcutta')
+        >>> idx
         DatetimeIndex(['2014-08-01 10:00:00+05:30',
                        '2014-08-01 11:00:00+05:30',
                        '2014-08-01 12:00:00+05:30'],
                         dtype='datetime64[ns, Asia/Calcutta]', freq=None)
-        >>> idx.normalize()  # doctest: +SKIP
+        >>> idx.normalize()
         DatetimeIndex(['2014-08-01 00:00:00+05:30',
                        '2014-08-01 00:00:00+05:30',
                        '2014-08-01 00:00:00+05:30'],
                        dtype='datetime64[ns, Asia/Calcutta]', freq=None)
         """
+        return DatetimeIndex(
+            query_compiler=self._query_compiler.dt_normalize(include_index=True)
+        )
 
     @datetime_index_not_implemented()
     def strftime(self, date_format: str) -> np.ndarray[np.object_]:

tests/integ/modin/index/test_datetime_index_methods.py

@@ -205,3 +205,15 @@ def test_day_month_name_negative(method):
     msg = f"Snowpark pandas method DatetimeIndex.{method} does not yet support the 'locale' parameter"
     with pytest.raises(NotImplementedError, match=msg):
         getattr(snow_index, method)(locale="pt_BR.utf8")
+
+
+@sql_count_checker(query_count=1)
+def test_normalize():
+    native_index = native_pd.date_range(start="2021-01-01", periods=5, freq="7h")
+    native_index = native_index.append(native_pd.DatetimeIndex([pd.NaT]))
+    snow_index = pd.DatetimeIndex(native_index)
+    eval_snowpark_pandas_result(
+        snow_index,
+        native_index,
+        lambda i: i.normalize(),
+    )

tests/integ/modin/series/test_dt_accessor.py

@@ -159,6 +159,19 @@ def test_floor_ceil_round_negative(func, freq, ambiguous, nonexistent):
         )
 
 
+@sql_count_checker(query_count=1)
+def test_normalize():
+    date_range = native_pd.date_range(start="2021-01-01", periods=5, freq="7h")
+    native_ser = native_pd.Series(date_range)
+    native_ser.iloc[2] = native_pd.NaT
+    snow_ser = pd.Series(native_ser)
+    eval_snowpark_pandas_result(
+        snow_ser,
+        native_ser,
+        lambda s: s.dt.normalize(),
+    )
+
+
 def test_isocalendar():
     with SqlCounter(query_count=1):
         date_range = native_pd.date_range("2020-05-01", periods=5, freq="4D")