SNOW-1637948: Add support for TimedeltaIndex floor, ceil and round

snowflakedb · Sep 5, 2024 · 71c6a2d · 71c6a2d
1 parent 4fae315
commit 71c6a2d
Show file tree

Hide file tree

Showing 7 changed files with 232 additions and 109 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -81,6 +81,7 @@
   - support for `GroupBy` filtrations `first` and `last`.
   - support for `TimedeltaIndex` attributes: `days`, `seconds`, `microseconds` and `nanoseconds`.
   - support for `diff` with timestamp columns on `axis=0` and `axis=1`
+  - support for `TimedeltaIndex` methods: `ceil`, `floor` and `round`.
 - Added support for index's arithmetic and comparison operators.
 - Added support for `Series.dt.round`.
 - Added documentation pages for `DatetimeIndex`.

diff --git a/docs/source/modin/supported/timedelta_index_supported.rst b/docs/source/modin/supported/timedelta_index_supported.rst
@@ -38,11 +38,11 @@ Methods
 +-----------------------------+---------------------------------+----------------------------------+-------------------------------------------+
 | ``to_pytimedelta``          | N                               |                                  |                                           |
 +-----------------------------+---------------------------------+----------------------------------+-------------------------------------------+
-| ``round``                   | N                               |                                  |                                           |
+| ``round``                   | Y                               |                                  |                                           |
 +-----------------------------+---------------------------------+----------------------------------+-------------------------------------------+
-| ``floor``                   | N                               |                                  |                                           |
+| ``floor``                   | Y                               |                                  |                                           |
 +-----------------------------+---------------------------------+----------------------------------+-------------------------------------------+
-| ``ceil``                    | N                               |                                  |                                           |
+| ``ceil``                    | Y                               |                                  |                                           |
 +-----------------------------+---------------------------------+----------------------------------+-------------------------------------------+
 | ``mean``                    | N                               |                                  |                                           |
 +-----------------------------+---------------------------------+----------------------------------+-------------------------------------------+
diff --git a/src/snowflake/snowpark/modin/plugin/_internal/frame.py b/src/snowflake/snowpark/modin/plugin/_internal/frame.py
@@ -1299,6 +1299,7 @@ def apply_snowpark_function_to_columns(
         self,
         snowpark_func: Callable[[Any], SnowparkColumn],
         include_index: bool = False,
+        return_type: Optional[SnowparkPandasType] = None,
     ) -> "InternalFrame":
         """
         Apply snowpark function callable to all data columns of an InternalFrame. If
@@ -1307,6 +1308,7 @@ def apply_snowpark_function_to_columns(
 
         Arguments:
             snowpark_func: Snowpark function to apply to columns of underlying snowpark df.
+            return_type: The optional SnowparkPandasType for the new column.
             include_index: Whether to apply the function to index columns as well.
 
         Returns:
@@ -1317,7 +1319,8 @@ def apply_snowpark_function_to_columns(
             snowflake_ids.extend(self.index_column_snowflake_quoted_identifiers)
 
         return self.update_snowflake_quoted_identifiers_with_expressions(
-            {col_id: snowpark_func(col(col_id)) for col_id in snowflake_ids}
+            {col_id: snowpark_func(col(col_id)) for col_id in snowflake_ids},
+            [return_type] * len(snowflake_ids) if return_type else None,
         ).frame
 
     def select_active_columns(self) -> "InternalFrame":

diff --git a/src/snowflake/snowpark/modin/plugin/_internal/timestamp_utils.py b/src/snowflake/snowpark/modin/plugin/_internal/timestamp_utils.py
@@ -11,7 +11,8 @@
 import numpy as np
 import pandas as native_pd
 from pandas._libs import lib
-from pandas._typing import DateTimeErrorChoices
+from pandas._libs.tslibs import to_offset
+from pandas._typing import DateTimeErrorChoices, Frequency
 from pandas.api.types import is_datetime64_any_dtype, is_float_dtype, is_integer_dtype
 
 from snowflake.snowpark import Column
@@ -168,11 +169,24 @@ def col_to_s(col: Column, unit: Literal["D", "s", "ms", "us", "ns"]) -> Column:
         return col / 10**9
 
 
+def timedelta_freq_to_nanos(freq: Frequency) -> int:
+    """
+    Convert a pandas frequency string to nanoseconds.
+
+    Args:
+        freq: Timedelta frequency string or offset.
+
+    Returns:
+        int: nanoseconds
+    """
+    return to_offset(freq).nanos
+
+
 def col_to_timedelta(col: Column, unit: str) -> Column:
     """
     Converts ``col`` (stored in the specified units) to timedelta nanoseconds.
     """
-    td_unit = VALID_PANDAS_TIMEDELTA_ABBREVS.get(unit)
+    td_unit = VALID_PANDAS_TIMEDELTA_ABBREVS.get(unit.lower())
     if not td_unit:
         # Same error as native pandas.
         raise ValueError(f"invalid unit abbreviation: {unit}")