3
3
from abc import ABC , abstractmethod
4
4
5
5
class DlpackDeviceType (enum .IntEnum ):
6
+ """Integer enum for device type codes matching DLPack."""
7
+
6
8
CPU = 1
7
9
CUDA = 2
8
10
CPU_PINNED = 3
@@ -13,6 +15,27 @@ class DlpackDeviceType(enum.IntEnum):
13
15
ROCM = 10
14
16
15
17
class DtypeKind (enum .IntEnum ):
18
+ """
19
+ Integer enum for data types.
20
+
21
+ Attributes
22
+ ----------
23
+ INT : int
24
+ Matches to signed integer data type.
25
+ UINT : int
26
+ Matches to unsigned integer data type.
27
+ FLOAT : int
28
+ Matches to floating point data type.
29
+ BOOL : int
30
+ Matches to boolean data type.
31
+ STRING : int
32
+ Matches to string data type (UTF-8 encoded).
33
+ DATETIME : int
34
+ Matches to datetime data type.
35
+ CATEGORICAL : int
36
+ Matches to categorical data type.
37
+ """
38
+
16
39
INT = 0
17
40
UINT = 1
18
41
FLOAT = 2
@@ -22,6 +45,23 @@ class DtypeKind(enum.IntEnum):
22
45
CATEGORICAL = 23
23
46
24
47
class ColumnNullType (enum .IntEnum ):
48
+ """
49
+ Integer enum for null type representation.
50
+
51
+ Attributes
52
+ ----------
53
+ NON_NULLABLE : int
54
+ Non-nullable column.
55
+ USE_NAN : int
56
+ Use explicit float NaN/NaT value.
57
+ USE_SENTINEL : int
58
+ Sentinel value besides NaN/NaT.
59
+ USE_BITMASK : int
60
+ The bit is set/unset representing a null on a certain position.
61
+ USE_BYTEMASK : int
62
+ The byte is set/unset representing a null on a certain position.
63
+ """
64
+
25
65
NON_NULLABLE = 0
26
66
USE_NAN = 1
27
67
USE_SENTINEL = 2
@@ -85,7 +125,7 @@ def __dlpack__(self):
85
125
raise NotImplementedError ("__dlpack__" )
86
126
87
127
@abstractmethod
88
- def __dlpack_device__ (self ) -> Tuple [DlpackDeviceType , int ]:
128
+ def __dlpack_device__ (self ) -> Tuple [DlpackDeviceType , Optional [ int ] ]:
89
129
"""
90
130
Device type and device ID for where the data in the buffer resides.
91
131
Uses device type codes matching DLPack.
@@ -184,7 +224,7 @@ def dtype(self) -> Tuple[DtypeKind, int, str, str]:
184
224
185
225
@property
186
226
@abstractmethod
187
- def describe_categorical (self ) -> Tuple [bool , bool , dict ]:
227
+ def describe_categorical (self ) -> Tuple [bool , bool , Optional [ dict ] ]:
188
228
"""
189
229
If the dtype is categorical, there are two options:
190
230
- There are only values in the data buffer.
0 commit comments