feat: Support reading with PyArrow-backed types (#2292)

Author: LeonLuttenberger

awswrangler/_config.py

@@ -8,6 +8,7 @@
 
 import botocore.config
 import pandas as pd
+from typing_extensions import Literal
 
 from awswrangler import exceptions
 from awswrangler.typing import AthenaCacheSettings
@@ -48,6 +49,7 @@ class _ConfigArg(NamedTuple):
     "workgroup": _ConfigArg(dtype=str, nullable=False, enforced=True),
     "chunksize": _ConfigArg(dtype=int, nullable=False, enforced=True),
     "suppress_warnings": _ConfigArg(dtype=bool, nullable=False, default=False, loaded=True),
+    "dtype_backend": _ConfigArg(dtype=str, nullable=True),
     # Endpoints URLs
     "s3_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True, loaded=True),
     "athena_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True, loaded=True),
@@ -411,6 +413,15 @@ def suppress_warnings(self) -> bool:
     def suppress_warnings(self, value: bool) -> None:
         self._set_config_value(key="suppress_warnings", value=value)
 
+    @property
+    def dtype_backend(self) -> Literal["numpy_nullable", "pyarrow", None]:
+        """Property dtype_backend."""
+        return cast(Literal["numpy_nullable", "pyarrow", None], self["dtype_backend"])
+
+    @dtype_backend.setter
+    def dtype_backend(self, value: Literal["numpy_nullable", "pyarrow", None]) -> None:
+        self._set_config_value(key="dtype_backend", value=value)
+
     @property
     def s3_endpoint_url(self) -> Optional[str]:
         """Property s3_endpoint_url."""

awswrangler/_data_types.py

@@ -342,33 +342,35 @@ def athena2pyarrow(dtype: str) -> pa.DataType:  # pylint: disable=too-many-retur
     raise exceptions.UnsupportedType(f"Unsupported Athena type: {dtype}")
 
 
-def athena2pandas(dtype: str) -> str:  # pylint: disable=too-many-branches,too-many-return-statements
+def athena2pandas(
+    dtype: str, dtype_backend: Optional[str] = None
+) -> str:  # pylint: disable=too-many-branches,too-many-return-statements
     """Athena to Pandas data types conversion."""
     dtype = dtype.lower()
     if dtype == "tinyint":
-        return "Int8"
+        return "Int8" if dtype_backend != "pyarrow" else "int8[pyarrow]"
     if dtype == "smallint":
-        return "Int16"
+        return "Int16" if dtype_backend != "pyarrow" else "int16[pyarrow]"
     if dtype in ("int", "integer"):
-        return "Int32"
+        return "Int32" if dtype_backend != "pyarrow" else "int32[pyarrow]"
     if dtype == "bigint":
-        return "Int64"
+        return "Int64" if dtype_backend != "pyarrow" else "int64[pyarrow]"
     if dtype in ("float", "real"):
-        return "float32"
+        return "float32" if dtype_backend != "pyarrow" else "double[pyarrow]"
     if dtype == "double":
-        return "float64"
+        return "float64" if dtype_backend != "pyarrow" else "double[pyarrow]"
     if dtype == "boolean":
-        return "boolean"
+        return "boolean" if dtype_backend != "pyarrow" else "bool[pyarrow]"
     if (dtype == "string") or dtype.startswith("char") or dtype.startswith("varchar"):
-        return "string"
+        return "string" if dtype_backend != "pyarrow" else "string[pyarrow]"
     if dtype in ("timestamp", "timestamp with time zone"):
-        return "datetime64"
+        return "datetime64" if dtype_backend != "pyarrow" else "date64[pyarrow]"
     if dtype == "date":
-        return "date"
+        return "date" if dtype_backend != "pyarrow" else "date32[pyarrow]"
     if dtype.startswith("decimal"):
-        return "decimal"
+        return "decimal" if dtype_backend != "pyarrow" else "double[pyarrow]"
     if dtype in ("binary", "varbinary"):
-        return "bytes"
+        return "bytes" if dtype_backend != "pyarrow" else "binary[pyarrow]"
     if dtype in ("array", "row", "map"):
         return "object"
     raise exceptions.UnsupportedType(f"Unsupported Athena type: {dtype}")
@@ -465,6 +467,22 @@ def pyarrow2pandas_extension(  # pylint: disable=too-many-branches,too-many-retu
     return None
 
 
+def pyarrow2pyarrow_backed_pandas_extension(  # pylint: disable=too-many-branches,too-many-return-statements
+    dtype: pa.DataType,
+) -> Optional[pd.api.extensions.ExtensionDtype]:
+    """Pyarrow to Pandas PyArrow-backed data types conversion."""
+    return pd.ArrowDtype(dtype)
+
+
+def get_pyarrow2pandas_type_mapper(
+    dtype_backend: Optional[str] = None,
+) -> Callable[[pa.DataType], Optional[pd.api.extensions.ExtensionDtype]]:
+    if dtype_backend == "pyarrow":
+        return pyarrow2pyarrow_backed_pandas_extension
+
+    return pyarrow2pandas_extension
+
+
 @engine.dispatch_on_engine
 def pyarrow_types_from_pandas(  # pylint: disable=too-many-branches,too-many-statements
     df: pd.DataFrame, index: bool, ignore_cols: Optional[List[str]] = None, index_left: bool = False
@@ -550,14 +568,16 @@ def pyarrow_types_from_pandas(  # pylint: disable=too-many-branches,too-many-sta
     return columns_types
 
 
-def pyarrow2pandas_defaults(use_threads: Union[bool, int], kwargs: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+def pyarrow2pandas_defaults(
+    use_threads: Union[bool, int], kwargs: Optional[Dict[str, Any]] = None, dtype_backend: Optional[str] = None
+) -> Dict[str, Any]:
     """Return Pyarrow to Pandas default dictionary arguments."""
     default_kwargs = {
         "use_threads": use_threads,
         "split_blocks": True,
         "self_destruct": True,
         "ignore_metadata": False,
-        "types_mapper": pyarrow2pandas_extension,
+        "types_mapper": get_pyarrow2pandas_type_mapper(dtype_backend),
     }
     if kwargs:
         default_kwargs.update(kwargs)
@@ -685,7 +705,9 @@ def athena_types_from_pyarrow_schema(
     return columns_types, partitions_types
 
 
-def cast_pandas_with_athena_types(df: pd.DataFrame, dtype: Dict[str, str]) -> pd.DataFrame:
+def cast_pandas_with_athena_types(
+    df: pd.DataFrame, dtype: Dict[str, str], dtype_backend: Optional[str] = None
+) -> pd.DataFrame:
     """Cast columns in a Pandas DataFrame."""
     mutability_ensured: bool = False
     for col, athena_type in dtype.items():
@@ -695,7 +717,7 @@ def cast_pandas_with_athena_types(df: pd.DataFrame, dtype: Dict[str, str]) -> pd
             and (athena_type.startswith("struct") is False)
             and (athena_type.startswith("map") is False)
         ):
-            desired_type: str = athena2pandas(dtype=athena_type)
+            desired_type: str = athena2pandas(dtype=athena_type, dtype_backend=dtype_backend)
             current_type: str = _normalize_pandas_dtype_name(dtype=str(df[col].dtypes))
             if desired_type != current_type:  # Needs conversion
                 _logger.debug("current_type: %s -> desired_type: %s", current_type, desired_type)

awswrangler/_databases.py

@@ -7,6 +7,7 @@
 
 import boto3
 import pyarrow as pa
+from typing_extensions import Literal
 
 import awswrangler.pandas as pd
 from awswrangler import _data_types, _utils, exceptions, oracle, secretsmanager
@@ -153,6 +154,7 @@ def _records2df(
     safe: bool,
     dtype: Optional[Dict[str, pa.DataType]],
     timestamp_as_object: bool,
+    dtype_backend: Literal["numpy_nullable", "pyarrow"],
 ) -> pd.DataFrame:
     arrays: List[pa.Array] = []
     for col_values, col_name in zip(tuple(zip(*records)), cols_names):  # Transposing
@@ -183,7 +185,7 @@ def _records2df(
             self_destruct=True,
             integer_object_nulls=False,
             date_as_object=True,
-            types_mapper=_data_types.pyarrow2pandas_extension,
+            types_mapper=_data_types.get_pyarrow2pandas_type_mapper(dtype_backend=dtype_backend),
             safe=safe,
             timestamp_as_object=timestamp_as_object,
         )
@@ -207,6 +209,7 @@ def _iterate_results(
     safe: bool,
     dtype: Optional[Dict[str, pa.DataType]],
     timestamp_as_object: bool,
+    dtype_backend: Literal["numpy_nullable", "pyarrow"],
 ) -> Iterator[pd.DataFrame]:
     with con.cursor() as cursor:
         cursor.execute(*cursor_args)
@@ -230,6 +233,7 @@ def _iterate_results(
                 safe=safe,
                 dtype=dtype,
                 timestamp_as_object=timestamp_as_object,
+                dtype_backend=dtype_backend,
             )
 
 
@@ -240,6 +244,7 @@ def _fetch_all_results(
     dtype: Optional[Dict[str, pa.DataType]] = None,
     safe: bool = True,
     timestamp_as_object: bool = False,
+    dtype_backend: Literal["numpy_nullable", "pyarrow"] = "pyarrow",
 ) -> pd.DataFrame:
     with con.cursor() as cursor:
         cursor.execute(*cursor_args)
@@ -259,6 +264,7 @@ def _fetch_all_results(
             dtype=dtype,
             safe=safe,
             timestamp_as_object=timestamp_as_object,
+            dtype_backend=dtype_backend,
         )
 
 
@@ -272,6 +278,7 @@ def read_sql_query(
     dtype: Optional[Dict[str, pa.DataType]] = ...,
     safe: bool = ...,
     timestamp_as_object: bool = ...,
+    dtype_backend: Literal["numpy_nullable", "pyarrow"] = ...,
 ) -> pd.DataFrame:
     ...
 
@@ -287,6 +294,7 @@ def read_sql_query(
     dtype: Optional[Dict[str, pa.DataType]] = ...,
     safe: bool = ...,
     timestamp_as_object: bool = ...,
+    dtype_backend: Literal["numpy_nullable", "pyarrow"] = ...,
 ) -> Iterator[pd.DataFrame]:
     ...
 
@@ -302,6 +310,7 @@ def read_sql_query(
     dtype: Optional[Dict[str, pa.DataType]] = ...,
     safe: bool = ...,
     timestamp_as_object: bool = ...,
+    dtype_backend: Literal["numpy_nullable", "pyarrow"] = ...,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     ...
 
@@ -315,6 +324,7 @@ def read_sql_query(
     dtype: Optional[Dict[str, pa.DataType]] = None,
     safe: bool = True,
     timestamp_as_object: bool = False,
+    dtype_backend: Literal["numpy_nullable", "pyarrow"] = "numpy_nullable",
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     """Read SQL Query (generic)."""
     args = _convert_params(sql, params)
@@ -327,6 +337,7 @@ def read_sql_query(
                 dtype=dtype,
                 safe=safe,
                 timestamp_as_object=timestamp_as_object,
+                dtype_backend=dtype_backend,
             )
 
         return _iterate_results(
@@ -337,6 +348,7 @@ def read_sql_query(
             dtype=dtype,
             safe=safe,
             timestamp_as_object=timestamp_as_object,
+            dtype_backend=dtype_backend,
         )
     except Exception as ex:
         con.rollback()

awswrangler/_utils.py

@@ -161,6 +161,11 @@ def inner(*args: Any, **kwargs: Any) -> Any:
                 set([key for key, value in kwargs.items() if value is not None])
             )
 
+            # Allow kwargs that didn't modify the default value
+            passed_unsupported_kwargs = {
+                key for key in passed_unsupported_kwargs if kwargs[key] != signature.parameters[key].default
+            }
+
             if condition_fn() and len(passed_unsupported_kwargs) > 0:
                 raise exceptions.InvalidArgument(f"{message} `{', '.join(passed_unsupported_kwargs)}`.")