ceil / floor / round + tests

databricks/koalas/indexes/datetimes.py

@@ -431,3 +431,99 @@ def days_in_month(self) -> Index:
         return Index(self.to_series().dt.days_in_month)
 
     days_in_month.__doc__ = daysinmonth.__doc__
+
+    # Methods
+    def ceil(self, freq, *args, **kwargs) -> "DatetimeIndex":
+        """
+        Perform floor operation on the data to the specified freq.
+
+        Parameters
+        ----------
+        freq : str or Offset
+            The frequency level to ceil the index to. Must be a fixed
+            frequency like 'S' (second) not 'ME' (month end).
+
+        Returns
+        -------
+        DatetimeIndex
+
+        Raises
+        ------
+        ValueError if the `freq` cannot be converted.
+
+        Examples
+        --------
+        >>> rng = ks.date_range('1/1/2018 11:59:00', periods=3, freq='min')
+        >>> rng.ceil('H')  # doctest: +NORMALIZE_WHITESPACE
+        DatetimeIndex(['2018-01-01 12:00:00', '2018-01-01 12:00:00',
+                       '2018-01-01 13:00:00'],
+                      dtype='datetime64[ns]', freq=None)
+        """
+        disallow_nanoseconds(freq)
+
+        return DatetimeIndex(self.to_series().dt.ceil(freq, *args, **kwargs))
+
+    def floor(self, freq, *args, **kwargs) -> "DatetimeIndex":
+        """
+        Perform floor operation on the data to the specified freq.
+
+        Parameters
+        ----------
+        freq : str or Offset
+            The frequency level to floor the index to. Must be a fixed
+            frequency like 'S' (second) not 'ME' (month end).
+
+        Returns
+        -------
+        DatetimeIndex
+
+        Raises
+        ------
+        ValueError if the `freq` cannot be converted.
+
+        Examples
+        --------
+        >>> rng = ks.date_range('1/1/2018 11:59:00', periods=3, freq='min')
+        >>> rng.floor("H")  # doctest: +NORMALIZE_WHITESPACE
+        DatetimeIndex(['2018-01-01 11:00:00', '2018-01-01 12:00:00',
+                       '2018-01-01 12:00:00'],
+                      dtype='datetime64[ns]', freq=None)
+        """
+        disallow_nanoseconds(freq)
+
+        return DatetimeIndex(self.to_series().dt.floor(freq, *args, **kwargs))
+
+    def round(self, freq, *args, **kwargs) -> "DatetimeIndex":
+        """
+        Perform round operation on the data to the specified freq.
+
+        Parameters
+        ----------
+        freq : str or Offset
+            The frequency level to round the index to. Must be a fixed
+            frequency like 'S' (second) not 'ME' (month end).
+
+        Returns
+        -------
+        DatetimeIndex
+
+        Raises
+        ------
+        ValueError if the `freq` cannot be converted.
+
+        Examples
+        --------
+        >>> rng = ks.date_range('1/1/2018 11:59:00', periods=3, freq='min')
+        >>> rng.round("H")  # doctest: +NORMALIZE_WHITESPACE
+        DatetimeIndex(['2018-01-01 12:00:00', '2018-01-01 12:00:00',
+                       '2018-01-01 12:00:00'],
+                      dtype='datetime64[ns]', freq=None)
+        """
+        disallow_nanoseconds(freq)
+
+        return DatetimeIndex(self.to_series().dt.round(freq, *args, **kwargs))
+
+
+def disallow_nanoseconds(freq):
+    if freq in ["N", "ns"]:
+        raise ValueError("nanoseconds is not supported")

databricks/koalas/tests/indexes/test_datetime.py

@@ -85,3 +85,28 @@ def test_properties(self):
             if LooseVersion(pd.__version__) >= LooseVersion("1.2.0"):
                 self.assert_eq(kidx.day_of_year, pidx.day_of_year)
                 self.assert_eq(kidx.day_of_week, pidx.day_of_week)
+
+    def test_ceil(self):
+        for kidx, pidx in self.idx_pairs:
+            for freq in self.fixed_freqs:
+                self.assert_eq(kidx.ceil(freq), pidx.ceil(freq))
+
+        self.disallow_nanoseconds(self.kidxs[0].ceil)
+
+    def test_floor(self):
+        for kidx, pidx in self.idx_pairs:
+            for freq in self.fixed_freqs:
+                self.assert_eq(kidx.floor(freq), pidx.floor(freq))
+
+        self.disallow_nanoseconds(self.kidxs[0].floor)
+
+    def test_round(self):
+        for kidx, pidx in self.idx_pairs:
+            for freq in self.fixed_freqs:
+                self.assert_eq(kidx.round(freq), pidx.round(freq))
+
+        self.disallow_nanoseconds(self.kidxs[0].round)
+
+    def disallow_nanoseconds(self, f):
+        self.assertRaises(ValueError, lambda: f(freq="ns"))
+        self.assertRaises(ValueError, lambda: f(freq="N"))