Extend cuda_array_interface_wrapper

python/cudf/cudf/core/buffer/__init__.py

@@ -1,4 +1,4 @@
 # Copyright (c) 2022, NVIDIA CORPORATION.
 
-from cudf.core.buffer.buffer import Buffer
+from cudf.core.buffer.buffer import Buffer, cuda_array_interface_wrapper
 from cudf.core.buffer.utils import as_buffer

python/cudf/cudf/core/buffer/buffer.py

@@ -18,9 +18,18 @@
 T = TypeVar("T", bound="Buffer")
 
 
-def cuda_array_interface_wrapper(ptr: int, size: int, owner: object = None):
+def cuda_array_interface_wrapper(
+    ptr: int,
+    size: int,
+    owner: object = None,
+    readonly=False,
+    typestr="|u1",
+    version=0,
+):
     """Wrap device pointer in an object that exposes `__cuda_array_interface__`
 
+    See <https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html>
+
     Parameters
     ----------
     ptr : int
@@ -30,6 +39,14 @@ def cuda_array_interface_wrapper(ptr: int, size: int, owner: object = None):
     owner : object, optional
         Python object to which the lifetime of the memory allocation is tied.
         A reference to this object is kept in the returned wrapper object.
+    readonly: bool, optional
+        Mark the interface read-only.
+    typestr: str, optional
+        The type string of the interface. By default this is "|u1", which
+        means "an unsigned integer with a not relevant byteorder". See:
+        <https://numpy.org/doc/stable/reference/arrays.interface.html>
+    version : bool, optional
+        The version of the interface.
 
     Return
     ------
@@ -43,11 +60,11 @@ def cuda_array_interface_wrapper(ptr: int, size: int, owner: object = None):
 
     return SimpleNamespace(
         __cuda_array_interface__={
-            "data": (ptr, False),
+            "data": (ptr, readonly),
             "shape": (size,),
             "strides": None,
-            "typestr": "|u1",
-            "version": 0,
+            "typestr": typestr,
+            "version": version,
         },
         owner=owner,
     )

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

@@ -6,7 +6,6 @@
 import locale
 import re
 from locale import nl_langinfo
-from types import SimpleNamespace
 from typing import Any, Mapping, Sequence, cast
 
 import numpy as np
@@ -23,7 +22,7 @@
 )
 from cudf.api.types import is_datetime64_dtype, is_scalar, is_timedelta64_dtype
 from cudf.core._compat import PANDAS_GE_120
-from cudf.core.buffer import Buffer
+from cudf.core.buffer import Buffer, cuda_array_interface_wrapper
 from cudf.core.column import ColumnBase, as_column, column, string
 from cudf.core.column.timedelta import _unit_to_nanoseconds_conversion
 from cudf.utils.utils import _fillna_natwise
@@ -289,20 +288,16 @@ def __cuda_array_interface__(self) -> Mapping[str, Any]:
         }
 
         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"] = cuda_array_interface_wrapper(
+                ptr=self.mask_ptr,
+                size=len(self),
+                owner=self.mask,
+                readonly=True,
+                typestr="<t1",
             )
-            output["mask"] = mask
-
         return output
 
     def as_datetime_column(self, dtype: Dtype, **kwargs) -> DatetimeColumn:

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

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-from types import SimpleNamespace
 from typing import (
     Any,
     Callable,
@@ -36,7 +35,7 @@
     is_number,
     is_scalar,
 )
-from cudf.core.buffer import Buffer, as_buffer
+from cudf.core.buffer import Buffer, as_buffer, cuda_array_interface_wrapper
 from cudf.core.column import (
     ColumnBase,
     as_column,
@@ -175,19 +174,16 @@ def __cuda_array_interface__(self) -> Mapping[str, Any]:
         }
 
         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"] = cuda_array_interface_wrapper(
+                ptr=self.mask_ptr,
+                size=len(self),
+                owner=self.mask,
+                readonly=True,
+                typestr="<t1",
             )
-            output["mask"] = mask
 
         return output