Skip to content

Commit

Permalink
Fix Index.repeat for datetime64 types (#15722)
Browse files Browse the repository at this point in the history
Fixes: #15720 

This PR fixes `Index.repeat` where the `freq` of `DatetimeIndex` needs to be reset.

Authors:
  - GALI PREM SAGAR (https://github.com/galipremsagar)

Approvers:
  - Matthew Roeschke (https://github.com/mroeschke)

URL: #15722
  • Loading branch information
galipremsagar authored May 14, 2024
1 parent 4069c82 commit cbe2775
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 1 deletion.
5 changes: 5 additions & 0 deletions python/cudf/cudf/core/index.py
Original file line number Diff line number Diff line change
Expand Up @@ -2373,6 +2373,11 @@ def tz_convert(self, tz: str | None):
result_col = self._column.tz_convert(tz)
return DatetimeIndex._from_data({self.name: result_col})

def repeat(self, repeats, axis=None):
res = super().repeat(repeats, axis=axis)
res._freq = None
return res


class TimedeltaIndex(Index):
"""
Expand Down
5 changes: 4 additions & 1 deletion python/cudf/cudf/core/indexed_frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -4871,13 +4871,16 @@ def repeat(self, repeats, axis=None):
1 2
dtype: int64
"""
return self._from_columns_like_self(
res = self._from_columns_like_self(
Frame._repeat(
[*self._index._data.columns, *self._columns], repeats, axis
),
self._column_names,
self._index_names,
)
if isinstance(res.index, cudf.DatetimeIndex):
res.index._freq = None
return res

def astype(
self,
Expand Down
14 changes: 14 additions & 0 deletions python/cudf/cudf/tests/test_index.py
Original file line number Diff line number Diff line change
Expand Up @@ -3252,3 +3252,17 @@ def test_Index_init_with_nans():
assert gi.dtype == np.dtype("float64")
pi = pd.Index([1, 2, 3, np.nan])
assert_eq(pi, gi)


def test_index_datetime_repeat():
gidx = cudf.date_range("2021-01-01", periods=3, freq="D")
pidx = gidx.to_pandas()

actual = gidx.repeat(5)
expected = pidx.repeat(5)

assert_eq(actual, expected)

actual = gidx.to_frame().repeat(5)

assert_eq(actual.index, expected)

0 comments on commit cbe2775

Please sign in to comment.