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

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`.

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                               |                                  |                                           |
 +-----------------------------+---------------------------------+----------------------------------+-------------------------------------------+

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":

src/snowflake/snowpark/modin/plugin/_internal/resample_utils.py

@@ -8,6 +8,7 @@
 import pandas as native_pd
 from pandas._libs.lib import no_default
 from pandas._libs.tslibs import to_offset
+from pandas._typing import Frequency
 
 import snowflake.snowpark.modin.pandas as pd
 from snowflake.snowpark._internal.type_utils import ColumnOrName
@@ -108,15 +109,15 @@
 )
 
 
-def rule_to_snowflake_width_and_slice_unit(rule: str) -> tuple[int, str]:
+def rule_to_snowflake_width_and_slice_unit(rule: Frequency) -> tuple[int, str]:
     """
     Converts pandas resample bin rule to Snowflake's slice_width and slice_unit
     format.
 
     Parameters
     ----------
-    rule : str
-        The offset string representing resample bin size. For example: '1D', '2T', etc.
+    rule : Frequency
+        The offset or string representing resample bin size. For example: '1D', '2T', etc.
 
     Returns
     -------

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}")