1from __future__ import annotations
2
3import ctypes
4import re
5from typing import (
6 Any,
7 overload,
8)
9import warnings
10
11import numpy as np
12
13from pandas._config import using_string_dtype
14
15from pandas.compat._optional import import_optional_dependency
16from pandas.errors import Pandas4Warning
17from pandas.util._decorators import set_module
18from pandas.util._exceptions import find_stack_level
19
20import pandas as pd
21from pandas.core.interchange.dataframe_protocol import (
22 Buffer,
23 Column,
24 ColumnNullType,
25 DataFrame as DataFrameXchg,
26 DtypeKind,
27)
28from pandas.core.interchange.utils import (
29 ArrowCTypes,
30 Endianness,
31)
32
33_NP_DTYPES: dict[DtypeKind, dict[int, Any]] = {
34 DtypeKind.INT: {8: np.int8, 16: np.int16, 32: np.int32, 64: np.int64},
35 DtypeKind.UINT: {8: np.uint8, 16: np.uint16, 32: np.uint32, 64: np.uint64},
36 DtypeKind.FLOAT: {32: np.float32, 64: np.float64},
37 DtypeKind.BOOL: {1: bool, 8: bool},
38}
39
40
41@set_module("pandas.api.interchange")
42def from_dataframe(df, allow_copy: bool = True) -> pd.DataFrame:
43 """
44 Build a ``pd.DataFrame`` from any DataFrame supporting the interchange protocol.
45
46 .. note::
47
48 For new development, we highly recommend using the Arrow C Data Interface
49 alongside the Arrow PyCapsule Interface instead of the interchange protocol.
50 From pandas 3.0 onwards, `from_dataframe` uses the PyCapsule Interface,
51 only falling back to the interchange protocol if that fails.
52
53 From pandas 4.0 onwards, that fallback will no longer be available and only
54 the PyCapsule Interface will be used.
55
56 .. warning::
57
58 Due to severe implementation issues, we recommend only considering using the
59 interchange protocol in the following cases:
60
61 - converting to pandas: for pandas >= 2.0.3
62 - converting from pandas: for pandas >= 3.0.0
63
64 Parameters
65 ----------
66 df : ArrowStreamExportable or DataFrameXchg
67 Object supporting the Arrow PyCapsule Interface, i.e. an
68 `__arrow_c_stream__` method (``ArrowStreamExportable``), or an object
69 supporting the interchange protocol, i.e. a `__dataframe__` method
70 (``DataFrameXchg``).
71 allow_copy : bool, default: True
72 Whether to allow copying the memory to perform the conversion
73 (if false then zero-copy approach is requested).
74
75 Returns
76 -------
77 pd.DataFrame
78 A pandas DataFrame built from the provided interchange
79 protocol object.
80
81 See Also
82 --------
83 pd.DataFrame : DataFrame class which can be created from various input data
84 formats, including objects that support the interchange protocol.
85
86 Examples
87 --------
88 >>> df_not_necessarily_pandas = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
89 >>> interchange_object = df_not_necessarily_pandas.__dataframe__()
90 >>> interchange_object.column_names()
91 Index(['A', 'B'], dtype='str')
92 >>> df_pandas = pd.api.interchange.from_dataframe(
93 ... interchange_object.select_columns_by_name(["A"])
94 ... )
95 >>> df_pandas
96 A
97 0 1
98 1 2
99
100 These methods (``column_names``, ``select_columns_by_name``) should work
101 for any dataframe library which implements the interchange protocol.
102 """
103 if isinstance(df, pd.DataFrame):
104 return df
105
106 if hasattr(df, "__arrow_c_stream__"):
107 try:
108 pa = import_optional_dependency("pyarrow", min_version="14.0.0")
109 except ImportError:
110 # fallback to _from_dataframe
111 warnings.warn(
112 "Conversion using Arrow PyCapsule Interface failed due to "
113 "missing PyArrow>=14 dependency, falling back to (deprecated) "
114 "interchange protocol. We recommend that you install "
115 "PyArrow>=14.0.0.",
116 UserWarning,
117 stacklevel=find_stack_level(),
118 )
119 else:
120 try:
121 return pa.table(df).to_pandas(zero_copy_only=not allow_copy)
122 except pa.ArrowInvalid as e:
123 raise RuntimeError(e) from e
124
125 if not hasattr(df, "__dataframe__"):
126 raise ValueError("`df` does not support __dataframe__")
127
128 warnings.warn(
129 "The Dataframe Interchange Protocol is deprecated.\n"
130 "For dataframe-agnostic code, you may want to look into:\n"
131 "- Arrow PyCapsule Interface: https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html\n"
132 "- Narwhals: https://github.com/narwhals-dev/narwhals\n",
133 Pandas4Warning,
134 stacklevel=find_stack_level(),
135 )
136
137 return _from_dataframe(
138 df.__dataframe__(allow_copy=allow_copy), allow_copy=allow_copy
139 )
140
141
142def _from_dataframe(df: DataFrameXchg, allow_copy: bool = True) -> pd.DataFrame:
143 """
144 Build a ``pd.DataFrame`` from the DataFrame interchange object.
145
146 Parameters
147 ----------
148 df : DataFrameXchg
149 Object supporting the interchange protocol, i.e. `__dataframe__` method.
150 allow_copy : bool, default: True
151 Whether to allow copying the memory to perform the conversion
152 (if false then zero-copy approach is requested).
153
154 Returns
155 -------
156 pd.DataFrame
157 """
158 pandas_dfs = []
159 for chunk in df.get_chunks():
160 pandas_df = protocol_df_chunk_to_pandas(chunk)
161 pandas_dfs.append(pandas_df)
162
163 if not allow_copy and len(pandas_dfs) > 1:
164 raise RuntimeError(
165 "To join chunks a copy is required which is forbidden by allow_copy=False"
166 )
167 if not pandas_dfs:
168 pandas_df = protocol_df_chunk_to_pandas(df)
169 elif len(pandas_dfs) == 1:
170 pandas_df = pandas_dfs[0]
171 else:
172 pandas_df = pd.concat(pandas_dfs, axis=0, ignore_index=True, copy=False)
173
174 index_obj = df.metadata.get("pandas.index", None)
175 if index_obj is not None:
176 pandas_df.index = index_obj
177
178 return pandas_df
179
180
181def protocol_df_chunk_to_pandas(df: DataFrameXchg) -> pd.DataFrame:
182 """
183 Convert interchange protocol chunk to ``pd.DataFrame``.
184
185 Parameters
186 ----------
187 df : DataFrameXchg
188
189 Returns
190 -------
191 pd.DataFrame
192 """
193 columns: dict[str, Any] = {}
194 buffers = [] # hold on to buffers, keeps memory alive
195 for name in df.column_names():
196 if not isinstance(name, str):
197 raise ValueError(f"Column {name} is not a string")
198 if name in columns:
199 raise ValueError(f"Column {name} is not unique")
200 col = df.get_column_by_name(name)
201 dtype = col.dtype[0]
202 if dtype in (
203 DtypeKind.INT,
204 DtypeKind.UINT,
205 DtypeKind.FLOAT,
206 DtypeKind.BOOL,
207 ):
208 columns[name], buf = primitive_column_to_ndarray(col)
209 elif dtype == DtypeKind.CATEGORICAL:
210 columns[name], buf = categorical_column_to_series(col)
211 elif dtype == DtypeKind.STRING:
212 columns[name], buf = string_column_to_ndarray(col)
213 elif dtype == DtypeKind.DATETIME:
214 columns[name], buf = datetime_column_to_ndarray(col)
215 else:
216 raise NotImplementedError(f"Data type {dtype} not handled yet")
217
218 buffers.append(buf)
219
220 pandas_df = pd.DataFrame(columns)
221 pandas_df.attrs["_INTERCHANGE_PROTOCOL_BUFFERS"] = buffers
222 return pandas_df
223
224
225def primitive_column_to_ndarray(col: Column) -> tuple[np.ndarray, Any]:
226 """
227 Convert a column holding one of the primitive dtypes to a NumPy array.
228
229 A primitive type is one of: int, uint, float, bool.
230
231 Parameters
232 ----------
233 col : Column
234
235 Returns
236 -------
237 tuple
238 Tuple of np.ndarray holding the data and the memory owner object
239 that keeps the memory alive.
240 """
241 buffers = col.get_buffers()
242
243 data_buff, data_dtype = buffers["data"]
244 data = buffer_to_ndarray(
245 data_buff, data_dtype, offset=col.offset, length=col.size()
246 )
247
248 data = set_nulls(data, col, buffers["validity"])
249 return data, buffers
250
251
252def categorical_column_to_series(col: Column) -> tuple[pd.Series, Any]:
253 """
254 Convert a column holding categorical data to a pandas Series.
255
256 Parameters
257 ----------
258 col : Column
259
260 Returns
261 -------
262 tuple
263 Tuple of pd.Series holding the data and the memory owner object
264 that keeps the memory alive.
265 """
266 categorical = col.describe_categorical
267
268 if not categorical["is_dictionary"]:
269 raise NotImplementedError("Non-dictionary categoricals not supported yet")
270
271 cat_column = categorical["categories"]
272 if hasattr(cat_column, "_col"):
273 # Item "Column" of "Optional[Column]" has no attribute "_col"
274 # Item "None" of "Optional[Column]" has no attribute "_col"
275 categories = np.array(cat_column._col) # type: ignore[union-attr]
276 else:
277 raise NotImplementedError(
278 "Interchanging categorical columns isn't supported yet, and our "
279 "fallback of using the `col._col` attribute (a ndarray) failed."
280 )
281 buffers = col.get_buffers()
282
283 codes_buff, codes_dtype = buffers["data"]
284 codes = buffer_to_ndarray(
285 codes_buff, codes_dtype, offset=col.offset, length=col.size()
286 )
287
288 # Doing module in order to not get ``IndexError`` for
289 # out-of-bounds sentinel values in `codes`
290 if len(categories) > 0:
291 values = categories[codes % len(categories)]
292 else:
293 values = codes
294
295 cat = pd.Categorical(
296 values, categories=categories, ordered=categorical["is_ordered"]
297 )
298 data = pd.Series(cat)
299
300 data = set_nulls(data, col, buffers["validity"])
301 return data, buffers
302
303
304def string_column_to_ndarray(col: Column) -> tuple[np.ndarray, Any]:
305 """
306 Convert a column holding string data to a NumPy array.
307
308 Parameters
309 ----------
310 col : Column
311
312 Returns
313 -------
314 tuple
315 Tuple of np.ndarray holding the data and the memory owner object
316 that keeps the memory alive.
317 """
318 null_kind, sentinel_val = col.describe_null
319
320 if null_kind not in (
321 ColumnNullType.NON_NULLABLE,
322 ColumnNullType.USE_BITMASK,
323 ColumnNullType.USE_BYTEMASK,
324 ):
325 raise NotImplementedError(
326 f"{null_kind} null kind is not yet supported for string columns."
327 )
328
329 buffers = col.get_buffers()
330
331 assert buffers["offsets"], "String buffers must contain offsets"
332 # Retrieve the data buffer containing the UTF-8 code units
333 data_buff, _ = buffers["data"]
334 # We're going to reinterpret the buffer as uint8, so make sure we can do it safely
335 assert col.dtype[2] in (
336 ArrowCTypes.STRING,
337 ArrowCTypes.LARGE_STRING,
338 ) # format_str == utf-8
339 # Convert the buffers to NumPy arrays. In order to go from STRING to
340 # an equivalent ndarray, we claim that the buffer is uint8 (i.e., a byte array)
341 data_dtype = (
342 DtypeKind.UINT,
343 8,
344 ArrowCTypes.UINT8,
345 Endianness.NATIVE,
346 )
347 # Specify zero offset as we don't want to chunk the string data
348 data = buffer_to_ndarray(data_buff, data_dtype, offset=0, length=data_buff.bufsize)
349
350 # Retrieve the offsets buffer containing the index offsets demarcating
351 # the beginning and the ending of each string
352 offset_buff, offset_dtype = buffers["offsets"]
353 # Offsets buffer contains start-stop positions of strings in the data buffer,
354 # meaning that it has more elements than in the data buffer, do `col.size() + 1`
355 # here to pass a proper offsets buffer size
356 offsets = buffer_to_ndarray(
357 offset_buff, offset_dtype, offset=col.offset, length=col.size() + 1
358 )
359
360 null_pos = None
361 if null_kind in (ColumnNullType.USE_BITMASK, ColumnNullType.USE_BYTEMASK):
362 validity = buffers["validity"]
363 if validity is not None:
364 valid_buff, valid_dtype = validity
365 null_pos = buffer_to_ndarray(
366 valid_buff, valid_dtype, offset=col.offset, length=col.size()
367 )
368 if sentinel_val == 0:
369 null_pos = ~null_pos
370
371 # Assemble the strings from the code units
372 str_list: list[float | str | None] = [None] * col.size()
373 for i in range(col.size()):
374 # Check for missing values
375 if null_pos is not None and null_pos[i]:
376 str_list[i] = np.nan
377 continue
378
379 # Extract a range of code units
380 units = data[offsets[i] : offsets[i + 1]]
381
382 # Convert the list of code units to bytes
383 str_bytes = bytes(units)
384
385 # Create the string
386 string = str_bytes.decode(encoding="utf-8")
387
388 # Add to our list of strings
389 str_list[i] = string
390
391 if using_string_dtype():
392 res = pd.Series(str_list, dtype="str")
393 else:
394 res = np.asarray(str_list, dtype="object") # type: ignore[assignment]
395
396 return res, buffers # type: ignore[return-value]
397
398
399def parse_datetime_format_str(format_str, data) -> pd.Series | np.ndarray:
400 """Parse datetime `format_str` to interpret the `data`."""
401 # timestamp 'ts{unit}:tz'
402 timestamp_meta = re.match(r"ts([smun]):(.*)", format_str)
403 if timestamp_meta:
404 unit, tz = timestamp_meta.group(1), timestamp_meta.group(2)
405 if unit != "s":
406 # the format string describes only a first letter of the unit, so
407 # add one extra letter to convert the unit to numpy-style:
408 # 'm' -> 'ms', 'u' -> 'us', 'n' -> 'ns'
409 unit += "s"
410 data = data.astype(f"datetime64[{unit}]")
411 if tz != "":
412 data = pd.Series(data).dt.tz_localize("UTC").dt.tz_convert(tz)
413 return data
414
415 # date 'td{Days/Ms}'
416 date_meta = re.match(r"td([Dm])", format_str)
417 if date_meta:
418 unit = date_meta.group(1)
419 if unit == "D":
420 # NumPy doesn't support DAY unit, so converting days to seconds
421 # (converting to uint64 to avoid overflow)
422 data = (data.astype(np.uint64) * (24 * 60 * 60)).astype("datetime64[s]")
423 elif unit == "m":
424 data = data.astype("datetime64[ms]")
425 else:
426 raise NotImplementedError(f"Date unit is not supported: {unit}")
427 return data
428
429 raise NotImplementedError(f"DateTime kind is not supported: {format_str}")
430
431
432def datetime_column_to_ndarray(col: Column) -> tuple[np.ndarray | pd.Series, Any]:
433 """
434 Convert a column holding DateTime data to a NumPy array.
435
436 Parameters
437 ----------
438 col : Column
439
440 Returns
441 -------
442 tuple
443 Tuple of np.ndarray holding the data and the memory owner object
444 that keeps the memory alive.
445 """
446 buffers = col.get_buffers()
447
448 _, col_bit_width, format_str, _ = col.dtype
449 dbuf, _ = buffers["data"]
450 # Consider dtype being `uint` to get number of units passed since the 01.01.1970
451
452 data = buffer_to_ndarray(
453 dbuf,
454 (
455 DtypeKind.INT,
456 col_bit_width,
457 getattr(ArrowCTypes, f"INT{col_bit_width}"),
458 Endianness.NATIVE,
459 ),
460 offset=col.offset,
461 length=col.size(),
462 )
463
464 data = parse_datetime_format_str(format_str, data) # type: ignore[assignment]
465 data = set_nulls(data, col, buffers["validity"])
466 return data, buffers
467
468
469def buffer_to_ndarray(
470 buffer: Buffer,
471 dtype: tuple[DtypeKind, int, str, str],
472 *,
473 length: int,
474 offset: int = 0,
475) -> np.ndarray:
476 """
477 Build a NumPy array from the passed buffer.
478
479 Parameters
480 ----------
481 buffer : Buffer
482 Buffer to build a NumPy array from.
483 dtype : tuple
484 Data type of the buffer conforming protocol dtypes format.
485 offset : int, default: 0
486 Number of elements to offset from the start of the buffer.
487 length : int, optional
488 If the buffer is a bit-mask, specifies a number of bits to read
489 from the buffer. Has no effect otherwise.
490
491 Returns
492 -------
493 np.ndarray
494
495 Notes
496 -----
497 The returned array doesn't own the memory. The caller of this function is
498 responsible for keeping the memory owner object alive as long as
499 the returned NumPy array is being used.
500 """
501 kind, bit_width, _, _ = dtype
502
503 column_dtype = _NP_DTYPES.get(kind, {}).get(bit_width, None)
504 if column_dtype is None:
505 raise NotImplementedError(f"Conversion for {dtype} is not yet supported.")
506
507 # TODO: No DLPack yet, so need to construct a new ndarray from the data pointer
508 # and size in the buffer plus the dtype on the column. Use DLPack as NumPy supports
509 # it since https://github.com/numpy/numpy/pull/19083
510 ctypes_type = np.ctypeslib.as_ctypes_type(column_dtype)
511
512 if bit_width == 1:
513 assert length is not None, "`length` must be specified for a bit-mask buffer."
514 pa = import_optional_dependency("pyarrow")
515 arr = pa.BooleanArray.from_buffers(
516 pa.bool_(),
517 length,
518 [None, pa.foreign_buffer(buffer.ptr, length)],
519 offset=offset,
520 )
521 return np.asarray(arr)
522 else:
523 data_pointer = ctypes.cast(
524 buffer.ptr + (offset * bit_width // 8), ctypes.POINTER(ctypes_type)
525 )
526 if length > 0:
527 return np.ctypeslib.as_array(data_pointer, shape=(length,))
528 return np.array([], dtype=ctypes_type)
529
530
531@overload
532def set_nulls(
533 data: np.ndarray,
534 col: Column,
535 validity: tuple[Buffer, tuple[DtypeKind, int, str, str]] | None,
536 allow_modify_inplace: bool = ...,
537) -> np.ndarray: ...
538
539
540@overload
541def set_nulls(
542 data: pd.Series,
543 col: Column,
544 validity: tuple[Buffer, tuple[DtypeKind, int, str, str]] | None,
545 allow_modify_inplace: bool = ...,
546) -> pd.Series: ...
547
548
549@overload
550def set_nulls(
551 data: np.ndarray | pd.Series,
552 col: Column,
553 validity: tuple[Buffer, tuple[DtypeKind, int, str, str]] | None,
554 allow_modify_inplace: bool = ...,
555) -> np.ndarray | pd.Series: ...
556
557
558def set_nulls(
559 data: np.ndarray | pd.Series,
560 col: Column,
561 validity: tuple[Buffer, tuple[DtypeKind, int, str, str]] | None,
562 allow_modify_inplace: bool = True,
563) -> np.ndarray | pd.Series:
564 """
565 Set null values for the data according to the column null kind.
566
567 Parameters
568 ----------
569 data : np.ndarray or pd.Series
570 Data to set nulls in.
571 col : Column
572 Column object that describes the `data`.
573 validity : tuple(Buffer, dtype) or None
574 The return value of ``col.buffers()``. We do not access the ``col.buffers()``
575 here to not take the ownership of the memory of buffer objects.
576 allow_modify_inplace : bool, default: True
577 Whether to modify the `data` inplace when zero-copy is possible (True) or always
578 modify a copy of the `data` (False).
579
580 Returns
581 -------
582 np.ndarray or pd.Series
583 Data with the nulls being set.
584 """
585 if validity is None:
586 return data
587 null_kind, sentinel_val = col.describe_null
588 null_pos = None
589
590 if null_kind == ColumnNullType.USE_SENTINEL:
591 null_pos = pd.Series(data) == sentinel_val
592 elif null_kind in (ColumnNullType.USE_BITMASK, ColumnNullType.USE_BYTEMASK):
593 valid_buff, valid_dtype = validity
594 null_pos = buffer_to_ndarray(
595 valid_buff, valid_dtype, offset=col.offset, length=col.size()
596 )
597 if sentinel_val == 0:
598 null_pos = ~null_pos
599 elif null_kind in (ColumnNullType.NON_NULLABLE, ColumnNullType.USE_NAN):
600 pass
601 else:
602 raise NotImplementedError(f"Null kind {null_kind} is not yet supported.")
603
604 if null_pos is not None and np.any(null_pos):
605 if not allow_modify_inplace:
606 data = data.copy()
607 try:
608 data[null_pos] = None
609 except TypeError:
610 # TypeError happens if the `data` dtype appears to be non-nullable
611 # in numpy notation (bool, int, uint). If this happens,
612 # cast the `data` to nullable float dtype.
613 data = data.astype(float)
614 data[null_pos] = None
615
616 return data