Support more units in cudf.DateOffset

python/cudf/cudf/core/column/datetime.py

@@ -1,4 +1,5 @@
-# Copyright (c) 2019-2020, NVIDIA CORPORATION.
+# Copyright (c) 2019-2021, NVIDIA CORPORATION.
+
 from __future__ import annotations
 
 import datetime as dt
@@ -252,7 +253,7 @@ def binary_operator(
         reflect: bool = False,
     ) -> ColumnBase:
         if isinstance(rhs, cudf.DateOffset):
-            return binop_offset(self, rhs, op)
+            return rhs._datetime_binop(self, op, reflect=reflect)
         lhs, rhs = self, rhs
         if op in ("eq", "ne", "lt", "gt", "le", "ge"):
             out_dtype = np.bool

python/cudf/cudf/core/tools/datetimes.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2019-2020, NVIDIA CORPORATION.
+# Copyright (c) 2019-2021, NVIDIA CORPORATION.
 
 import warnings
 
@@ -7,6 +7,7 @@
 from pandas.core.tools.datetimes import _unit_map
 
 import cudf
+from cudf import _lib as libcudf
 from cudf._lib.strings.char_types import is_integer as cpp_is_integer
 from cudf.core import column
 from cudf.core.index import as_index
@@ -334,9 +335,8 @@ def get_units(value):
     return value
 
 
-class _DateOffsetScalars(object):
-    def __init__(self, scalars):
-        self._gpu_scalars = scalars
+class _DateOffsetScalars(dict):
+    pass
 
 
 class _UndoOffsetMeta(pd._libs.tslibs.offsets.OffsetMeta):
@@ -434,10 +434,6 @@ def __init__(self, n=1, normalize=False, **kwds):
                 "normalize not yet supported for DateOffset"
             )
 
-        # TODO: Pandas supports combinations
-        if len(kwds) > 1:
-            raise NotImplementedError("Multiple time units not yet supported")
-
         all_possible_kwargs = {
             "years",
             "months",
@@ -456,36 +452,101 @@ def __init__(self, n=1, normalize=False, **kwds):
             "minute",
             "second",
             "microsecond",
-            "millisecond" "nanosecond",
+            "millisecond",
+            "nanosecond",
         }
 
-        supported_kwargs = {"months"}
+        supported_kwargs = {
+            "years",
+            "months",
+            "weeks",
+            "days",
+            "hours",
+            "minutes",
+            "seconds",
+            "microseconds",
+            "nanoseconds",
+        }
+
+        super().__init__(n=n, normalize=normalize, **kwds)
+
+        kwds = self._combine_months_and_years(**kwds)
+        kwds = self._combine_timedeltas_to_nanos(**kwds)
 
         scalars = {}
         for k, v in kwds.items():
             if k in all_possible_kwargs:
                 # Months must be int16
-                dtype = "int16" if k == "months" else None
+                if k == "months":
+                    dtype = "int16"
+                elif k == "nanoseconds":
+                    dtype = "timedelta64[ns]"
+                else:
+                    dtype = None
                 scalars[k] = cudf.Scalar(v, dtype=dtype)
 
-        super().__init__(n=n, normalize=normalize, **kwds)
-
         wrong_kwargs = set(kwds.keys()).difference(supported_kwargs)
         if len(wrong_kwargs) > 0:
-            raise ValueError(
+            raise NotImplementedError(
                 f"Keyword arguments '{','.join(list(wrong_kwargs))}'"
                 " are not yet supported in cuDF DateOffsets"
             )
         self._scalars = _DateOffsetScalars(scalars)
 
-    def _generate_column(self, size, op):
-        months = self._scalars._gpu_scalars["months"]
+    def _combine_months_and_years(self, **kwargs):
+        kwargs["months"] = kwargs.pop("years", 0) * 12 + kwargs.pop(
+            "months", 0
+        )
+        return kwargs
+
+    def _combine_timedeltas_to_nanos(self, **kwargs):
+        weeks = kwargs.pop("weeks", 0)
+        days = kwargs.pop("days", 0) + 7 * weeks
+        hours = kwargs.pop("hours", 0) + days * 24
+        minutes = kwargs.pop("minutes", 0) + hours * 60
+        seconds = kwargs.pop("seconds", 0) + minutes * 60
+        milliseconds = kwargs.pop("milliseconds", 0) + seconds * 1000
+        microseconds = kwargs.pop("microseconds", 0) + milliseconds * 1000
+        nanoseconds = kwargs.pop("nanoseconds", 0) + microseconds * 1000
+        kwargs["nanoseconds"] = nanoseconds
+        return kwargs
+
+    def _datetime_binop(self, datetime_col, op, reflect=False):
+        if reflect and op == "sub":
+            raise TypeError(
+                f"Can not subtract a {type(datetime_col).__name__}"
+                f" from a {type(self).__name__}"
+            )
+        if op not in {"add", "sub"}:
+            raise TypeError(
+                f"{op} not supported between {type(self).__name__}"
+                f" and {type(datetime_col).__name__}"
+            )
+        if self._is_no_op:
+            return datetime_col
+        else:
+            if "months" in self._scalars:
+                rhs = self._generate_months_column(len(datetime_col), op)
+                datetime_col = libcudf.datetime.add_months(datetime_col, rhs)
+            if "nanoseconds" in self._scalars:
+                datetime_col = datetime_col + self._generate_nanos_column(
+                    len(datetime_col), op
+                )
+            return datetime_col
+
+    def _generate_months_column(self, size, op):
+        months = self._scalars["months"]
         months = -months if op == "sub" else months
         # TODO: pass a scalar instead of constructing a column
         # https://github.com/rapidsai/cudf/issues/6990
         col = cudf.core.column.as_column(months, length=size)
         return col
 
+    def _generate_nanos_column(self, size, op):
+        nanos = self._scalars["nanoseconds"]
+        nanos = -nanos if op == "sub" else nanos
+        return cudf.core.column.as_column(nanos, length=size)
+
     @property
     def _is_no_op(self):
         # some logic could be implemented here for more complex cases
@@ -497,3 +558,7 @@ def __setattr__(self, name, value):
             raise AttributeError("DateOffset objects are immutable.")
         else:
             object.__setattr__(self, name, value)
+
+    def __neg__(self):
+        new_scalars = {k: -v for k, v in self.kwds.items()}
+        return DateOffset(**new_scalars)

python/cudf/cudf/tests/test_binops.py

@@ -1474,12 +1474,32 @@ def test_scalar_power_invalid(dtype_l, dtype_r):
     ],
 )
 @pytest.mark.parametrize("n_periods", [0, 1, -1, 12, -12])
-@pytest.mark.parametrize("frequency", ["months"])
+@pytest.mark.parametrize(
+    "frequency",
+    [
+        "months",
+        "years",
+        "days",
+        "hours",
+        "minutes",
+        "seconds",
+        "microseconds",
+        pytest.param(
+            "nanoseconds",
+            marks=pytest.mark.xfail(
+                reason="https://github.com/pandas-dev/pandas/issues/36589"
+            ),
+        ),
+    ],
+)
 @pytest.mark.parametrize(
     "dtype",
     ["datetime64[ns]", "datetime64[us]", "datetime64[ms]", "datetime64[s]"],
 )
-def test_datetime_dateoffset_binaryop(date_col, n_periods, frequency, dtype):
+@pytest.mark.parametrize("op", [operator.add, operator.sub])
+def test_datetime_dateoffset_binaryop(
+    date_col, n_periods, frequency, dtype, op
+):
     gsr = cudf.Series(date_col, dtype=dtype)
     psr = gsr.to_pandas()  # converts to nanos
 
@@ -1488,16 +1508,165 @@ def test_datetime_dateoffset_binaryop(date_col, n_periods, frequency, dtype):
     goffset = cudf.DateOffset(**kwargs)
     poffset = pd.DateOffset(**kwargs)
 
-    expect = psr + poffset
-    got = gsr + goffset
+    expect = op(psr, poffset)
+    got = op(gsr, goffset)
+
+    utils.assert_eq(expect, got)
+
+    expect = op(psr, -poffset)
+    got = op(gsr, -goffset)
+
+    utils.assert_eq(expect, got)
+
+
+@pytest.mark.parametrize(
+    "date_col",
+    [
+        [
+            "2000-01-01 00:00:00.012345678",
+            "2000-01-31 00:00:00.012345678",
+            "2000-02-29 00:00:00.012345678",
+        ]
+    ],
+)
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"weeks": 0.5},
+        {"days": 0.5},
+        {"hours": 0.5},
+        {"minutes": 0.5},
+        {"seconds": 0.5},
+        utils.xfail_param({"seconds": 0.5e-6}),
+        utils.xfail_param({"microseconds": 0.5}),
+    ],
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["datetime64[ns]", "datetime64[us]", "datetime64[ms]", "datetime64[s]"],
+)
+@pytest.mark.parametrize("op", [operator.add, operator.sub])
+def test_datetime_dateoffset_binaryop_fractional_periods(
+    date_col, kwargs, dtype, op
+):
+    gsr = cudf.Series(date_col, dtype=dtype)
+    psr = gsr.to_pandas()  # converts to nanos
+
+    goffset = cudf.DateOffset(**kwargs)
+    poffset = pd.DateOffset(**kwargs)
+
+    expect = op(psr, poffset)
+    got = op(gsr, goffset)
+
+    utils.assert_eq(expect, got)
+
+
+@pytest.mark.parametrize(
+    "date_col",
+    [
+        [
+            "2000-01-01 00:00:00.012345678",
+            "2000-01-31 00:00:00.012345678",
+            "2000-02-29 00:00:00.012345678",
+        ]
+    ],
+)
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"months": 2, "years": 5},
+        {"microseconds": 1, "seconds": 1},
+        {"months": 2, "years": 5, "seconds": 923, "microseconds": 481},
+        pytest.param(
+            {"milliseconds": 4},
+            marks=pytest.mark.xfail(
+                reason="Pandas gets the wrong answer for milliseconds"
+            ),
+        ),
+        pytest.param(
+            {"milliseconds": 4, "years": 2},
+            marks=pytest.mark.xfail(
+                reason="Pandas construction fails with these keywords"
+            ),
+        ),
+        pytest.param(
+            {"nanoseconds": 12},
+            marks=pytest.mark.xfail(
+                reason="Pandas gets the wrong answer for nanoseconds"
+            ),
+        ),
+    ],
+)
+@pytest.mark.parametrize("op", [operator.add, operator.sub])
+def test_datetime_dateoffset_binaryop_multiple(date_col, kwargs, op):
+
+    gsr = cudf.Series(date_col, dtype="datetime64[ns]")
+    psr = gsr.to_pandas()
+
+    poffset = pd.DateOffset(**kwargs)
+    goffset = cudf.DateOffset(**kwargs)
+
+    expect = op(psr, poffset)
+    got = op(gsr, goffset)
 
     utils.assert_eq(expect, got)
 
-    expect = psr - poffset
-    got = gsr - goffset
+
+@pytest.mark.parametrize(
+    "date_col",
+    [
+        [
+            "2000-01-01 00:00:00.012345678",
+            "2000-01-31 00:00:00.012345678",
+            "2000-02-29 00:00:00.012345678",
+        ]
+    ],
+)
+@pytest.mark.parametrize("n_periods", [0, 1, -1, 12, -12])
+@pytest.mark.parametrize(
+    "frequency",
+    [
+        "months",
+        "years",
+        "days",
+        "hours",
+        "minutes",
+        "seconds",
+        "microseconds",
+        pytest.param(
+            "nanoseconds",
+            marks=pytest.mark.xfail(
+                reason="https://github.com/pandas-dev/pandas/issues/36589"
+            ),
+        ),
+    ],
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["datetime64[ns]", "datetime64[us]", "datetime64[ms]", "datetime64[s]"],
+)
+def test_datetime_dateoffset_binaryop_reflected(
+    date_col, n_periods, frequency, dtype
+):
+    gsr = cudf.Series(date_col, dtype=dtype)
+    psr = gsr.to_pandas()  # converts to nanos
+
+    kwargs = {frequency: n_periods}
+
+    goffset = cudf.DateOffset(**kwargs)
+    poffset = pd.DateOffset(**kwargs)
+
+    expect = poffset + psr
+    got = goffset + gsr
 
     utils.assert_eq(expect, got)
 
+    with pytest.raises(TypeError):
+        poffset - psr
+
+    with pytest.raises(TypeError):
+        goffset - gsr
+
 
 @pytest.mark.parametrize("frame", [cudf.Series, cudf.Index, cudf.DataFrame])
 @pytest.mark.parametrize(

python/cudf/cudf/tests/test_offset.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.
+
+import pytest
+
+from cudf import DateOffset
+
+
+@pytest.mark.parametrize("period", [1.5, 0.5, "string", "1", "1.0"])
+@pytest.mark.parametrize("freq", ["years", "months"])
+def test_construction_invalid(period, freq):
+    kwargs = {freq: period}
+    with pytest.raises(ValueError):
+        DateOffset(**kwargs)