Repository: dask/dask
Pull request: Support `dtype_backend="pandas|pyarrow"` configuration
This PR updates the `use_nullable_dtypes=` keyword in `read_parquet` to accept `"pandas"` and `"pyarrow"` as valid inputs. The equivalent in `pandas` would be `use_nullable_dtypes=True` + the new `io.nullable_backend` `pandas` config option that's coming in `pandas=2.0`. I like `use_nullable_dtypes="pandas|pyarrow"` in Dask because

1. I find it a bit more ergonomic (my subjective opinion)
2. This functionality could be made available to Dask users quickly (actually, I'm not sure when `pandas=2.0` is scheduled for release -- I'm just guessing several `dask` releases will happen beforehand)
3. It let's us remain de-coupled from `pandas` config system

We might consider going with this over https://github.com/dask/dask/pull/9711 for having `read_parquet` support reading in pyarrow-backed dtypes

EDIT: The primary downside to this PR I see is that it's a deviation away from `pandas` API. The good news is that the logic is well isolated enough that it would be very easy to deprecate in the future, should we want to align on `io.nullable_backend` at some point. 

cc @rjzamora @mroeschke for visibility 

Closes https://github.com/dask/dask/issues/9631

Problem statement:
Read Parquet directly into string[pyarrow]
I would like to read parquet data directly into string pyarrow dtypes

So I'm trying this naively on a dataset:

```python
import dask.dataframe as dd

df = dd.read_parquet(
    "s3://nyc-tlc/trip data/fhvhv_tripdata_2022-06.parquet", 
    split_row_groups=True, 
    use_nullable_dtypes=True,
)
df.dtypes
```
```
hvfhs_license_num               object
dispatching_base_num            object
originating_base_num            object
request_datetime        datetime64[ns]
on_scene_datetime       datetime64[ns]
pickup_datetime         datetime64[ns]
dropoff_datetime        datetime64[ns]
PULocationID                     int64
DOLocationID                     int64
trip_miles                     float64
trip_time                        int64
base_passenger_fare            float64
tolls                          float64
bcf                            float64
sales_tax                      float64
congestion_surcharge           float64
airport_fee                    float64
tips                           float64
driver_pay                     float64
shared_request_flag             object
shared_match_flag               object
access_a_ride_flag              object
wav_request_flag                object
wav_match_flag                  object
dtype: object
```

This is especially important to me because in this case one row group is 10GB when stored naively.  The data is somewhat unreadable in its current state on modest machines.  I also suspect that I'm spending almost all of my time just creating and then destroying Python objects. 

Any thoughts @rjzamora @ian-r-rose ?

Hints from issue discussion:
Thanks @mrocklin. This sounds sensible. Something like `use_nullable_dtypes=True` + telling `pandas` to use `string[pyarrow]` as the default string type should do the trick.

Was the dataset originally written with extension string dtypes? Ideally I'd like to see us respect roundtripping 
> Was the dataset originally written with extension string dtypes

This one was not.  This one is maintained by nyc tlc and I don't think that they (yet!) use Dask for data prep.

> Something like use_nullable_dtypes=True + telling pandas to use string[pyarrow] as the default string type should do the trick

I would hope so, but I'm not sure that that's the case.  In this case Dask/Pandas/Arrow aren't sure that the object dtypes are strings, they might be some other serialized object.  In this case I'd like to explicitly tell the system "nope, this is definitely text data, please act accordingly".
> Was the dataset originally written with extension string dtypes? Ideally I'd like to see us respect roundtripping

@jrbourbeau We have data written with pyarrow strings, python strings and categorical (3 different sets), that we can use to test this if needed. 
>>Was the dataset originally written with extension string dtypes
>
>This one was not.

If I remember correctly (which probably don't), this may mean there is no way to do anything more efficient than using `astype("string[pyarrow]")` after the read. I was looking into this for https://github.com/dask/dask/issues/9476, but got distracted by other things.
I've got to believe that when Arrow reads this dataset they are able to see "yup, this is text" and read it in appropriately.  Spark also.  @jorisvandenbossche any suggestions?
> I've got to believe that when Arrow reads this dataset they are able to see "yup, this is text" and read it in appropriately.

I think that's right. See https://github.com/pandas-dev/pandas/issues/42664#issuecomment-885058289

@mrocklin are you able to set `pd.option_context("string_storage", "pyarrow")` on the workers before reading the data? And maybe doing it just on the client would be sufficient for checking the `meta`.

I'm not sure if that has a trip from arrow -> python objects -> arrow, but hopefully it's zero-copy from arrow memory to pandas.
Still no dice on the problem above:

```python
In [1]: import pandas as pd

In [2]: pd.option_context("string_storage", "pyarrow")
Out[2]: <pandas._config.config.option_context at 0x10295d690>

In [3]: import dask.dataframe as dd
   ...:
   ...: df = dd.read_parquet(
   ...:     "s3://nyc-tlc/trip data/fhvhv_tripdata_2022-06.parquet",
   ...:     split_row_groups=True,
   ...:     use_nullable_dtypes=True,
   ...: )
   ...: df.dtypes
   ...:
Out[3]:
hvfhs_license_num               object
dispatching_base_num            object
originating_base_num            object
request_datetime        datetime64[ns]
on_scene_datetime       datetime64[ns]
pickup_datetime         datetime64[ns]
dropoff_datetime        datetime64[ns]
PULocationID                     int64
DOLocationID                     int64
trip_miles                     float64
trip_time                        int64
base_passenger_fare            float64
tolls                          float64
bcf                            float64
sales_tax                      float64
congestion_surcharge           float64
airport_fee                    float64
tips                           float64
driver_pay                     float64
shared_request_flag             object
shared_match_flag               object
access_a_ride_flag              object
wav_request_flag                object
wav_match_flag                  object
dtype: object
```

I think that the problem here is less about the `string[pyarrow]` vs `string[python]` distinction as it is about the `string` vs `object` distinction. 
I didn't test in combination with dask, but just using pandas, the following seems to work:

```
import pyarrow as pa
import pyarrow.parquet as pq
pq.write_table(pa.table({"col": ['a', 'b', 'c']}), "test_string.parquet")

In [5]: pd.read_parquet("test_string.parquet").dtypes
Out[5]: 
col    object
dtype: object

In [6]: pd.read_parquet("test_string.parquet", use_nullable_dtypes=True).dtypes
Out[6]: 
col    string
dtype: object

In [7]: pd.options.mode.string_storage = "pyarrow"

In [8]: pd.read_parquet("test_string.parquet", use_nullable_dtypes=True).dtypes
Out[8]: 
col    string
dtype: object

In [9]: pd.read_parquet("test_string.parquet", use_nullable_dtypes=True).dtypes[0]
Out[9]: string[pyarrow]
```

So it might be some issue in how dask is propagating those keywords to pandas?

I was suspicious that this might be due to PyArrow writing metadata differently than however this data in the wild was written so I read the dataset with Pandas.  I received the following:

```python
In [1]: import pandas as pd

In [2]: pd.option_context("string_storage", "pyarrow")
Out[2]: <pandas._config.config.option_context at 0x106be1720>

In [3]: %time df = pd.read_parquet("s3://nyc-tlc/trip data/fhvhv_tripdata_2022-0
   ...: 1.parquet", use_nullable_dtypes=True)
CPU times: user 10.3 s, sys: 4.89 s, total: 15.2 s
Wall time: 59.5 s

In [4]: df.dtypes
Out[4]:
hvfhs_license_num               string
dispatching_base_num            string
originating_base_num            string
request_datetime        datetime64[ns]
on_scene_datetime       datetime64[ns]
pickup_datetime         datetime64[ns]
dropoff_datetime        datetime64[ns]
PULocationID                     Int64
DOLocationID                     Int64
trip_miles                     Float64
trip_time                        Int64
base_passenger_fare            Float64
tolls                          Float64
bcf                            Float64
sales_tax                      Float64
congestion_surcharge           Float64
airport_fee                    Float64
tips                           Float64
driver_pay                     Float64
shared_request_flag             string
shared_match_flag               string
access_a_ride_flag              string
wav_request_flag                string
wav_match_flag                  string
dtype: object

In [5]: df.hvfhs_license_num.head()
Out[5]:
0    HV0003
1    HV0003
2    HV0003
3    HV0003
4    HV0003
Name: hvfhs_license_num, dtype: string

In [6]: df.hvfhs_license_num.head().values
Out[6]:
<StringArray>
['HV0003', 'HV0003', 'HV0003', 'HV0003', 'HV0003']
Length: 5, dtype: string

In [9]: df.memory_usage(deep=True)
Out[9]:
Index                         128
hvfhs_license_num       929350233
dispatching_base_num    929350233
originating_base_num    839144900
request_datetime        118012728
on_scene_datetime       118012728
pickup_datetime         118012728
dropoff_datetime        118012728
PULocationID            132764319
DOLocationID            132764319
trip_miles              132764319
trip_time               132764319
base_passenger_fare     132764319
tolls                   132764319
bcf                     132764319
sales_tax               132764319
congestion_surcharge    132764319
airport_fee             132764319
tips                    132764319
driver_pay              132764319
shared_request_flag     855592278
shared_match_flag       855592278
access_a_ride_flag      855592278
wav_request_flag        855592278
wav_match_flag          855592278
dtype: int64
```

It looks like Pandas is properly getting the string dtype (🎉).

It's not clear to me that it's using pyarrow here for storage.  Those storage sizes seem large to me given the values.
I think you want to use `pd.set_option("mode.string_storage", "pyarrow")` instead of `pd.option_context("string_storage", "pyarrow")`

```python
In [1]: import pandas as pd

In [2]: pd.set_option("mode.string_storage", "pyarrow")

In [3]: df = pd.read_parquet("s3://nyc-tlc/trip data/fhvhv_tripdata_2022-01.parquet", use_nullable_dtypes=True)

In [4]: df.hvfhs_license_num.dtype
Out[4]: string[pyarrow]
```

Commit message: Use config option

Patch to review:
diff --git a/dask/dask-schema.yaml b/dask/dask-schema.yaml
index 8573c0e5b66..6be1a5b3bba 100644
--- a/dask/dask-schema.yaml
+++ b/dask/dask-schema.yaml
@@ -72,6 +72,14 @@ properties:
               task when reading a parquet dataset from a REMOTE file system.
               Specifying 0 will result in serial execution on the client.
 
+      nullable_backend:
+        enum:
+          - pandas
+          - pyarrow
+        description: |
+          The nullable dtype implementation to use. Must be either "pandas" or
+          "pyarrow". Default is "pandas".
+
   array:
     type: object
     properties:
diff --git a/dask/dask.yaml b/dask/dask.yaml
index 4b649c80a81..50d0aed9e24 100644
--- a/dask/dask.yaml
+++ b/dask/dask.yaml
@@ -12,6 +12,7 @@ dataframe:
   parquet:
     metadata-task-size-local: 512  # Number of files per local metadata-processing task
     metadata-task-size-remote: 16  # Number of files per remote metadata-processing task
+  nullable_backend: "pandas"  # Nullable dtype implementation to use
 
 array:
   backend: "numpy"  # Backend array library for input IO and data creation
diff --git a/dask/dataframe/io/parquet/arrow.py b/dask/dataframe/io/parquet/arrow.py
index 369de1707e1..e00cf050c8e 100644
--- a/dask/dataframe/io/parquet/arrow.py
+++ b/dask/dataframe/io/parquet/arrow.py
@@ -1580,19 +1580,31 @@ def _arrow_table_to_pandas(
         _kwargs.update({"use_threads": False, "ignore_metadata": False})
 
         if use_nullable_dtypes:
+            # Determine is `pandas` or `pyarrow`-backed dtypes should be used
+            if use_nullable_dtypes == "pandas":
+                default_types_mapper = PYARROW_NULLABLE_DTYPE_MAPPING.get
+            else:
+                # use_nullable_dtypes == "pyarrow"
+
+                def default_types_mapper(pyarrow_dtype):  # type: ignore
+                    # Special case pyarrow strings to use more feature complete dtype
+                    # See https://github.com/pandas-dev/pandas/issues/50074
+                    if pyarrow_dtype == pa.string():
+                        return pd.StringDtype("pyarrow")
+                    else:
+                        return pd.ArrowDtype(pyarrow_dtype)
+
             if "types_mapper" in _kwargs:
-                # User-provided entries take priority over PYARROW_NULLABLE_DTYPE_MAPPING
+                # User-provided entries take priority over default_types_mapper
                 types_mapper = _kwargs["types_mapper"]
 
                 def _types_mapper(pa_type):
-                    return types_mapper(pa_type) or PYARROW_NULLABLE_DTYPE_MAPPING.get(
-                        pa_type
-                    )
+                    return types_mapper(pa_type) or default_types_mapper(pa_type)
 
                 _kwargs["types_mapper"] = _types_mapper
 
             else:
-                _kwargs["types_mapper"] = PYARROW_NULLABLE_DTYPE_MAPPING.get
+                _kwargs["types_mapper"] = default_types_mapper
 
         return arrow_table.to_pandas(categories=categories, **_kwargs)
 
diff --git a/dask/dataframe/io/parquet/core.py b/dask/dataframe/io/parquet/core.py
index 8ed92b6a892..78da7e51dab 100644
--- a/dask/dataframe/io/parquet/core.py
+++ b/dask/dataframe/io/parquet/core.py
@@ -185,7 +185,7 @@ def read_parquet(
     index=None,
     storage_options=None,
     engine="auto",
-    use_nullable_dtypes=False,
+    use_nullable_dtypes: bool = False,
     calculate_divisions=None,
     ignore_metadata_file=False,
     metadata_task_size=None,
@@ -257,6 +257,22 @@ def read_parquet(
     engine : {'auto', 'pyarrow', 'fastparquet'}, default 'auto'
         Parquet library to use. Defaults to 'auto', which uses ``pyarrow`` if
         it is installed, and falls back to ``fastparquet`` otherwise.
+    use_nullable_dtypes : {False, True}
+        Whether to use extension dtypes for the resulting ``DataFrame``.
+        ``use_nullable_dtypes=True`` is only supported when ``engine="pyarrow"``.
+
+        .. note::
+
+            Use the ``dataframe.nullable_backend`` config option to select which
+            dtype implementation to use.
+
+            ``dataframe.nullable_backend="pandas"`` (the default) will use
+            pandas' ``numpy``-backed nullable dtypes (e.g. ``Int64``,
+            ``string[python]``, etc.) while ``dataframe.nullable_backend="pyarrow"``
+            will use ``pyarrow``-backed extension dtypes (e.g. ``int64[pyarrow]``,
+            ``string[pyarrow]``, etc.). ``dataframe.nullable_backend="pyarrow"``
+            requires ``pandas`` 1.5+.
+
     calculate_divisions : bool, default False
         Whether to use min/max statistics from the footer metadata (or global
         ``_metadata`` file) to calculate divisions for the output DataFrame
@@ -366,6 +382,9 @@ def read_parquet(
     pyarrow.parquet.ParquetDataset
     """
 
+    if use_nullable_dtypes:
+        use_nullable_dtypes = dask.config.get("dataframe.nullable_backend")
+
     # "Pre-deprecation" warning for `chunksize`
     if chunksize:
         warnings.warn(
@@ -552,7 +571,7 @@ def read_parquet(
     if "retries" not in annotations and not _is_local_fs(fs):
         ctx = dask.annotate(retries=5)
     else:
-        ctx = contextlib.nullcontext()
+        ctx = contextlib.nullcontext()  # type: ignore
 
     with ctx:
         # Construct the output collection with from_map
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
index c12f956f35f..2ab5ae3e65f 100644
--- a/dask/dataframe/io/tests/test_parquet.py
+++ b/dask/dataframe/io/tests/test_parquet.py
@@ -15,7 +15,12 @@
 import dask.dataframe as dd
 import dask.multiprocessing
 from dask.blockwise import Blockwise, optimize_blockwise
-from dask.dataframe._compat import PANDAS_GT_110, PANDAS_GT_121, PANDAS_GT_130
+from dask.dataframe._compat import (
+    PANDAS_GT_110,
+    PANDAS_GT_121,
+    PANDAS_GT_130,
+    PANDAS_GT_150,
+)
 from dask.dataframe.io.parquet.core import get_engine
 from dask.dataframe.io.parquet.utils import _parse_pandas_metadata
 from dask.dataframe.optimize import optimize_dataframe_getitem
@@ -618,17 +623,37 @@ def test_roundtrip_nullable_dtypes(tmp_path, write_engine, read_engine):
 
 
 @PYARROW_MARK
-def test_use_nullable_dtypes(tmp_path, engine):
+@pytest.mark.parametrize(
+    "nullable_backend",
+    [
+        "pandas",
+        pytest.param(
+            "pyarrow",
+            marks=pytest.mark.skipif(
+                not PANDAS_GT_150, reason="Requires pyarrow-backed nullable dtypes"
+            ),
+        ),
+    ],
+)
+def test_use_nullable_dtypes(tmp_path, engine, nullable_backend):
     """
     Test reading a parquet file without pandas metadata,
     but forcing use of nullable dtypes where appropriate
     """
+
+    if nullable_backend == "pandas":
+        dtype_extra = ""
+    else:
+        # nullable_backend == "pyarrow"
+        dtype_extra = "[pyarrow]"
     df = pd.DataFrame(
         {
-            "a": pd.Series([1, 2, pd.NA, 3, 4], dtype="Int64"),
-            "b": pd.Series([True, pd.NA, False, True, False], dtype="boolean"),
-            "c": pd.Series([0.1, 0.2, 0.3, pd.NA, 0.4], dtype="Float64"),
-            "d": pd.Series(["a", "b", "c", "d", pd.NA], dtype="string"),
+            "a": pd.Series([1, 2, pd.NA, 3, 4], dtype=f"Int64{dtype_extra}"),
+            "b": pd.Series(
+                [True, pd.NA, False, True, False], dtype=f"boolean{dtype_extra}"
+            ),
+            "c": pd.Series([0.1, 0.2, 0.3, pd.NA, 0.4], dtype=f"Float64{dtype_extra}"),
+            "d": pd.Series(["a", "b", "c", "d", pd.NA], dtype=f"string{dtype_extra}"),
         }
     )
     ddf = dd.from_pandas(df, npartitions=2)
@@ -644,21 +669,24 @@ def write_partition(df, i):
     partitions = ddf.to_delayed()
     dask.compute([write_partition(p, i) for i, p in enumerate(partitions)])
 
-    # Not supported by fastparquet
-    if engine == "fastparquet":
-        with pytest.raises(ValueError, match="`use_nullable_dtypes` is not supported"):
-            dd.read_parquet(tmp_path, engine=engine, use_nullable_dtypes=True)
+    with dask.config.set({"dataframe.nullable_backend": nullable_backend}):
+        # Not supported by fastparquet
+        if engine == "fastparquet":
+            with pytest.raises(
+                ValueError, match="`use_nullable_dtypes` is not supported"
+            ):
+                dd.read_parquet(tmp_path, engine=engine, use_nullable_dtypes=True)
 
-    # Works in pyarrow
-    else:
-        # Doesn't round-trip by default when we aren't using nullable dtypes
-        with pytest.raises(AssertionError):
-            ddf2 = dd.read_parquet(tmp_path, engine=engine)
-            assert_eq(df, ddf2)
-
-        # Round trip works when we use nullable dtypes
-        ddf2 = dd.read_parquet(tmp_path, engine=engine, use_nullable_dtypes=True)
-        assert_eq(df, ddf2, check_index=False)
+        # Works in pyarrow
+        else:
+            # Doesn't round-trip by default when we aren't using nullable dtypes
+            with pytest.raises(AssertionError):
+                ddf2 = dd.read_parquet(tmp_path, engine=engine)
+                assert_eq(df, ddf2)
+
+            # Round trip works when we use nullable dtypes
+            ddf2 = dd.read_parquet(tmp_path, engine=engine, use_nullable_dtypes=True)
+            assert_eq(df, ddf2, check_index=False)
 
 
 @pytest.mark.xfail(


Code review comments:
diff_hunk: @@ -12,6 +12,7 @@ dataframe:
   parquet:
     metadata-task-size-local: 512  # Number of files per local metadata-processing task
     metadata-task-size-remote: 16  # Number of files per remote metadata-processing task
+  nullable_backend: "pandas"  # Nullable dtype implementation to use
path: dask/dask.yaml
text: @user1:
How do you expect this option (and its default) to interact with the corresponding pandas 2.0 config option? When pandas-2 is released, should the default just correspond to whatever the pandas default is?

For example, it would be nice if we were able to use a test like this for pandas-2:

```python
with pd.option_context("io.nullable_backend", "pyarrow"):
    df = pd.read_parquet("tmp.parquet", engine="pyarrow", use_nullable_dtypes=True)
    ddf = dd.read_parquet("tmp.parquet", engine="pyarrow", use_nullable_dtypes=True)
assert_eq(dd, ddf)
```
Does client vs worker config options make this a challenge?


@author:
> Does client vs worker config options make this a challenge?

It's something we'll definitely need to account for. My guess is the most pleasant user experience will be if we pull the corresponding config value on the client and then embed it into the task graph (like we're doing in this PR). That way users won't need to worry about setting config options on the workers. Regardless, I suspect the implementation will be the same whether we pull the `pandas` or `dask` config option (see https://github.com/dask/dask/pull/9711 for an example). 

The downside to supporting `pandas` config options is that we wouldn't support _all_ the config options. We could explicitly document which ones we do support, and when, but still might be a source of confusion. 

Either way, I think this is a good question to ask. But I'm not too concerned because there is a smooth path in either direction. If we don't support the `pandas` option, then no changes are needed. If we do, then we can either update the default for the `dask` config value to pull in the current `pandas` option, or we deprecate the `dask` config value altogether.