Skip to content

Commit

Permalink
Add missing __cuda_array_interface__ for datetime objects.
Browse files Browse the repository at this point in the history
  • Loading branch information
vyasr committed Apr 27, 2021
1 parent 51e6bea commit 0c2c7a1
Showing 1 changed file with 30 additions and 1 deletion.
31 changes: 30 additions & 1 deletion python/cudf/cudf/core/column/datetime.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,12 @@

from __future__ import annotations

import builtins
import datetime as dt
import re
from numbers import Number
from typing import Any, Sequence, Union, cast
from types import SimpleNamespace
from typing import Any, Mapping, Sequence, Union, cast

import numpy as np
import pandas as pd
Expand Down Expand Up @@ -199,6 +201,33 @@ def as_numerical(self) -> "cudf.core.column.NumericalColumn":
),
)

@property
def __cuda_array_interface__(self) -> Mapping[builtins.str, Any]:
output = {
"shape": (len(self),),
"strides": (self.dtype.itemsize,),
"typestr": self.dtype.str,
"data": (self.data_ptr, False),
"version": 1,
}

if self.nullable and self.has_nulls:

# Create a simple Python object that exposes the
# `__cuda_array_interface__` attribute here since we need to modify
# some of the attributes from the numba device array
mask = SimpleNamespace(
__cuda_array_interface__={
"shape": (len(self),),
"typestr": "<t1",
"data": (self.mask_ptr, True),
"version": 1,
}
)
output["mask"] = mask

return output

def as_datetime_column(self, dtype: Dtype, **kwargs) -> DatetimeColumn:
dtype = np.dtype(dtype)
if dtype == self.dtype:
Expand Down

0 comments on commit 0c2c7a1

Please sign in to comment.