Interval index and interval_range

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 #6885 Share `factorize` implementation with Index and cudf module
 - PR #6775 Implement cudf.DateOffset for months
 

python/cudf/cudf/__init__.py

@@ -18,6 +18,8 @@
 from cudf.core import (
     NA,
     CategoricalIndex,
+    interval_range,
+    IntervalIndex,
     DataFrame,
     DatetimeIndex,
     Float32Index,

python/cudf/cudf/core/__init__.py

@@ -5,6 +5,8 @@
 from cudf.core.dataframe import DataFrame, from_pandas, merge
 from cudf.core.index import (
     CategoricalIndex,
+    interval_range,
+    IntervalIndex,
     DatetimeIndex,
     Float32Index,
     Float64Index,

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

@@ -22,3 +22,4 @@
 from cudf.core.column.struct import StructColumn  # noqa: F401
 from cudf.core.column.timedelta import TimeDeltaColumn  # noqa: F401
 from cudf.core.column.decimal import DecimalColumn  # noqa: F401
+from cudf.core.column.interval import IntervalColumn  # noqa: F401

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

@@ -25,6 +25,7 @@
 from cudf.core.abc import Serializable
 from cudf.core.buffer import Buffer
 from cudf.core.dtypes import CategoricalDtype
+from cudf.core.dtypes import IntervalDtype
 from cudf.utils import ioutils, utils
 from cudf.utils.dtypes import (
     NUMERIC_TYPES,
@@ -38,6 +39,7 @@
     is_scalar,
     is_string_dtype,
     is_struct_dtype,
+    is_interval_dtype,
     min_signed_type,
     min_unsigned_type,
     np_to_pa_dtype,
@@ -133,6 +135,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)
 
@@ -371,7 +377,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(
@@ -407,6 +412,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"
@@ -1015,6 +1024,12 @@ def astype(self, dtype, **kwargs):
                     "Casting list columns not currently supported"
                 )
             return self
+        elif is_interval_dtype(self.dtype):
+            if not self.dtype == dtype:
+                raise NotImplementedError(
+                    "Casting interval columns not currently supported"
+                )
+            return self
         elif np.issubdtype(dtype, np.datetime64):
             return self.as_datetime_column(dtype, **kwargs)
         elif np.issubdtype(dtype, np.timedelta64):
@@ -1480,6 +1495,15 @@ def build_column(
             null_count=null_count,
             children=children,
         )
+    elif is_interval_dtype(dtype):
+        return cudf.core.column.IntervalColumn(
+            dtype=dtype,
+            mask=mask,
+            size=size,
+            offset=offset,
+            children=children,
+            null_count=null_count,
+        )
     elif is_struct_dtype(dtype):
         return cudf.core.column.StructColumn(
             data=data,
@@ -1535,14 +1559,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,
@@ -1554,6 +1576,48 @@ def build_categorical_column(
     )
 
 
+def build_interval_column(
+    left_col,
+    right_col,
+    mask=None,
+    size=None,
+    offset=0,
+    null_count=None,
+    closed="right",
+):
+    """
+    Build an IntervalColumn
+
+    Parameters
+    ----------
+    left_col : Column
+        Column of values representing the left of the interval
+    right_col : Column
+        Column of representing the right of the interval
+    mask : Buffer
+        Null mask
+    size : int, optional
+    offset : int, optional
+    closed : str
+        Indicates which side if any is closed in the interval set
+    """
+    left = as_column(left_col)
+    right = as_column(right_col)
+
+    dtype = IntervalDtype(left_col.dtype, closed)
+    size = len(left)
+    # breakpoint()
+    return build_column(
+        data=None,
+        dtype=dtype,
+        mask=mask,
+        size=size,
+        offset=offset,
+        null_count=null_count,
+        children=(left, right),
+    )
+
+
 def as_column(arbitrary, nan_as_null=None, dtype=None, length=None):
     """Create a Column from an arbitrary object
 
@@ -1675,6 +1739,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"):
@@ -1794,9 +1860,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.
@@ -1824,7 +1898,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(
@@ -1880,7 +1963,7 @@ def as_column(arbitrary, nan_as_null=None, dtype=None, length=None):
                         )
                         return cudf.core.column.DecimalColumn.from_arrow(data)
                     dtype = pd.api.types.pandas_dtype(dtype)
-                    if is_categorical_dtype(dtype):
+                    if is_categorical_dtype(dtype) or is_interval_dtype(dtype):
                         raise TypeError
                     else:
                         np_type = np.dtype(dtype).type
@@ -1906,6 +1989,9 @@ 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 is_interval_dtype(dtype):
+                    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

@@ -418,11 +418,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
 
@@ -214,7 +215,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})"
 
     def __hash__(self):
         return hash(self._typ)
@@ -299,3 +300,29 @@ def _validate(cls, precision, scale=0):
             )
         if abs(scale) > precision:
             raise ValueError(f"scale={scale} exceeds precision={precision}")
+
+
+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

@@ -2402,9 +2402,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)