1"""
2A verbatim copy (vendored) of the spec from https://github.com/data-apis/dataframe-api
3"""
4
5from __future__ import annotations
6
7from abc import (
8 ABC,
9 abstractmethod,
10)
11import enum
12from typing import (
13 TYPE_CHECKING,
14 Any,
15 TypedDict,
16)
17
18from pandas.util._decorators import set_module
19
20if TYPE_CHECKING:
21 from collections.abc import (
22 Iterable,
23 Sequence,
24 )
25
26
27class DlpackDeviceType(enum.IntEnum):
28 """Integer enum for device type codes matching DLPack."""
29
30 CPU = 1
31 CUDA = 2
32 CPU_PINNED = 3
33 OPENCL = 4
34 VULKAN = 7
35 METAL = 8
36 VPI = 9
37 ROCM = 10
38
39
40class DtypeKind(enum.IntEnum):
41 """
42 Integer enum for data types.
43
44 Attributes
45 ----------
46 INT : int
47 Matches to signed integer data type.
48 UINT : int
49 Matches to unsigned integer data type.
50 FLOAT : int
51 Matches to floating point data type.
52 BOOL : int
53 Matches to boolean data type.
54 STRING : int
55 Matches to string data type (UTF-8 encoded).
56 DATETIME : int
57 Matches to datetime data type.
58 CATEGORICAL : int
59 Matches to categorical data type.
60 """
61
62 INT = 0
63 UINT = 1
64 FLOAT = 2
65 BOOL = 20
66 STRING = 21 # UTF-8
67 DATETIME = 22
68 CATEGORICAL = 23
69
70
71class ColumnNullType(enum.IntEnum):
72 """
73 Integer enum for null type representation.
74
75 Attributes
76 ----------
77 NON_NULLABLE : int
78 Non-nullable column.
79 USE_NAN : int
80 Use explicit float NaN value.
81 USE_SENTINEL : int
82 Sentinel value besides NaN/NaT.
83 USE_BITMASK : int
84 The bit is set/unset representing a null on a certain position.
85 USE_BYTEMASK : int
86 The byte is set/unset representing a null on a certain position.
87 """
88
89 NON_NULLABLE = 0
90 USE_NAN = 1
91 USE_SENTINEL = 2
92 USE_BITMASK = 3
93 USE_BYTEMASK = 4
94
95
96class ColumnBuffers(TypedDict):
97 # first element is a buffer containing the column data;
98 # second element is the data buffer's associated dtype
99 data: tuple[Buffer, Any]
100
101 # first element is a buffer containing mask values indicating missing data;
102 # second element is the mask value buffer's associated dtype.
103 # None if the null representation is not a bit or byte mask
104 validity: tuple[Buffer, Any] | None
105
106 # first element is a buffer containing the offset values for
107 # variable-size binary data (e.g., variable-length strings);
108 # second element is the offsets buffer's associated dtype.
109 # None if the data buffer does not have an associated offsets buffer
110 offsets: tuple[Buffer, Any] | None
111
112
113class CategoricalDescription(TypedDict):
114 # whether the ordering of dictionary indices is semantically meaningful
115 is_ordered: bool
116 # whether a dictionary-style mapping of categorical values to other objects exists
117 is_dictionary: bool
118 # Python-level only (e.g. ``{int: str}``).
119 # None if not a dictionary-style categorical.
120 categories: Column | None
121
122
123class Buffer(ABC):
124 """
125 Data in the buffer is guaranteed to be contiguous in memory.
126
127 Note that there is no dtype attribute present, a buffer can be thought of
128 as simply a block of memory. However, if the column that the buffer is
129 attached to has a dtype that's supported by DLPack and ``__dlpack__`` is
130 implemented, then that dtype information will be contained in the return
131 value from ``__dlpack__``.
132
133 This distinction is useful to support both data exchange via DLPack on a
134 buffer and (b) dtypes like variable-length strings which do not have a
135 fixed number of bytes per element.
136 """
137
138 @property
139 @abstractmethod
140 def bufsize(self) -> int:
141 """
142 Buffer size in bytes.
143 """
144
145 @property
146 @abstractmethod
147 def ptr(self) -> int:
148 """
149 Pointer to start of the buffer as an integer.
150 """
151
152 @abstractmethod
153 def __dlpack__(self):
154 """
155 Produce DLPack capsule (see array API standard).
156
157 Raises:
158
159 - TypeError : if the buffer contains unsupported dtypes.
160 - NotImplementedError : if DLPack support is not implemented
161
162 Useful to have to connect to array libraries. Support optional because
163 it's not completely trivial to implement for a Python-only library.
164 """
165 raise NotImplementedError("__dlpack__")
166
167 @abstractmethod
168 def __dlpack_device__(self) -> tuple[DlpackDeviceType, int | None]:
169 """
170 Device type and device ID for where the data in the buffer resides.
171 Uses device type codes matching DLPack.
172 Note: must be implemented even if ``__dlpack__`` is not.
173 """
174
175
176class Column(ABC):
177 """
178 A column object, with only the methods and properties required by the
179 interchange protocol defined.
180
181 A column can contain one or more chunks. Each chunk can contain up to three
182 buffers - a data buffer, a mask buffer (depending on null representation),
183 and an offsets buffer (if variable-size binary; e.g., variable-length
184 strings).
185
186 TBD: Arrow has a separate "null" dtype, and has no separate mask concept.
187 Instead, it seems to use "children" for both columns with a bit mask,
188 and for nested dtypes. Unclear whether this is elegant or confusing.
189 This design requires checking the null representation explicitly.
190
191 The Arrow design requires checking:
192 1. the ARROW_FLAG_NULLABLE (for sentinel values)
193 2. if a column has two children, combined with one of those children
194 having a null dtype.
195
196 Making the mask concept explicit seems useful. One null dtype would
197 not be enough to cover both bit and byte masks, so that would mean
198 even more checking if we did it the Arrow way.
199
200 TBD: there's also the "chunk" concept here, which is implicit in Arrow as
201 multiple buffers per array (= column here). Semantically it may make
202 sense to have both: chunks were meant for example for lazy evaluation
203 of data which doesn't fit in memory, while multiple buffers per column
204 could also come from doing a selection operation on a single
205 contiguous buffer.
206
207 Given these concepts, one would expect chunks to be all of the same
208 size (say a 10,000 row dataframe could have 10 chunks of 1,000 rows),
209 while multiple buffers could have data-dependent lengths. Not an issue
210 in pandas if one column is backed by a single NumPy array, but in
211 Arrow it seems possible.
212 Are multiple chunks *and* multiple buffers per column necessary for
213 the purposes of this interchange protocol, or must producers either
214 reuse the chunk concept for this or copy the data?
215
216 Note: this Column object can only be produced by ``__dataframe__``, so
217 doesn't need its own version or ``__column__`` protocol.
218 """
219
220 @abstractmethod
221 def size(self) -> int:
222 """
223 Size of the column, in elements.
224
225 Corresponds to DataFrame.num_rows() if column is a single chunk;
226 equal to size of this current chunk otherwise.
227 """
228
229 @property
230 @abstractmethod
231 def offset(self) -> int:
232 """
233 Offset of first element.
234
235 May be > 0 if using chunks; for example for a column with N chunks of
236 equal size M (only the last chunk may be shorter),
237 ``offset = n * M``, ``n = 0 .. N-1``.
238 """
239
240 @property
241 @abstractmethod
242 def dtype(self) -> tuple[DtypeKind, int, str, str]:
243 """
244 Dtype description as a tuple ``(kind, bit-width, format string, endianness)``.
245
246 Bit-width : the number of bits as an integer
247 Format string : data type description format string in Apache Arrow C
248 Data Interface format.
249 Endianness : current only native endianness (``=``) is supported
250
251 Notes:
252 - Kind specifiers are aligned with DLPack where possible (hence the
253 jump to 20, leave enough room for future extension)
254 - Masks must be specified as boolean with either bit width 1 (for bit
255 masks) or 8 (for byte masks).
256 - Dtype width in bits was preferred over bytes
257 - Endianness isn't too useful, but included now in case in the future
258 we need to support non-native endianness
259 - Went with Apache Arrow format strings over NumPy format strings
260 because they're more complete from a dataframe perspective
261 - Format strings are mostly useful for datetime specification, and
262 for categoricals.
263 - For categoricals, the format string describes the type of the
264 categorical in the data buffer. In case of a separate encoding of
265 the categorical (e.g. an integer to string mapping), this can
266 be derived from ``self.describe_categorical``.
267 - Data types not included: complex, Arrow-style null, binary, decimal,
268 and nested (list, struct, map, union) dtypes.
269 """
270
271 @property
272 @abstractmethod
273 def describe_categorical(self) -> CategoricalDescription:
274 """
275 If the dtype is categorical, there are two options:
276 - There are only values in the data buffer.
277 - There is a separate non-categorical Column encoding for categorical values.
278
279 Raises TypeError if the dtype is not categorical
280
281 Returns the dictionary with description on how to interpret the data buffer:
282 - "is_ordered" : bool, whether the ordering of dictionary indices is
283 semantically meaningful.
284 - "is_dictionary" : bool, whether a mapping of
285 categorical values to other objects exists
286 - "categories" : Column representing the (implicit) mapping of indices to
287 category values (e.g. an array of cat1, cat2, ...).
288 None if not a dictionary-style categorical.
289
290 TBD: are there any other in-memory representations that are needed?
291 """
292
293 @property
294 @abstractmethod
295 def describe_null(self) -> tuple[ColumnNullType, Any]:
296 """
297 Return the missing value (or "null") representation the column dtype
298 uses, as a tuple ``(kind, value)``.
299
300 Value : if kind is "sentinel value", the actual value. If kind is a bit
301 mask or a byte mask, the value (0 or 1) indicating a missing value. None
302 otherwise.
303 """
304
305 @property
306 @abstractmethod
307 def null_count(self) -> int | None:
308 """
309 Number of null elements, if known.
310
311 Note: Arrow uses -1 to indicate "unknown", but None seems cleaner.
312 """
313
314 @property
315 @abstractmethod
316 def metadata(self) -> dict[str, Any]:
317 """
318 The metadata for the column. See `DataFrame.metadata` for more details.
319 """
320
321 @abstractmethod
322 def num_chunks(self) -> int:
323 """
324 Return the number of chunks the column consists of.
325 """
326
327 @abstractmethod
328 def get_chunks(self, n_chunks: int | None = None) -> Iterable[Column]:
329 """
330 Return an iterator yielding the chunks.
331
332 See `DataFrame.get_chunks` for details on ``n_chunks``.
333 """
334
335 @abstractmethod
336 def get_buffers(self) -> ColumnBuffers:
337 """
338 Return a dictionary containing the underlying buffers.
339
340 The returned dictionary has the following contents:
341
342 - "data": a two-element tuple whose first element is a buffer
343 containing the data and whose second element is the data
344 buffer's associated dtype.
345 - "validity": a two-element tuple whose first element is a buffer
346 containing mask values indicating missing data and
347 whose second element is the mask value buffer's
348 associated dtype. None if the null representation is
349 not a bit or byte mask.
350 - "offsets": a two-element tuple whose first element is a buffer
351 containing the offset values for variable-size binary
352 data (e.g., variable-length strings) and whose second
353 element is the offsets buffer's associated dtype. None
354 if the data buffer does not have an associated offsets
355 buffer.
356 """
357
358
359# def get_children(self) -> Iterable[Column]:
360# """
361# Children columns underneath the column, each object in this iterator
362# must adhere to the column specification.
363# """
364# pass
365
366
367@set_module("pandas.api.interchange")
368class DataFrame(ABC):
369 """
370 A data frame class, with only the methods required by the interchange
371 protocol defined.
372
373 A "data frame" represents an ordered collection of named columns.
374 A column's "name" must be a unique string.
375 Columns may be accessed by name or by position.
376
377 This could be a public data frame class, or an object with the methods and
378 attributes defined on this DataFrame class could be returned from the
379 ``__dataframe__`` method of a public data frame class in a library adhering
380 to the dataframe interchange protocol specification.
381 """
382
383 version = 0 # version of the protocol
384
385 @abstractmethod
386 def __dataframe__(self, nan_as_null: bool = False, allow_copy: bool = True):
387 """Construct a new interchange object, potentially changing the parameters."""
388
389 @property
390 @abstractmethod
391 def metadata(self) -> dict[str, Any]:
392 """
393 The metadata for the data frame, as a dictionary with string keys. The
394 contents of `metadata` may be anything, they are meant for a library
395 to store information that it needs to, e.g., roundtrip losslessly or
396 for two implementations to share data that is not (yet) part of the
397 interchange protocol specification. For avoiding collisions with other
398 entries, please add name the keys with the name of the library
399 followed by a period and the desired name, e.g, ``pandas.indexcol``.
400 """
401
402 @abstractmethod
403 def num_columns(self) -> int:
404 """
405 Return the number of columns in the DataFrame.
406 """
407
408 @abstractmethod
409 def num_rows(self) -> int | None:
410 # TODO: not happy with Optional, but need to flag it may be expensive
411 # why include it if it may be None - what do we expect consumers
412 # to do here?
413 """
414 Return the number of rows in the DataFrame, if available.
415 """
416
417 @abstractmethod
418 def num_chunks(self) -> int:
419 """
420 Return the number of chunks the DataFrame consists of.
421 """
422
423 @abstractmethod
424 def column_names(self) -> Iterable[str]:
425 """
426 Return an iterator yielding the column names.
427 """
428
429 @abstractmethod
430 def get_column(self, i: int) -> Column:
431 """
432 Return the column at the indicated position.
433 """
434
435 @abstractmethod
436 def get_column_by_name(self, name: str) -> Column:
437 """
438 Return the column whose name is the indicated name.
439 """
440
441 @abstractmethod
442 def get_columns(self) -> Iterable[Column]:
443 """
444 Return an iterator yielding the columns.
445 """
446
447 @abstractmethod
448 def select_columns(self, indices: Sequence[int]) -> DataFrame:
449 """
450 Create a new DataFrame by selecting a subset of columns by index.
451 """
452
453 @abstractmethod
454 def select_columns_by_name(self, names: Sequence[str]) -> DataFrame:
455 """
456 Create a new DataFrame by selecting a subset of columns by name.
457 """
458
459 @abstractmethod
460 def get_chunks(self, n_chunks: int | None = None) -> Iterable[DataFrame]:
461 """
462 Return an iterator yielding the chunks.
463
464 By default (None), yields the chunks that the data is stored as by the
465 producer. If given, ``n_chunks`` must be a multiple of
466 ``self.num_chunks()``, meaning the producer must subdivide each chunk
467 before yielding it.
468 """