Adding Interval Dtype

CHANGELOG.md

@@ -7,6 +7,7 @@
 - PR #6814 Implement `cudf::reduce` for `decimal32` and `decimal64` (part 1)
 - PR #6929 Add `Index.set_names` api
 - PR #6907 Add `replace_null` API with `replace_policy` parameter, `fixed_width` column support
+- PR #6984 Add `interval` dtype to cudf
 
 - PR #6775 Implement cudf.DateOffset for months
 

python/cudf/cudf/core/column/__init__.py

@@ -21,3 +21,4 @@
 from cudf.core.column.string import StringColumn  # noqa: F401
 from cudf.core.column.struct import StructColumn  # noqa: F401
 from cudf.core.column.timedelta import TimeDeltaColumn  # noqa: F401
+from cudf.core.column.interval import IntervalColumn  # noqa: F401

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

@@ -37,6 +37,7 @@
     is_scalar,
     is_string_dtype,
     is_struct_dtype,
+    is_interval_dtype,
     min_signed_type,
     min_unsigned_type,
     np_to_pa_dtype,
@@ -132,6 +133,10 @@ def to_pandas(self, index=None, nullable=False, **kwargs):
             pd_series = pd.Series(pandas_array, copy=False)
         elif str(self.dtype) in NUMERIC_TYPES and self.null_count == 0:
             pd_series = pd.Series(cupy.asnumpy(self.values), copy=False)
+        elif is_interval_dtype(self.dtype):
+            pd_series = pd.Series(
+                pd.IntervalDtype().__from_arrow__(self.to_arrow())
+            )
         else:
             pd_series = self.to_arrow().to_pandas(**kwargs)
 
@@ -364,7 +369,6 @@ def from_arrow(cls, array):
         """
         if not isinstance(array, (pa.Array, pa.ChunkedArray)):
             raise TypeError("array should be PyArrow array or chunked array")
-
         data = pa.table([array], [None])
         if isinstance(array.type, pa.DictionaryType):
             indices_table = pa.table(
@@ -400,6 +404,10 @@ def from_arrow(cls, array):
             )
         elif isinstance(array.type, pa.StructType):
             return cudf.core.column.StructColumn.from_arrow(array)
+        elif isinstance(
+            array.type, pd.core.arrays._arrow_utils.ArrowIntervalType
+        ):
+            return cudf.core.column.IntervalColumn.from_arrow(array)
 
         return libcudf.interop.from_arrow(data, data.column_names)._data[
             "None"
@@ -1006,6 +1014,8 @@ def astype(self, dtype, **kwargs):
                     "Casting list columns not currently supported"
                 )
             return self
+        elif is_interval_dtype(self.dtype):
+            return self
         elif np.issubdtype(dtype, np.datetime64):
             return self.as_datetime_column(dtype, **kwargs)
         elif np.issubdtype(dtype, np.timedelta64):
@@ -1478,6 +1488,15 @@ def build_column(
             null_count=null_count,
             children=children,
         )
+    elif is_interval_dtype(dtype):
+        return cudf.core.column.IntervalColumn(
+            data=data,
+            dtype=dtype,
+            mask=mask,
+            size=size,
+            offset=offset,
+            null_count=null_count,
+        )
     else:
         return cudf.core.column.NumericalColumn(
             data=data,
@@ -1515,14 +1534,12 @@ def build_categorical_column(
     ordered : bool
         Indicates whether the categories are ordered
     """
-
     codes_dtype = min_unsigned_type(len(categories))
     codes = as_column(codes)
     if codes.dtype != codes_dtype:
         codes = codes.astype(codes_dtype)
 
     dtype = CategoricalDtype(categories=as_column(categories), ordered=ordered)
-
     return build_column(
         data=None,
         dtype=dtype,
@@ -1655,6 +1672,8 @@ def as_column(arbitrary, nan_as_null=None, dtype=None, length=None):
             return as_column(arbitrary.array)
         if is_categorical_dtype(arbitrary):
             data = as_column(pa.array(arbitrary, from_pandas=True))
+        elif is_interval_dtype(arbitrary.dtype):
+            data = as_column(pa.array(arbitrary, from_pandas=True))
         elif arbitrary.dtype == np.bool:
             data = as_column(cupy.asarray(arbitrary), dtype=arbitrary.dtype)
         elif arbitrary.dtype.kind in ("f"):
@@ -1774,9 +1793,17 @@ def as_column(arbitrary, nan_as_null=None, dtype=None, length=None):
                 data=buffer, mask=mask, dtype=arbitrary.dtype
             )
         elif arb_dtype.kind in ("O", "U"):
-            data = as_column(
-                pa.Array.from_pandas(arbitrary), dtype=arbitrary.dtype
-            )
+            if isinstance(arbitrary[0], pd._libs.interval.Interval):
+                # changing from pd array to series,possible arrow bug
+                interval_series = pd.Series(arbitrary)
+                data = as_column(
+                    pa.Array.from_pandas(interval_series),
+                    dtype=arbitrary.dtype,
+                )
+            else:
+                data = as_column(
+                    pa.Array.from_pandas(arbitrary), dtype=arbitrary.dtype
+                )
             # There is no cast operation available for pa.Array from int to
             # str, Hence instead of handling in pa.Array block, we
             # will have to type-cast here.
@@ -1804,7 +1831,16 @@ def as_column(arbitrary, nan_as_null=None, dtype=None, length=None):
                 if arb_dtype != arbitrary.dtype.numpy_dtype:
                     arbitrary = arbitrary.astype(arb_dtype)
         if arb_dtype.kind in ("O", "U"):
-            data = as_column(pa.Array.from_pandas(arbitrary), dtype=arb_dtype)
+            if isinstance(arbitrary[0], pd._libs.interval.Interval):
+                # changing from pd array to series,possible arrow bug
+                interval_series = pd.Series(arbitrary)
+                data = as_column(
+                    pa.Array.from_pandas(interval_series), dtype=arb_dtype
+                )
+            else:
+                data = as_column(
+                    pa.Array.from_pandas(arbitrary), dtype=arb_dtype
+                )
         else:
             data = as_column(
                 pa.array(
@@ -1852,7 +1888,10 @@ def as_column(arbitrary, nan_as_null=None, dtype=None, length=None):
                             )
                         return as_column(data, nan_as_null=nan_as_null)
                     dtype = pd.api.types.pandas_dtype(dtype)
-                    if is_categorical_dtype(dtype):
+                    if (
+                        is_categorical_dtype(dtype)
+                        or type(dtype) == pd.core.dtypes.dtypes.IntervalDtype
+                    ):
                         raise TypeError
                     else:
                         np_type = np.dtype(dtype).type
@@ -1878,6 +1917,10 @@ def as_column(arbitrary, nan_as_null=None, dtype=None, length=None):
                 elif np_type == np.str_:
                     sr = pd.Series(arbitrary, dtype="str")
                     data = as_column(sr, nan_as_null=nan_as_null)
+                elif type(dtype) == pd.core.dtypes.dtypes.IntervalDtype:
+                    # breakpoint()
+                    sr = pd.Series(arbitrary, dtype="interval")
+                    data = as_column(sr, nan_as_null=nan_as_null)
                 else:
                     native_dtype = dtype
                     if dtype is None and pd.api.types.infer_dtype(

python/cudf/cudf/core/column/interval.py

@@ -0,0 +1,58 @@
+import pyarrow as pa
+import cudf
+from cudf.core.column import StructColumn
+
+
+class IntervalColumn(StructColumn):
+    def __init__(
+        self,
+        dtype,
+        mask=None,
+        size=None,
+        offset=0,
+        null_count=None,
+        children=(),
+        closed="right",
+    ):
+
+        super().__init__(
+            data=None,
+            dtype=dtype,
+            mask=mask,
+            size=size,
+            offset=offset,
+            null_count=null_count,
+            children=children,
+        )
+        self.closed = closed
+
+    @classmethod
+    def from_arrow(self, data):
+        new_col = super().from_arrow(data.storage)
+        size = len(data)
+        dtype = cudf.core.dtypes.IntervalDtype.from_arrow(data.type)
+        mask = data.buffers()[0]
+        if mask is not None:
+            mask = cudf.utils.utils.pa_mask_buffer_to_mask(mask, len(data))
+
+        offset = data.offset
+        null_count = data.null_count
+        children = new_col.children
+        closed = dtype.closed
+
+        return IntervalColumn(
+            size=size,
+            dtype=dtype,
+            mask=mask,
+            offset=offset,
+            null_count=null_count,
+            children=children,
+            closed=closed,
+        )
+
+    def to_arrow(self):
+        typ = self.dtype.to_arrow()
+        return pa.ExtensionArray.from_storage(typ, super().to_arrow())
+
+    def copy(self, deep=True):
+        return super().copy(deep=deep).as_interval_column(self.closed)

python/cudf/cudf/core/column/struct.py

@@ -71,6 +71,21 @@ def copy(self, deep=True):
             result = result._rename_fields(self.dtype.fields.keys())
         return result
 
+    def as_interval_column(self, closed="right"):
+        from cudf.core.column import IntervalColumn
+
+        return IntervalColumn(
+            size=self.size,
+            dtype=cudf.core.dtypes.IntervalDtype(
+                self.dtype.fields["left"], closed
+            ),
+            mask=self.mask,
+            offset=self.offset,
+            null_count=self.null_count,
+            children=self.children,
+            closed=closed,
+        )
+
     def _rename_fields(self, names):
         """
         Return a StructColumn with the same field values as this StructColumn,

python/cudf/cudf/core/dataframe.py

@@ -416,11 +416,14 @@ def _init_from_list_like(self, data, index=None, columns=None):
             index = as_index(index)
 
         self._index = as_index(index)
-
         # list-of-dicts case
         if len(data) > 0 and isinstance(data[0], dict):
             data = DataFrame.from_pandas(pd.DataFrame(data))
             self._data = data._data
+        # interval in a list
+        elif len(data) > 0 and isinstance(data[0], pd._libs.interval.Interval):
+            data = DataFrame.from_pandas(pd.DataFrame(data))
+            self._data = data._data
         else:
             data = list(itertools.zip_longest(*data))
 

python/cudf/cudf/core/dtypes.py

@@ -7,6 +7,7 @@
 import pandas as pd
 import pyarrow as pa
 from pandas.api.extensions import ExtensionDtype
+from pandas.core.arrays._arrow_utils import ArrowIntervalType
 
 import cudf
 
@@ -211,7 +212,7 @@ def __eq__(self, other):
         return self._typ.equals(other._typ)
 
     def __repr__(self):
-        return f"StructDtype({self.fields})"
+        return f"{type(self).__name__}({self.fields})"
 
 
 class DecimalDtype(ExtensionDtype):
@@ -241,3 +242,29 @@ def to_arrow(self):
     @classmethod
     def from_arrow(cls, typ):
         return cls(typ.precision, typ.scale)
+
+
+class IntervalDtype(StructDtype):
+    name = "interval"
+
+    def __init__(self, subtype, closed="right"):
+        """
+        subtype: str, np.dtype
+            The dtype of the Interval bounds.
+        """
+        super().__init__(fields={"left": subtype, "right": subtype})
+        self.closed = closed
+
+    @property
+    def subtype(self):
+        return self.fields["left"]
+
+    @classmethod
+    def from_arrow(cls, typ):
+        return IntervalDtype(typ.subtype.to_pandas_dtype(), typ.closed)
+
+    def to_arrow(self):
+
+        return ArrowIntervalType(
+            pa.from_numpy_dtype(self.subtype), self.closed
+        )

python/cudf/cudf/core/frame.py

@@ -2205,9 +2205,17 @@ def _copy_struct_names(self, other, include_index=True):
 
         return self
 
+    def _copy_interval_data(self, other, include_index=True):
+        for name, col, other_col in zip(
+            self._data.keys(), self._data.values(), other._data.values()
+        ):
+            if isinstance(other_col, cudf.core.column.IntervalColumn):
+                self._data[name] = col.as_interval_column()
+
     def _postprocess_columns(self, other, include_index=True):
         self._copy_categories(other, include_index=include_index)
         self._copy_struct_names(other, include_index=include_index)
+        self._copy_interval_data(other, include_index=include_index)
 
     def _unaryop(self, op):
         data_columns = (col.unary_operator(op) for col in self._columns)

python/cudf/cudf/tests/test_dtypes.py

@@ -11,6 +11,7 @@
     ListDtype,
     StructDtype,
     DecimalDtype,
+    IntervalDtype,
 )
 from cudf.tests.utils import assert_eq
 
@@ -139,3 +140,12 @@ def test_decimal_dtype():
     dt = DecimalDtype(4, 2)
     assert dt.to_arrow() == pa.decimal128(4, 2)
     assert dt == DecimalDtype.from_arrow(pa.decimal128(4, 2))
+
+
+@pytest.mark.parametrize("fields", ["int64", "int32"])
+@pytest.mark.parametrize("closed", ["left", "right", "both", "neither"])
+def test_interval_dtype_pyarrow_round_trip(fields, closed):
+    pa_array = pd.core.arrays._arrow_utils.ArrowIntervalType(fields, closed)
+    expect = pa_array
+    got = IntervalDtype.from_arrow(expect).to_arrow()
+    assert expect.equals(got)