Interval index and interval_range

python/cudf/cudf/core/index.py

@@ -11,34 +11,36 @@
 import pandas as pd
 from nvtx import annotate
 from pandas._config import get_option
-from cudf._lib.filling import sequence
 
 import cudf
+from cudf._lib.filling import sequence
 from cudf._typing import DtypeObj
 from cudf.core.abc import Serializable
 from cudf.core.column import (
     CategoricalColumn,
-    IntervalColumn,
     ColumnBase,
     DatetimeColumn,
+    IntervalColumn,
     NumericalColumn,
     StringColumn,
     TimeDeltaColumn,
-    column,
     arange,
+    column,
 )
 from cudf.core.column.string import StringMethods as StringMethods
 from cudf.core.dtypes import IntervalDtype
 from cudf.core.frame import Frame
 from cudf.utils import ioutils, utils
 from cudf.utils.docutils import copy_docstring
 from cudf.utils.dtypes import (
+    find_common_type,
     is_categorical_dtype,
+    is_interval_dtype,
     is_list_like,
     is_mixed_with_object_dtype,
+    is_numerical_dtype,
     is_scalar,
     numeric_normalize_types,
-    is_interval_dtype,
 )
 from cudf.utils.utils import cached_property, search_range
 
@@ -2811,57 +2813,80 @@ def interval_range(
             "Of the four parameters: start, end, periods, and "
             "freq, exactly three must be specified"
         )
-    if not isinstance(start or freq or end, int) and not isinstance(
-        start or freq or end, float
+    args = [start, end, freq, periods]
+    *args, periods = [cudf.Scalar(x) if x is not None else None for x in args]
+    if any(
+        [
+            not is_numerical_dtype(x.dtype) if x is not None else False
+            for x in args
+        ]
     ):
-        raise NotImplementedError("Non-numeric values not yet supported")
-    elif periods and not freq:
+        raise ValueError("start, end, freq must be numeric values.")
+    common_dtype = find_common_type([x.dtype for x in args if x])
+    start, end, freq = args
+
+    if periods and not freq:
         # if statement for mypy to pass
         if end is not None and start is not None:
             # determine if periods are float or integer
-            periods = int(periods)
-            quotient, remainder = divmod((end - start), periods)
+            periods = periods._as_host_type("int64")
+            # divmod only supported on host side scalars
+            quotient, remainder = divmod((end - start).value, periods.value)
             if remainder:
                 freq_step = cudf.Scalar((end - start) / periods)
             else:
                 freq_step = cudf.Scalar(quotient)
-            start = cudf.Scalar(start)
+            common_dtype = find_common_type([common_dtype, freq_step.dtype])
             if start.dtype != freq_step.dtype:
-                start = cudf.Scalar(start.value.astype(freq_step.dtype))
-            start = start.device_value
-            freq_step = freq_step.device_value
-            bin_edges = sequence(size=periods + 1, init=start, step=freq_step,)
-            left_col = bin_edges[:-1]
-            right_col = bin_edges[1:]
+                start = start._as_host_type(freq_step.dtype)
+            bin_edges = sequence(
+                size=periods + 1,
+                init=start.device_value,
+                step=freq_step.device_value,
+            )
+            left_col = bin_edges[:-1].astype(common_dtype)
+            right_col = bin_edges[1:].astype(common_dtype)
     elif freq and periods:
         if end:
             start = end - (freq * periods)
         if start:
             end = freq * periods + start
         if end is not None and start is not None:
-            left_col = arange(start, end, freq)
+            left_col = arange(
+                start.value, end.value, freq.value, dtype=common_dtype
+            )
             end = end + 1
             start = start + freq
-            right_col = arange(start, end, freq)
+            right_col = arange(
+                start.value, end.value, freq.value, dtype=common_dtype
+            )
     elif freq and not periods:
         if end is not None and start is not None:
             end = end - freq + 1
-            left_col = arange(start, end, freq)
+            left_col = arange(
+                start.value, end.value, freq.value, dtype=common_dtype
+            )
             end = end + freq + 1
             start = start + freq
-            right_col = arange(start, end, freq)
+            right_col = arange(
+                start.value, end.value, freq.value, dtype=common_dtype
+            )
     elif start is not None and end is not None:
         # if statements for mypy to pass
         if freq:
-            left_col = arange(start, end, freq)
+            left_col = arange(
+                start.value, end.value, freq.value, dtype=common_dtype
+            )
         else:
-            left_col = arange(start, end)
+            left_col = arange(start.value, end.value, dtype=common_dtype)
         start = start + 1
         end = end + 1
         if freq:
-            right_col = arange(start, end, freq)
+            right_col = arange(
+                start.value, end.value, freq.value, dtype=common_dtype
+            )
         else:
-            right_col = arange(start, end)
+            right_col = arange(start.value, end.value, dtype=common_dtype)
     else:
         raise ValueError(
             "Of the four parameters: start, end, periods, and "

python/cudf/cudf/core/scalar.py

@@ -360,6 +360,9 @@ def _dispatch_scalar_unaop(self, op):
     def astype(self, dtype):
         return Scalar(self.device_value, dtype)
 
+    def _as_host_type(self, dtype):
+        return Scalar(self.value, dtype)
+
 
 class _NAType(object):
     def __init__(self):

python/cudf/cudf/tests/test_index.py

@@ -1361,6 +1361,18 @@ def test_categorical_index_basic(data, categories, dtype, ordered, name):
     assert_eq(pindex, gindex)
 
 
+INTERVAL_BOUNDARY_TYPES = [
+    int,
+    np.int8,
+    np.int16,
+    np.int32,
+    np.int64,
+    np.float32,
+    np.float64,
+    cudf.Scalar,
+]
+
+
 @pytest.mark.parametrize("closed", ["left", "right", "both", "neither"])
 @pytest.mark.parametrize("start", [0, 1, 2, 3])
 @pytest.mark.parametrize("end", [4, 5, 6, 7])
@@ -1371,6 +1383,18 @@ def test_interval_range_basic(start, end, closed):
     assert_eq(pindex, gindex)
 
 
+@pytest.mark.parametrize("start_t", INTERVAL_BOUNDARY_TYPES)
+@pytest.mark.parametrize("end_t", INTERVAL_BOUNDARY_TYPES)
+def test_interval_range_dtype_basic(start_t, end_t):
+    start, end = start_t(24), end_t(42)
+    start_val = start.value if isinstance(start, cudf.Scalar) else start
+    end_val = end.value if isinstance(end, cudf.Scalar) else end
+    pindex = pd.interval_range(start=start_val, end=end_val, closed="left")
+    gindex = cudf.interval_range(start=start, end=end, closed="left")
+
+    assert_eq(pindex, gindex)
+
+
 @pytest.mark.parametrize("closed", ["left", "right", "both", "neither"])
 @pytest.mark.parametrize("start", [0])
 @pytest.mark.parametrize("end", [0])
@@ -1394,6 +1418,24 @@ def test_interval_range_freq_basic(start, end, freq, closed):
     assert_eq(pindex, gindex)
 
 
+@pytest.mark.parametrize("start_t", INTERVAL_BOUNDARY_TYPES)
+@pytest.mark.parametrize("end_t", INTERVAL_BOUNDARY_TYPES)
+@pytest.mark.parametrize("freq_t", INTERVAL_BOUNDARY_TYPES)
+def test_interval_range_freq_basic_dtype(start_t, end_t, freq_t):
+    start, end, freq = start_t(5), end_t(70), freq_t(3)
+    start_val = start.value if isinstance(start, cudf.Scalar) else start
+    end_val = end.value if isinstance(end, cudf.Scalar) else end
+    freq_val = freq.value if isinstance(freq, cudf.Scalar) else freq
+    pindex = pd.interval_range(
+        start=start_val, end=end_val, freq=freq_val, closed="left"
+    )
+    gindex = cudf.interval_range(
+        start=start, end=end, freq=freq, closed="left"
+    )
+
+    assert_eq(pindex, gindex)
+
+
 @pytest.mark.parametrize("closed", ["left", "right", "both", "neither"])
 @pytest.mark.parametrize("periods", [1, 1.0, 2, 2.0, 3.0, 3])
 @pytest.mark.parametrize("start", [0, 0.0, 1.0, 1, 2, 2.0, 3.0, 3])
@@ -1409,6 +1451,26 @@ def test_interval_range_periods_basic(start, end, periods, closed):
     assert_eq(pindex, gindex)
 
 
+@pytest.mark.parametrize("start_t", INTERVAL_BOUNDARY_TYPES)
+@pytest.mark.parametrize("end_t", INTERVAL_BOUNDARY_TYPES)
+@pytest.mark.parametrize("periods_t", INTERVAL_BOUNDARY_TYPES)
+def test_interval_range_periods_basic_dtype(start_t, end_t, periods_t):
+    start, end, periods = start_t(0), end_t(4), periods_t(1.0)
+    start_val = start.value if isinstance(start, cudf.Scalar) else start
+    end_val = end.value if isinstance(end, cudf.Scalar) else end
+    periods_val = (
+        periods.value if isinstance(periods, cudf.Scalar) else periods
+    )
+    pindex = pd.interval_range(
+        start=start_val, end=end_val, periods=periods_val, closed="left"
+    )
+    gindex = cudf.interval_range(
+        start=start, end=end, periods=periods, closed="left"
+    )
+
+    assert_eq(pindex, gindex)
+
+
 @pytest.mark.parametrize("closed", ["left", "right", "both", "neither"])
 @pytest.mark.parametrize("periods", [1, 2, 3])
 @pytest.mark.parametrize("freq", [1, 2, 3, 4])
@@ -1424,6 +1486,26 @@ def test_interval_range_periods_freq_end(end, freq, periods, closed):
     assert_eq(pindex, gindex)
 
 
+@pytest.mark.parametrize("periods_t", INTERVAL_BOUNDARY_TYPES)
+@pytest.mark.parametrize("freq_t", INTERVAL_BOUNDARY_TYPES)
+@pytest.mark.parametrize("end_t", INTERVAL_BOUNDARY_TYPES)
+def test_interval_range_periods_freq_end_dtype(periods_t, freq_t, end_t):
+    periods, freq, end = periods_t(2), freq_t(3), end_t(10)
+    freq_val = freq.value if isinstance(freq, cudf.Scalar) else freq
+    end_val = end.value if isinstance(end, cudf.Scalar) else end
+    periods_val = (
+        periods.value if isinstance(periods, cudf.Scalar) else periods
+    )
+    pindex = pd.interval_range(
+        end=end_val, freq=freq_val, periods=periods_val, closed="left"
+    )
+    gindex = cudf.interval_range(
+        end=end, freq=freq, periods=periods, closed="left"
+    )
+
+    assert_eq(pindex, gindex)
+
+
 @pytest.mark.parametrize("closed", ["left", "right", "both", "neither"])
 @pytest.mark.parametrize("periods", [1, 2, 3])
 @pytest.mark.parametrize("freq", [1, 2, 3, 4])
@@ -1439,6 +1521,26 @@ def test_interval_range_periods_freq_start(start, freq, periods, closed):
     assert_eq(pindex, gindex)
 
 
+@pytest.mark.parametrize("periods_t", INTERVAL_BOUNDARY_TYPES)
+@pytest.mark.parametrize("freq_t", INTERVAL_BOUNDARY_TYPES)
+@pytest.mark.parametrize("start_t", INTERVAL_BOUNDARY_TYPES)
+def test_interval_range_periods_freq_start_dtype(periods_t, freq_t, start_t):
+    periods, freq, start = periods_t(2), freq_t(3), start_t(9)
+    freq_val = freq.value if isinstance(freq, cudf.Scalar) else freq
+    start_val = start.value if isinstance(start, cudf.Scalar) else start
+    periods_val = (
+        periods.value if isinstance(periods, cudf.Scalar) else periods
+    )
+    pindex = pd.interval_range(
+        start=start_val, freq=freq_val, periods=periods_val, closed="left"
+    )
+    gindex = cudf.interval_range(
+        start=start, freq=freq, periods=periods, closed="left"
+    )
+
+    assert_eq(pindex, gindex)
+
+
 @pytest.mark.parametrize("closed", ["right", "left", "both", "neither"])
 @pytest.mark.parametrize(
     "data",