11from __future__ import annotations
22
3- from typing import TYPE_CHECKING
3+ from typing import (
4+ TYPE_CHECKING ,
5+ Literal ,
6+ )
47
58import numpy as np
69
10+ from pandas ._config import using_string_dtype
11+
12+ from pandas ._libs import lib
713from pandas .compat import pa_version_under18p0
814from pandas .compat ._optional import import_optional_dependency
915
1218if TYPE_CHECKING :
1319 from collections .abc import Callable
1420
21+ import pyarrow
22+
23+ from pandas ._typing import DtypeBackend
24+
1525
1626def _arrow_dtype_mapping () -> dict :
1727 pa = import_optional_dependency ("pyarrow" )
@@ -33,7 +43,7 @@ def _arrow_dtype_mapping() -> dict:
3343 }
3444
3545
36- def arrow_string_types_mapper () -> Callable :
46+ def _arrow_string_types_mapper () -> Callable :
3747 pa = import_optional_dependency ("pyarrow" )
3848
3949 mapping = {
@@ -44,3 +54,31 @@ def arrow_string_types_mapper() -> Callable:
4454 mapping [pa .string_view ()] = pd .StringDtype (na_value = np .nan )
4555
4656 return mapping .get
57+
58+
59+ def arrow_table_to_pandas (
60+ table : pyarrow .Table ,
61+ dtype_backend : DtypeBackend | Literal ["numpy" ] | lib .NoDefault = lib .no_default ,
62+ null_to_int64 : bool = False ,
63+ ) -> pd .DataFrame :
64+ pa = import_optional_dependency ("pyarrow" )
65+
66+ types_mapper : type [pd .ArrowDtype ] | None | Callable
67+ if dtype_backend == "numpy_nullable" :
68+ mapping = _arrow_dtype_mapping ()
69+ if null_to_int64 :
70+ # Modify the default mapping to also map null to Int64
71+ # (to match other engines - only for CSV parser)
72+ mapping [pa .null ()] = pd .Int64Dtype ()
73+ types_mapper = mapping .get
74+ elif dtype_backend == "pyarrow" :
75+ types_mapper = pd .ArrowDtype
76+ elif using_string_dtype ():
77+ types_mapper = _arrow_string_types_mapper ()
78+ elif dtype_backend is lib .no_default or dtype_backend == "numpy" :
79+ types_mapper = None
80+ else :
81+ raise NotImplementedError
82+
83+ df = table .to_pandas (types_mapper = types_mapper )
84+ return df
0 commit comments