Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3import ctypes
4import re
5from typing import Any
7import numpy as np
9from pandas.compat._optional import import_optional_dependency
11import pandas as pd
12from pandas.core.interchange.dataframe_protocol import (
13 Buffer,
14 Column,
15 ColumnNullType,
16 DataFrame as DataFrameXchg,
17 DtypeKind,
18)
19from pandas.core.interchange.utils import (
20 ArrowCTypes,
21 Endianness,
22)
24_NP_DTYPES: dict[DtypeKind, dict[int, Any]] = {
25 DtypeKind.INT: {8: np.int8, 16: np.int16, 32: np.int32, 64: np.int64},
26 DtypeKind.UINT: {8: np.uint8, 16: np.uint16, 32: np.uint32, 64: np.uint64},
27 DtypeKind.FLOAT: {32: np.float32, 64: np.float64},
28 DtypeKind.BOOL: {1: bool, 8: bool},
29}
32def from_dataframe(df, allow_copy: bool = True) -> pd.DataFrame:
33 """
34 Build a ``pd.DataFrame`` from any DataFrame supporting the interchange protocol.
36 Parameters
37 ----------
38 df : DataFrameXchg
39 Object supporting the interchange protocol, i.e. `__dataframe__` method.
40 allow_copy : bool, default: True
41 Whether to allow copying the memory to perform the conversion
42 (if false then zero-copy approach is requested).
44 Returns
45 -------
46 pd.DataFrame
47 """
48 if isinstance(df, pd.DataFrame):
49 return df
51 if not hasattr(df, "__dataframe__"):
52 raise ValueError("`df` does not support __dataframe__")
54 return _from_dataframe(df.__dataframe__(allow_copy=allow_copy))
57def _from_dataframe(df: DataFrameXchg, allow_copy: bool = True):
58 """
59 Build a ``pd.DataFrame`` from the DataFrame interchange object.
61 Parameters
62 ----------
63 df : DataFrameXchg
64 Object supporting the interchange protocol, i.e. `__dataframe__` method.
65 allow_copy : bool, default: True
66 Whether to allow copying the memory to perform the conversion
67 (if false then zero-copy approach is requested).
69 Returns
70 -------
71 pd.DataFrame
72 """
73 pandas_dfs = []
74 for chunk in df.get_chunks():
75 pandas_df = protocol_df_chunk_to_pandas(chunk)
76 pandas_dfs.append(pandas_df)
78 if not allow_copy and len(pandas_dfs) > 1:
79 raise RuntimeError(
80 "To join chunks a copy is required which is forbidden by allow_copy=False"
81 )
82 if len(pandas_dfs) == 1:
83 pandas_df = pandas_dfs[0]
84 else:
85 pandas_df = pd.concat(pandas_dfs, axis=0, ignore_index=True, copy=False)
87 index_obj = df.metadata.get("pandas.index", None)
88 if index_obj is not None:
89 pandas_df.index = index_obj
91 return pandas_df
94def protocol_df_chunk_to_pandas(df: DataFrameXchg) -> pd.DataFrame:
95 """
96 Convert interchange protocol chunk to ``pd.DataFrame``.
98 Parameters
99 ----------
100 df : DataFrameXchg
102 Returns
103 -------
104 pd.DataFrame
105 """
106 # We need a dict of columns here, with each column being a NumPy array (at
107 # least for now, deal with non-NumPy dtypes later).
108 columns: dict[str, Any] = {}
109 buffers = [] # hold on to buffers, keeps memory alive
110 for name in df.column_names():
111 if not isinstance(name, str):
112 raise ValueError(f"Column {name} is not a string")
113 if name in columns:
114 raise ValueError(f"Column {name} is not unique")
115 col = df.get_column_by_name(name)
116 dtype = col.dtype[0]
117 if dtype in (
118 DtypeKind.INT,
119 DtypeKind.UINT,
120 DtypeKind.FLOAT,
121 DtypeKind.BOOL,
122 ):
123 columns[name], buf = primitive_column_to_ndarray(col)
124 elif dtype == DtypeKind.CATEGORICAL:
125 columns[name], buf = categorical_column_to_series(col)
126 elif dtype == DtypeKind.STRING:
127 columns[name], buf = string_column_to_ndarray(col)
128 elif dtype == DtypeKind.DATETIME:
129 columns[name], buf = datetime_column_to_ndarray(col)
130 else:
131 raise NotImplementedError(f"Data type {dtype} not handled yet")
133 buffers.append(buf)
135 pandas_df = pd.DataFrame(columns)
136 pandas_df.attrs["_INTERCHANGE_PROTOCOL_BUFFERS"] = buffers
137 return pandas_df
140def primitive_column_to_ndarray(col: Column) -> tuple[np.ndarray, Any]:
141 """
142 Convert a column holding one of the primitive dtypes to a NumPy array.
144 A primitive type is one of: int, uint, float, bool.
146 Parameters
147 ----------
148 col : Column
150 Returns
151 -------
152 tuple
153 Tuple of np.ndarray holding the data and the memory owner object
154 that keeps the memory alive.
155 """
156 buffers = col.get_buffers()
158 data_buff, data_dtype = buffers["data"]
159 data = buffer_to_ndarray(
160 data_buff, data_dtype, offset=col.offset, length=col.size()
161 )
163 data = set_nulls(data, col, buffers["validity"])
164 return data, buffers
167def categorical_column_to_series(col: Column) -> tuple[pd.Series, Any]:
168 """
169 Convert a column holding categorical data to a pandas Series.
171 Parameters
172 ----------
173 col : Column
175 Returns
176 -------
177 tuple
178 Tuple of pd.Series holding the data and the memory owner object
179 that keeps the memory alive.
180 """
181 categorical = col.describe_categorical
183 if not categorical["is_dictionary"]:
184 raise NotImplementedError("Non-dictionary categoricals not supported yet")
186 cat_column = categorical["categories"]
187 if hasattr(cat_column, "_col"):
188 # Item "Column" of "Optional[Column]" has no attribute "_col"
189 # Item "None" of "Optional[Column]" has no attribute "_col"
190 categories = np.array(cat_column._col) # type: ignore[union-attr]
191 else:
192 raise NotImplementedError(
193 "Interchanging categorical columns isn't supported yet, and our "
194 "fallback of using the `col._col` attribute (a ndarray) failed."
195 )
196 buffers = col.get_buffers()
198 codes_buff, codes_dtype = buffers["data"]
199 codes = buffer_to_ndarray(
200 codes_buff, codes_dtype, offset=col.offset, length=col.size()
201 )
203 # Doing module in order to not get ``IndexError`` for
204 # out-of-bounds sentinel values in `codes`
205 if len(categories) > 0:
206 values = categories[codes % len(categories)]
207 else:
208 values = codes
210 cat = pd.Categorical(
211 values, categories=categories, ordered=categorical["is_ordered"]
212 )
213 data = pd.Series(cat)
215 data = set_nulls(data, col, buffers["validity"])
216 return data, buffers
219def string_column_to_ndarray(col: Column) -> tuple[np.ndarray, Any]:
220 """
221 Convert a column holding string data to a NumPy array.
223 Parameters
224 ----------
225 col : Column
227 Returns
228 -------
229 tuple
230 Tuple of np.ndarray holding the data and the memory owner object
231 that keeps the memory alive.
232 """
233 null_kind, sentinel_val = col.describe_null
235 if null_kind not in (
236 ColumnNullType.NON_NULLABLE,
237 ColumnNullType.USE_BITMASK,
238 ColumnNullType.USE_BYTEMASK,
239 ):
240 raise NotImplementedError(
241 f"{null_kind} null kind is not yet supported for string columns."
242 )
244 buffers = col.get_buffers()
246 assert buffers["offsets"], "String buffers must contain offsets"
247 # Retrieve the data buffer containing the UTF-8 code units
248 data_buff, protocol_data_dtype = buffers["data"]
249 # We're going to reinterpret the buffer as uint8, so make sure we can do it safely
250 assert protocol_data_dtype[1] == 8
251 assert protocol_data_dtype[2] in (
252 ArrowCTypes.STRING,
253 ArrowCTypes.LARGE_STRING,
254 ) # format_str == utf-8
255 # Convert the buffers to NumPy arrays. In order to go from STRING to
256 # an equivalent ndarray, we claim that the buffer is uint8 (i.e., a byte array)
257 data_dtype = (
258 DtypeKind.UINT,
259 8,
260 ArrowCTypes.UINT8,
261 Endianness.NATIVE,
262 )
263 # Specify zero offset as we don't want to chunk the string data
264 data = buffer_to_ndarray(data_buff, data_dtype, offset=0, length=data_buff.bufsize)
266 # Retrieve the offsets buffer containing the index offsets demarcating
267 # the beginning and the ending of each string
268 offset_buff, offset_dtype = buffers["offsets"]
269 # Offsets buffer contains start-stop positions of strings in the data buffer,
270 # meaning that it has more elements than in the data buffer, do `col.size() + 1`
271 # here to pass a proper offsets buffer size
272 offsets = buffer_to_ndarray(
273 offset_buff, offset_dtype, offset=col.offset, length=col.size() + 1
274 )
276 null_pos = None
277 if null_kind in (ColumnNullType.USE_BITMASK, ColumnNullType.USE_BYTEMASK):
278 assert buffers["validity"], "Validity buffers cannot be empty for masks"
279 valid_buff, valid_dtype = buffers["validity"]
280 null_pos = buffer_to_ndarray(
281 valid_buff, valid_dtype, offset=col.offset, length=col.size()
282 )
283 if sentinel_val == 0:
284 null_pos = ~null_pos
286 # Assemble the strings from the code units
287 str_list: list[None | float | str] = [None] * col.size()
288 for i in range(col.size()):
289 # Check for missing values
290 if null_pos is not None and null_pos[i]:
291 str_list[i] = np.nan
292 continue
294 # Extract a range of code units
295 units = data[offsets[i] : offsets[i + 1]]
297 # Convert the list of code units to bytes
298 str_bytes = bytes(units)
300 # Create the string
301 string = str_bytes.decode(encoding="utf-8")
303 # Add to our list of strings
304 str_list[i] = string
306 # Convert the string list to a NumPy array
307 return np.asarray(str_list, dtype="object"), buffers
310def parse_datetime_format_str(format_str, data):
311 """Parse datetime `format_str` to interpret the `data`."""
312 # timestamp 'ts{unit}:tz'
313 timestamp_meta = re.match(r"ts([smun]):(.*)", format_str)
314 if timestamp_meta:
315 unit, tz = timestamp_meta.group(1), timestamp_meta.group(2)
316 if tz != "":
317 raise NotImplementedError("Timezones are not supported yet")
318 if unit != "s":
319 # the format string describes only a first letter of the unit, so
320 # add one extra letter to convert the unit to numpy-style:
321 # 'm' -> 'ms', 'u' -> 'us', 'n' -> 'ns'
322 unit += "s"
323 data = data.astype(f"datetime64[{unit}]")
324 return data
326 # date 'td{Days/Ms}'
327 date_meta = re.match(r"td([Dm])", format_str)
328 if date_meta:
329 unit = date_meta.group(1)
330 if unit == "D":
331 # NumPy doesn't support DAY unit, so converting days to seconds
332 # (converting to uint64 to avoid overflow)
333 data = (data.astype(np.uint64) * (24 * 60 * 60)).astype("datetime64[s]")
334 elif unit == "m":
335 data = data.astype("datetime64[ms]")
336 else:
337 raise NotImplementedError(f"Date unit is not supported: {unit}")
338 return data
340 raise NotImplementedError(f"DateTime kind is not supported: {format_str}")
343def datetime_column_to_ndarray(col: Column) -> tuple[np.ndarray, Any]:
344 """
345 Convert a column holding DateTime data to a NumPy array.
347 Parameters
348 ----------
349 col : Column
351 Returns
352 -------
353 tuple
354 Tuple of np.ndarray holding the data and the memory owner object
355 that keeps the memory alive.
356 """
357 buffers = col.get_buffers()
359 _, _, format_str, _ = col.dtype
360 dbuf, dtype = buffers["data"]
361 # Consider dtype being `uint` to get number of units passed since the 01.01.1970
362 data = buffer_to_ndarray(
363 dbuf,
364 (
365 DtypeKind.UINT,
366 dtype[1],
367 getattr(ArrowCTypes, f"UINT{dtype[1]}"),
368 Endianness.NATIVE,
369 ),
370 offset=col.offset,
371 length=col.size(),
372 )
374 data = parse_datetime_format_str(format_str, data)
375 data = set_nulls(data, col, buffers["validity"])
376 return data, buffers
379def buffer_to_ndarray(
380 buffer: Buffer,
381 dtype: tuple[DtypeKind, int, str, str],
382 *,
383 length: int,
384 offset: int = 0,
385) -> np.ndarray:
386 """
387 Build a NumPy array from the passed buffer.
389 Parameters
390 ----------
391 buffer : Buffer
392 Buffer to build a NumPy array from.
393 dtype : tuple
394 Data type of the buffer conforming protocol dtypes format.
395 offset : int, default: 0
396 Number of elements to offset from the start of the buffer.
397 length : int, optional
398 If the buffer is a bit-mask, specifies a number of bits to read
399 from the buffer. Has no effect otherwise.
401 Returns
402 -------
403 np.ndarray
405 Notes
406 -----
407 The returned array doesn't own the memory. The caller of this function is
408 responsible for keeping the memory owner object alive as long as
409 the returned NumPy array is being used.
410 """
411 kind, bit_width, _, _ = dtype
413 column_dtype = _NP_DTYPES.get(kind, {}).get(bit_width, None)
414 if column_dtype is None:
415 raise NotImplementedError(f"Conversion for {dtype} is not yet supported.")
417 # TODO: No DLPack yet, so need to construct a new ndarray from the data pointer
418 # and size in the buffer plus the dtype on the column. Use DLPack as NumPy supports
419 # it since https://github.com/numpy/numpy/pull/19083
420 ctypes_type = np.ctypeslib.as_ctypes_type(column_dtype)
422 if bit_width == 1:
423 assert length is not None, "`length` must be specified for a bit-mask buffer."
424 pa = import_optional_dependency("pyarrow")
425 arr = pa.BooleanArray.from_buffers(
426 pa.bool_(),
427 length,
428 [None, pa.foreign_buffer(buffer.ptr, length)],
429 offset=offset,
430 )
431 return np.asarray(arr)
432 else:
433 data_pointer = ctypes.cast(
434 buffer.ptr + (offset * bit_width // 8), ctypes.POINTER(ctypes_type)
435 )
436 return np.ctypeslib.as_array(
437 data_pointer,
438 shape=(length,),
439 )
442def set_nulls(
443 data: np.ndarray | pd.Series,
444 col: Column,
445 validity: tuple[Buffer, tuple[DtypeKind, int, str, str]] | None,
446 allow_modify_inplace: bool = True,
447):
448 """
449 Set null values for the data according to the column null kind.
451 Parameters
452 ----------
453 data : np.ndarray or pd.Series
454 Data to set nulls in.
455 col : Column
456 Column object that describes the `data`.
457 validity : tuple(Buffer, dtype) or None
458 The return value of ``col.buffers()``. We do not access the ``col.buffers()``
459 here to not take the ownership of the memory of buffer objects.
460 allow_modify_inplace : bool, default: True
461 Whether to modify the `data` inplace when zero-copy is possible (True) or always
462 modify a copy of the `data` (False).
464 Returns
465 -------
466 np.ndarray or pd.Series
467 Data with the nulls being set.
468 """
469 null_kind, sentinel_val = col.describe_null
470 null_pos = None
472 if null_kind == ColumnNullType.USE_SENTINEL:
473 null_pos = pd.Series(data) == sentinel_val
474 elif null_kind in (ColumnNullType.USE_BITMASK, ColumnNullType.USE_BYTEMASK):
475 assert validity, "Expected to have a validity buffer for the mask"
476 valid_buff, valid_dtype = validity
477 null_pos = buffer_to_ndarray(
478 valid_buff, valid_dtype, offset=col.offset, length=col.size()
479 )
480 if sentinel_val == 0:
481 null_pos = ~null_pos
482 elif null_kind in (ColumnNullType.NON_NULLABLE, ColumnNullType.USE_NAN):
483 pass
484 else:
485 raise NotImplementedError(f"Null kind {null_kind} is not yet supported.")
487 if null_pos is not None and np.any(null_pos):
488 if not allow_modify_inplace:
489 data = data.copy()
490 try:
491 data[null_pos] = None
492 except TypeError:
493 # TypeError happens if the `data` dtype appears to be non-nullable
494 # in numpy notation (bool, int, uint). If this happens,
495 # cast the `data` to nullable float dtype.
496 data = data.astype(float)
497 data[null_pos] = None
499 return data