Raise NotImplementedError for Categoricals with timezones

rapidsai · Sep 1, 2023 · 39b42f9 · 39b42f9
1 parent d1fb671
commit 39b42f9
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 11 deletions.
diff --git a/python/cudf/cudf/core/column/column.py b/python/cudf/cudf/core/column/column.py
@@ -2033,9 +2033,19 @@ def as_column(
                     f"{arbitrary.dtype} is not supported. Convert first to "
                     f"{arbitrary.dtype.subtype}."
                 )
-        if is_categorical_dtype(arbitrary):
+        if is_categorical_dtype(arbitrary.dtype):
+            if isinstance(
+                arbitrary.dtype.categories.dtype, pd.DatetimeTZDtype
+            ):
+                raise NotImplementedError(
+                    "cuDF does not yet support timezone-aware datetimes"
+                )
             data = as_column(pa.array(arbitrary, from_pandas=True))
         elif is_interval_dtype(arbitrary.dtype):
+            if isinstance(arbitrary.dtype.subtype, pd.DatetimeTZDtype):
+                raise NotImplementedError(
+                    "cuDF does not yet support timezone-aware datetimes"
+                )
             data = as_column(pa.array(arbitrary, from_pandas=True))
         elif arbitrary.dtype == np.bool_:
             data = as_column(cupy.asarray(arbitrary), dtype=arbitrary.dtype)
@@ -2262,11 +2272,20 @@ def as_column(
     elif isinstance(arbitrary, pd.core.arrays.masked.BaseMaskedArray):
         data = as_column(pa.Array.from_pandas(arbitrary), dtype=dtype)
     elif (
-        isinstance(arbitrary, pd.DatetimeIndex)
-        and isinstance(arbitrary.dtype, pd.DatetimeTZDtype)
-    ) or (
-        isinstance(arbitrary, pd.IntervalIndex)
-        and is_datetime64tz_dtype(arbitrary.dtype.subtype)
+        (
+            isinstance(arbitrary, pd.DatetimeIndex)
+            and isinstance(arbitrary.dtype, pd.DatetimeTZDtype)
+        )
+        or (
+            isinstance(arbitrary, pd.IntervalIndex)
+            and is_datetime64tz_dtype(arbitrary.dtype.subtype)
+        )
+        or (
+            isinstance(arbitrary, pd.CategoricalIndex)
+            and isinstance(
+                arbitrary.dtype.categories.dtype, pd.DatetimeTZDtype
+            )
+        )
     ):
         raise NotImplementedError(
             "cuDF does not yet support timezone-aware datetimes"

diff --git a/python/cudf/cudf/tests/test_datetime.py b/python/cudf/cudf/tests/test_datetime.py
@@ -2095,6 +2095,8 @@ def test_construction_from_tz_timestamps(data):
         _ = cudf.Index(data)
     with pytest.raises(NotImplementedError):
         _ = cudf.DatetimeIndex(data)
+    with pytest.raises(NotImplementedError):
+        cudf.CategoricalIndex(data)
 
 
 @pytest.mark.parametrize("op", _cmpops)

diff --git a/python/cudf/cudf/tests/test_interval.py b/python/cudf/cudf/tests/test_interval.py
@@ -167,17 +167,18 @@ def test_interval_index_unique():
     assert_eq(expected, actual)
 
 
+@pytest.mark.parametrize("box", [pd.Series, pd.IntervalIndex])
 @pytest.mark.parametrize("tz", ["US/Eastern", None])
-def test_interval_with_datetime(tz):
+def test_interval_with_datetime(tz, box):
     dti = pd.date_range(
         start=pd.Timestamp("20180101", tz=tz),
         end=pd.Timestamp("20181231", tz=tz),
         freq="M",
     )
-    pidx = pd.IntervalIndex.from_breaks(dti)
+    pobj = box(pd.IntervalIndex.from_breaks(dti))
     if tz is None:
-        gidx = cudf.from_pandas(pidx)
-        assert_eq(pidx, gidx)
+        gobj = cudf.from_pandas(pobj)
+        assert_eq(pobj, gobj)
     else:
         with pytest.raises(NotImplementedError):
-            cudf.from_pandas(pidx)
+            cudf.from_pandas(pobj)