Adding feature Truncate to DataFrame and Series

python/cudf/cudf/core/dataframe.py

@@ -2886,6 +2886,112 @@ def axes(self):
         """
         return [self._index, self._data.to_pandas_index()]
 
+    @_cudf_nvtx_annotate
+    def truncate(self, before=None, after=None, axis=0, copy=True):
+        """
+        Truncate a DataFrame before and after some index value.
+        This is a useful shorthand for boolean indexing based on index
+        values above or below certain thresholds.
+
+        Parameters
+        ----------
+        before : date, str, int
+            Truncate all rows before this index value.
+        after : date, str, int
+            Truncate all rows after this index value.
+        axis : {0 or 'index', 1 or 'columns'}, optional
+            Axis to truncate. Truncates the index (rows) by default.
+        copy : bool, default is True,
+            Return a copy of the truncated section.
+
+        Returns
+        -------
+            The truncated Series or DataFrame.
+
+        Notes
+        -----
+        If the index being truncated contains only datetime values,
+        `before` and `after` may be specified as strings instead of
+        Timestamps.
+
+        Examples
+        --------
+        >>> cdf1 = cudf.DataFrame({'A': ['a', 'b', 'c', 'd', 'e'],
+        ...                       'B': ['f', 'g', 'h', 'i', 'j'],
+        ...                       'C': ['k', 'l', 'm', 'n', 'o']},
+        ...                             index=[1, 2, 3, 4, 5])
+        >>> cdf1
+            A	B	C
+        1	a	f	k
+        2	b	g	l
+        3	c	h	m
+        4	d	i	n
+        5	e	j	o
+
+        >>> cdf1.truncate(before=2, after=4)
+            A	B	C
+        2	b	g	l
+        3	c	h	m
+        4	d	i	n
+
+        >>> cdf1.truncate(before="A", after="B", axis="columns")
+            A	B
+        1	a	f
+        2	b	g
+        3	c	h
+        4	d	i
+        5	e	j
+
+        >>> import cudf
+        >>> dates = cudf.date_range('2021-01-01', '2021-01-02', freq='s')
+        >>> cdf2 = cudf.DataFrame(data={'A': 1, 'B': 2}, index=dates)
+        >>> cdf2.head()
+                            A	B
+        2021-01-01 00:00:00	1	2
+        2021-01-01 00:00:01	1	2
+        2021-01-01 00:00:02	1	2
+        2021-01-01 00:00:03	1	2
+        2021-01-01 00:00:04	1	2
+
+        >>> cdf2.truncate(
+        ... before="2021-01-01 23:45:18", after="2021-01-01 23:45:27")
+                            A	B
+        2021-01-01 23:45:18	1	2
+        2021-01-01 23:45:19	1	2
+        2021-01-01 23:45:20	1	2
+        2021-01-01 23:45:21	1	2
+        2021-01-01 23:45:22	1	2
+        2021-01-01 23:45:23	1	2
+        2021-01-01 23:45:24	1	2
+        2021-01-01 23:45:25	1	2
+        2021-01-01 23:45:26	1	2
+        2021-01-01 23:45:27	1	2
+        """
+        axis = self._get_axis_from_axis_arg(axis)
+        ax = self._index if axis == 0 else self.columns
+
+        if not ax.is_monotonic_increasing and not ax.is_monotonic_decreasing:
+            raise ValueError("truncate requires a sorted index")
+
+        if type(ax) is cudf.core.index.DatetimeIndex:
+            before = pd.to_datetime(before)
+            after = pd.to_datetime(after)
+
+        if before is not None and after is not None and before > after:
+            raise ValueError(f"Truncate: {after} must be after {before}")
+
+        if len(ax) > 1 and ax.is_monotonic_decreasing and ax.nunique() > 1:
+            before, after = after, before
+
+        slicer = [slice(None, None)] * self.ndim
+        slicer[axis] = slice(before, after)
+        result = self.loc[tuple(slicer)]
+
+        if copy:
+            result = result.copy()
+
+        return result
+
     def diff(self, periods=1, axis=0):
         """
         First discrete difference of element.

python/cudf/cudf/core/series.py

@@ -661,6 +661,115 @@ def axes(self):
         """
         return [self.index]
 
+    @_cudf_nvtx_annotate
+    def truncate(self, before=None, after=None, axis=0, copy=True):
+        """
+        Truncate a Series before and after some index value.
+        This is a useful shorthand for boolean indexing based on index
+        values above or below certain thresholds.
+        Parameters
+        ----------
+        before : date, str, int
+            Truncate all rows before this index value.
+        after : date, str, int
+            Truncate all rows after this index value.
+        axis : {0 or 'index'}, optional
+            Axis to truncate. Truncates the index (rows) by default.
+        copy : bool, default is True,
+            Return a copy of the truncated section.
+        Returns
+        -------
+            The truncated Series.
+        Notes
+        -----
+        If the index being truncated contains only datetime values,
+        `before` and `after` may be specified as strings instead of
+        Timestamps.
+        Examples
+        --------
+        >>> import cudf
+        >>> cs1 = cudf.Series([1, 2, 3, 4])
+        >>> cs1
+        0    1
+        1    2
+        2    3
+        3    4
+        dtype: int64
+
+        >>> cs1.truncate(before=1, after=2)
+        1    2
+        2    3
+        dtype: int64
+
+        >>> import cudf
+        >>> dates = cudf.date_range('2021-01-01', '2021-01-02', freq='s')
+        >>> cs2 = cudf.Series(range(len(dates)), index=dates)
+        >>> cs2
+        2021-01-01 00:00:00        0
+        2021-01-01 00:00:01        1
+        2021-01-01 00:00:02        2
+        2021-01-01 00:00:03        3
+        2021-01-01 00:00:04        4
+        ...
+        2021-01-01 23:59:55    86395
+        2021-01-01 23:59:56    86396
+        2021-01-01 23:59:57    86397
+        2021-01-01 23:59:58    86398
+        2021-01-01 23:59:59    86399
+        Length: 86400, dtype: int64
+
+        >>> cs2.truncate(
+        ... before="2021-01-01 23:45:18", after="2021-01-01 23:45:27")
+        2021-01-01 23:45:18    85518
+        2021-01-01 23:45:19    85519
+        2021-01-01 23:45:20    85520
+        2021-01-01 23:45:21    85521
+        2021-01-01 23:45:22    85522
+        2021-01-01 23:45:23    85523
+        2021-01-01 23:45:24    85524
+        2021-01-01 23:45:25    85525
+        2021-01-01 23:45:26    85526
+        2021-01-01 23:45:27    85527
+        dtype: int64
+
+        >>> cs3 = cudf.Series({'A':1, 'B':2, 'C':3, 'D':4})
+        >>> cs3
+        A    1
+        B    2
+        C    3
+        D    4
+        dtype: int64
+
+        >>> cs3.truncate(before='B', after='C')
+        B    2
+        C    3
+        dtype: int64
+        """
+        if axis not in (0, "index"):
+            raise ValueError(f"No axis named {axis} for object type Series")
+        ax = self.index
+
+        if not ax.is_monotonic_increasing and not ax.is_monotonic_decreasing:
+            raise ValueError("truncate requires a sorted index")
+
+        if ax is cudf.core.index.DatetimeIndex:
+            before = pd.to_datetime(before)
+            after = pd.to_datetime(after)
+
+        if before is not None and after is not None and before > after:
+            raise ValueError(f"Truncate: {after} must be after {before}")
+
+        if len(ax) > 1 and ax.is_monotonic_decreasing and ax.nunique() > 1:
+            before, after = after, before
+
+        slicer = slice(before, after)
+        result = self.loc[slicer]
+
+        if copy:
+            result = result.copy()
+
+        return result
+
     @_cudf_nvtx_annotate
     def serialize(self):
         header, frames = super().serialize()

python/cudf/cudf/tests/test_dataframe.py

@@ -288,6 +288,54 @@ def test_axes(data):
         assert_eq(e, a)
 
 
+def test_dataframe_truncate_axis_0():
+    cdf1 = cudf.DataFrame(
+        {
+            "A": ["a", "b", "c", "d", "e"],
+            "B": ["f", "g", "h", "i", "j"],
+            "C": ["k", "l", "m", "n", "o"],
+        },
+        index=[1, 2, 3, 4, 5],
+    )
+    pdf1 = cdf1.to_pandas()
+    expected = pdf1.truncate(before=2, after=4, axis="index")
+    actual = cdf1.truncate(before=2, after=4, axis="index")
+
+    assert_eq(actual, expected)
+    assert_eq(pdf1.truncate(1, 4, 0), cdf1.truncate(1, 4, 0))
+
+
+def test_dataframe_truncate_axis_1():
+    cdf1 = cudf.DataFrame(
+        {
+            "A": ["a", "b", "c", "d", "e"],
+            "B": ["f", "g", "h", "i", "j"],
+            "C": ["k", "l", "m", "n", "o"],
+        },
+        index=[1, 2, 3, 4, 5],
+    )
+    pdf1 = cdf1.to_pandas()
+    expected = pdf1.truncate(before="A", after="B", axis="columns")
+    actual = cdf1.truncate(before="A", after="B", axis="columns")
+
+    assert_eq(actual, expected)
+    assert_eq(pdf1.truncate("A", "B", 1), cdf1.truncate("A", "B", 1))
+
+
+def test_dataframe_truncate_datetimeindex():
+    dates = cudf.date_range("2021-01-01", "2021-01-02", freq="s")
+    cdf1 = cudf.DataFrame(data={"A": 1, "B": 2}, index=dates)
+    pdf1 = cdf1.to_pandas()
+    expected = pdf1.truncate(
+        before="2021-01-01 23:45:18", after="2021-01-01 23:45:27"
+    )
+    actual = cdf1.truncate(
+        before="2021-01-01 23:45:18", after="2021-01-01 23:45:27"
+    )
+
+    assert_eq(actual, expected)
+
+
 def test_series_init_none():
 
     # test for creating empty series

python/cudf/cudf/tests/test_series.py

@@ -1614,6 +1614,37 @@ def test_axes(data):
         assert_eq(e, a)
 
 
+def test_series_truncate():
+    csr = cudf.Series([1, 2, 3, 4])
+    psr = csr.to_pandas()
+
+    assert_eq(csr.truncate(), psr.truncate())
+    assert_eq(csr.truncate(1, 2), psr.truncate(1, 2))
+    assert_eq(csr.truncate(before=1, after=2), psr.truncate(before=1, after=2))
+
+
+def test_series_truncate_invalid_axis():
+    csr = cudf.Series([1, 2, 3, 4])
+    msg = "No axis named 1 for object type Series"
+    with pytest.raises(ValueError, match=msg):
+        csr.truncate(axis=1)
+
+
+def test_series_truncate_datetimeindex():
+    dates = cudf.date_range("2021-01-01", "2021-01-02", freq="s")
+    csr = cudf.Series(range(len(dates)), index=dates)
+    psr = csr.to_pandas()
+
+    assert_eq(
+        csr.truncate(
+            before="2021-01-01 23:45:18", after="2021-01-01 23:45:27"
+        ),
+        psr.truncate(
+            before="2021-01-01 23:45:18", after="2021-01-01 23:45:27"
+        ),
+    )
+
+
 @pytest.mark.parametrize(
     "data",
     [