Repository: dask/dask
Pull request: Support arrow dtypes in `Index` meta creation
- [x] Closes #10160
- [x] Tests added / passed
- [x] Passes `pre-commit run --all-files`

Sits on top of #10169

Not sure if there is a better way, this supports all dtypes that we support in pandas as of right now.
I let some of the cases fall through to the regular index check (pa.bool and pa.string), hence we have to check for ArrowDtype in all branches

Problem statement:
`ddf.set_index('col')` fails with pyarrow timestamp dtype column
**Describe the issue**:
The `set_index` operation fails on dask dataframes when the column dtype is a pyarrow timestamp.

**Minimal Complete Verifiable Example**:

```python
import pandas as pd
import pyarrow as pa
import dask.dataframe as dd

ts_series = pd.Series([pd.Timestamp('2022-01-01')], dtype='datetime64[ns]')
pa_ts_series = pd.Series(
    data=[pd.Timestamp('2022-01-01')],
    dtype=pd.ArrowDtype(pa.timestamp('ns')),
)

df = pd.DataFrame({'ts': ts_series, 'pa_ts': pa_ts_series})

df.set_index('ts')  # works
df.set_index('pa_ts')  # works

ddf = dd.from_pandas(df, npartitions=1)

ddf.set_index('ts')  # works
ddf.set_index('pa_ts')  # fails
```

**Environment**:

- Pandas version: 2.0.0
- Dask version: 2023.3.2
- Python version: 3.10

**Traceback**:
```
Traceback (most recent call last):
  File "/Users/bphillips/workspace/exos-chronos/zz.py", line 20, in <module>
    ddf.set_index('pa_ts')  # fails
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/dataframe/core.py", line 5199, in set_index
    return set_index(
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/dataframe/shuffle.py", line 236, in set_index
    return result.map_partitions(M.sort_index)
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/dataframe/core.py", line 998, in map_partitions
    return map_partitions(func, self, *args, **kwargs)
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/dataframe/core.py", line 6863, in map_partitions
    meta = _get_meta_map_partitions(args, dfs, func, kwargs, meta, parent_meta)
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/dataframe/core.py", line 6974, in _get_meta_map_partitions
    meta = _emulate(func, *args, udf=True, **kwargs)
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/dataframe/core.py", line 6792, in _emulate
    with raise_on_meta_error(funcname(func), udf=udf), check_numeric_only_deprecation():
  File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/contextlib.py", line 153, in __exit__
    self.gen.throw(typ, value, traceback)
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/dataframe/utils.py", line 214, in raise_on_meta_error
    raise ValueError(msg) from e
ValueError: Metadata inference failed in `sort_index`.

You have supplied a custom function and Dask is unable to 
determine the type of output that that function returns. 

To resolve this please provide a meta= keyword.
The docstring of the Dask function you ran should have more information.

Original error is below:
------------------------
ArrowTypeError("object of type <class 'str'> cannot be converted to int")

Traceback:
---------
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/dataframe/utils.py", line 193, in raise_on_meta_error
    yield
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/dataframe/core.py", line 6793, in _emulate
    return func(*_extract_meta(args, True), **_extract_meta(kwargs, True))
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/dataframe/core.py", line 6773, in _extract_meta
    return tuple(_extract_meta(_x, nonempty) for _x in x)
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/dataframe/core.py", line 6773, in <genexpr>
    return tuple(_extract_meta(_x, nonempty) for _x in x)
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/dataframe/core.py", line 6769, in _extract_meta
    return x._meta_nonempty if nonempty else x._meta
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/dataframe/core.py", line 558, in _meta_nonempty
    return meta_nonempty(self._meta)
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/utils.py", line 642, in __call__
    return meth(arg, *args, **kwargs)
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/dataframe/backends.py", line 314, in meta_nonempty_dataframe
    idx = meta_nonempty(x.index)
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/utils.py", line 642, in __call__
    return meth(arg, *args, **kwargs)
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/dask/dataframe/backends.py", line 343, in _nonempty_index
    return pd.Index(["a", "b"], name=idx.name, dtype=idx.dtype)
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/pandas/core/indexes/base.py", line 551, in __new__
    arr = sanitize_array(data, None, dtype=dtype, copy=copy)
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/pandas/core/construction.py", line 559, in sanitize_array
    subarr = cls._from_sequence(data, dtype=dtype, copy=copy)
  File "/Users/bphillips/workspace/exos-chronos/.venv/lib/python3.10/site-packages/pandas/core/arrays/arrow/array.py", line 253, in _from_sequence
    scalars = pa.array(scalars, type=pa_dtype, from_pandas=True)
  File "pyarrow/array.pxi", line 320, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 39, in pyarrow.lib._sequence_to_array
  File "pyarrow/error.pxi", line 144, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 123, in pyarrow.lib.check_status
```

Hints from issue discussion:
A related example, although happy to move into another issue if appropriate. Seems to be tied to meta generation when the index is a pyarrow timestamp dtype

```python
index = pd.Series(
    data=[pd.Timestamp('2022-01-01')],
    dtype=pd.ArrowDtype(pa.timestamp('ns')),
)
str_series = pd.Series(['a'], index=index, dtype=pd.ArrowDtype(pa.string()))

df = pd.DataFrame(data={'str': str_series}, index=index)
ddf = dd.from_pandas(df, npartitions=1)

df['str'] = df['str'].astype(str)  # works
ddf['str'] = ddf['str'].astype(str)  # fails
```

Thanks for report @bphillips-exos. I'm able to reproduce.

cc @phofl for when you're back online 

EDIT: 

> Seems to be tied to meta generation when the index is a pyarrow timestamp dtype

Yeah, that's correct. We need to generalize Index `meta` creation to accommodate those dtypes

Commit message: Fix compatibility

Patch to review:
diff --git a/dask/dataframe/_compat.py b/dask/dataframe/_compat.py
index 03ab63d5b46..dd014b7ca59 100644
--- a/dask/dataframe/_compat.py
+++ b/dask/dataframe/_compat.py
@@ -181,3 +181,9 @@ def is_any_real_numeric_dtype(arr_or_dtype) -> bool:
             and not is_complex_dtype(arr_or_dtype)
             and not is_bool_dtype(arr_or_dtype)
         )
+
+
+def is_arrow_dtype(dtype) -> bool:
+    if not PANDAS_GT_150:
+        return False
+    return isinstance(dtype, pd.ArrowDtype)
diff --git a/dask/dataframe/backends.py b/dask/dataframe/backends.py
index 74b182ba1b9..4ac17d918f8 100644
--- a/dask/dataframe/backends.py
+++ b/dask/dataframe/backends.py
@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import warnings
+from datetime import date, time
+from decimal import Decimal
 from typing import Iterable
 
 import numpy as np
@@ -19,7 +21,7 @@
 from dask.array.dispatch import percentile_lookup
 from dask.array.percentile import _percentile
 from dask.backends import CreationDispatch, DaskBackendEntrypoint
-from dask.dataframe._compat import is_any_real_numeric_dtype
+from dask.dataframe._compat import is_any_real_numeric_dtype, is_arrow_dtype
 from dask.dataframe.core import DataFrame, Index, Scalar, Series, _Frame
 from dask.dataframe.dispatch import (
     categorical_dtype_dispatch,
@@ -50,6 +52,11 @@
 from dask.sizeof import SimpleSizeof, sizeof
 from dask.utils import is_arraylike, is_series_like, typename
 
+try:
+    import pyarrow as pa
+except ImportError:
+    pa = None
+
 
 class DataFrameBackendEntrypoint(DaskBackendEntrypoint):
     """Dask-DataFrame version of ``DaskBackendEntrypoint``
@@ -333,6 +340,31 @@ def _nonempty_index(idx):
         return pd.RangeIndex(2, name=idx.name, dtype=idx.dtype)
     elif is_any_real_numeric_dtype(idx):
         return typ([1, 2], name=idx.name, dtype=idx.dtype)
+    elif is_arrow_dtype(idx.dtype) and pa.types.is_timestamp(idx.dtype.pyarrow_dtype):
+        return pd.Index(
+            [pd.Timestamp("1970-01-01"), pd.Timestamp("1970-01-02")],
+            dtype=idx.dtype,
+            name=idx.name,
+        )
+    elif is_arrow_dtype(idx.dtype) and pa.types.is_date(idx.dtype.pyarrow_dtype):
+        return pd.Index(
+            [date(1970, 1, 1), date(1970, 1, 2)], dtype=idx.dtype, name=idx.name
+        )
+    elif is_arrow_dtype(idx.dtype) and (
+        pa.types.is_binary(idx.dtype.pyarrow_dtype)
+        or pa.types.is_large_binary(idx.dtype.pyarrow_dtype)
+    ):
+        return pd.Index([b"a", b"b"], dtype=idx.dtype, name=idx.name)
+    elif is_arrow_dtype(idx.dtype) and pa.types.is_decimal(idx.dtype.pyarrow_dtype):
+        return pd.Index([Decimal("1"), Decimal("0.0")], dtype=idx.dtype, name=idx.name)
+    elif is_arrow_dtype(idx.dtype) and pa.types.is_duration(idx.dtype.pyarrow_dtype):
+        return pd.Index(
+            [pd.Timedelta("1 day"), pd.Timedelta("2 days")],
+            dtype=idx.dtype,
+            name=idx.name,
+        )
+    elif is_arrow_dtype(idx.dtype) and pa.types.is_time(idx.dtype.pyarrow_dtype):
+        return pd.Index([time(12, 0), time(0, 12)], dtype=idx.dtype, name=idx.name)
     elif typ is pd.Index:
         if idx.dtype == bool:
             # pd 1.5 introduce bool dtypes and respect non-uniqueness
diff --git a/dask/dataframe/partitionquantiles.py b/dask/dataframe/partitionquantiles.py
index 128493b5dd6..94e4fe1c32e 100644
--- a/dask/dataframe/partitionquantiles.py
+++ b/dask/dataframe/partitionquantiles.py
@@ -383,6 +383,9 @@ def process_val_weights(vals_and_weights, npartitions, dtype_info):
     elif "datetime64" in str(dtype):
         rv = pd.DatetimeIndex(rv, dtype=dtype)
     elif rv.dtype != dtype:
+        if is_integer_dtype(dtype):
+            # pandas EA raises instead of truncating
+            rv = np.floor(rv)
         rv = pd.array(rv, dtype=dtype)
     return rv
 
diff --git a/dask/dataframe/tests/test_pyarrow.py b/dask/dataframe/tests/test_pyarrow.py
index f3dd4e629c4..7b41ca3218c 100644
--- a/dask/dataframe/tests/test_pyarrow.py
+++ b/dask/dataframe/tests/test_pyarrow.py
@@ -1,9 +1,14 @@
+from datetime import date, time
+
 import numpy as np
 import pandas as pd
 import pytest
+from _decimal import Decimal
 from pandas.tests.extension.decimal.array import DecimalDtype
 
-from dask.dataframe._compat import PANDAS_GT_140, PANDAS_GT_150
+import dask.dataframe as dd
+from dask.dataframe import assert_eq
+from dask.dataframe._compat import PANDAS_GT_140, PANDAS_GT_150, PANDAS_GT_200
 from dask.dataframe._pyarrow import (
     is_object_string_dataframe,
     is_object_string_dtype,
@@ -210,3 +215,34 @@ def test_is_object_string_series(series, expected):
 )
 def tests_is_object_string_dataframe(series, expected):
     assert is_object_string_dataframe(series) is expected
+
+
+@pytest.mark.skipif(not PANDAS_GT_200, reason="dtype support not good before 2.0")
+@pytest.mark.parametrize(
+    "data, arrow_dtype",
+    [
+        (["a", "b"], pa.string()),
+        ([b"a", b"b"], pa.binary()),
+        # Should probably fix upstream, https://github.com/pandas-dev/pandas/issues/52590
+        # (["a", "b"], pa.large_string()),
+        # ([b"a", b"b"], pa.large_binary()),
+        ([1, 2], pa.int64()),
+        ([1, 2], pa.float64()),
+        ([1, 2], pa.uint64()),
+        ([date(2022, 1, 1), date(1999, 12, 31)], pa.date32()),
+        (
+            [pd.Timestamp("2022-01-01"), pd.Timestamp("2023-01-02")],
+            pa.timestamp(unit="ns"),
+        ),
+        ([Decimal("5"), Decimal("6.24")], pa.decimal128(10, 2)),
+        ([pd.Timedelta("1 day"), pd.Timedelta("20 days")], pa.duration("ns")),
+        ([time(12, 0), time(0, 12)], pa.time64("ns")),
+    ],
+)
+def test_set_index_pyarrow_dtype(data, arrow_dtype):
+    dtype = pd.ArrowDtype(arrow_dtype)
+    pdf = pd.DataFrame({"a": 1, "arrow_col": pd.Series(data, dtype=dtype)})
+    ddf = dd.from_pandas(pdf, npartitions=2)
+    pdf_result = pdf.set_index("arrow_col")
+    ddf_result = ddf.set_index("arrow_col")
+    assert_eq(ddf_result, pdf_result)
diff --git a/dask/dataframe/tests/test_shuffle.py b/dask/dataframe/tests/test_shuffle.py
index 009356b1ba4..0bf28d35c3e 100644
--- a/dask/dataframe/tests/test_shuffle.py
+++ b/dask/dataframe/tests/test_shuffle.py
@@ -1082,6 +1082,15 @@ def test_set_index_timestamp():
     assert_eq(df2, ddf.set_index("A"), check_freq=False)
 
 
+@pytest.mark.skipif(not PANDAS_GT_140, reason="EA Indexes not supported before")
+def test_set_index_ea_dtype():
+    pdf = pd.DataFrame({"a": 1, "b": pd.Series([1, 2], dtype="Int64")})
+    ddf = dd.from_pandas(pdf, npartitions=2)
+    pdf_result = pdf.set_index("b")
+    ddf_result = ddf.set_index("b")
+    assert_eq(ddf_result, pdf_result)
+
+
 @pytest.mark.parametrize("compression", [None, "ZLib"])
 def test_disk_shuffle_with_compression_option(compression):
     # test if dataframe shuffle works both with and without compression


Code review comments:
diff_hunk: @@ -1,9 +1,14 @@
+from datetime import date, time
+
 import numpy as np
 import pandas as pd
 import pytest
+from _decimal import Decimal
path: dask/dataframe/tests/test_pyarrow.py
text: @user1:
Why `_decimal` and not `decimal`? 

@author:
This must have gotten In somehow (auto-import maybe?), changed

diff_hunk: @@ -333,6 +340,31 @@ def _nonempty_index(idx):
         return pd.RangeIndex(2, name=idx.name, dtype=idx.dtype)
     elif is_any_real_numeric_dtype(idx):
         return typ([1, 2], name=idx.name, dtype=idx.dtype)
+    elif is_arrow_dtype(idx.dtype) and pa.types.is_timestamp(idx.dtype.pyarrow_dtype):
path: dask/dataframe/backends.py
text: @user1:
I'm a little surprised we need to have the `is_arrow_dtype` check here. Naively I was expecting something like 

```python
        return pd.Index(
            [pd.Timestamp("1970-01-01"), pd.Timestamp("1970-01-02")],
            dtype=idx.dtype,
            name=idx.name,
        )
```

To work for both `pyarrow`-backed and non-`pyarrow`-backed dtypes. With `pandas` handling the conversion internally. 

Why not use `pd.api.types.is_*_dtype(...)` checks instead? Maybe there isn't a corresponding check in `pd.api.types` for these dtypes? 

@user1:
Also, previously I added something similar to our non-empty meta Series creation that resulted in meta creation working, but with all `pd.NA` values

https://github.com/dask/dask/blob/35554159781b2f5b7fecbd967ea76e90f97cb63c/dask/dataframe/_dtypes.py#L24-L26

It might be worth revisiting / reusing some of the logic here for better series creation. Not saying that needs to happen in this PR. Just mentioned for context and as a possible follow-up 



@author:
I will look into the existing logic tomorrow.

Non-arrow DatetimeIndexes have a freq attribute which would require special casing too. I like the current version a bit better because it's explicit, but can change if you prefer.

Most of the non numeric is_*_dtype aren't arrow aware yet, ``is_datetime64_dtype`` for example isn't. They have quite different semantics, which makes supporting arrow there non-trivial. We ran into a couple of problems over the last few weeks there, not sure if/when we support arrow there. The whole date time-conversion logic in pandas makes this non-trivial.

@author:
I refactored the existing logic so that we pull proper values for Series and Index.

diff_hunk: @@ -210,3 +215,34 @@ def test_is_object_string_series(series, expected):
 )
 def tests_is_object_string_dataframe(series, expected):
     assert is_object_string_dataframe(series) is expected
+
+
+@pytest.mark.skipif(not PANDAS_GT_200, reason="dtype support not good before 2.0")
+@pytest.mark.parametrize(
+    "data, arrow_dtype",
+    [
+        (["a", "b"], pa.string()),
+        ([b"a", b"b"], pa.binary()),
+        # Should probably fix upstream, https://github.com/pandas-dev/pandas/issues/52590
+        # (["a", "b"], pa.large_string()),
+        # ([b"a", b"b"], pa.large_binary()),
+        ([1, 2], pa.int64()),
+        ([1, 2], pa.float64()),
+        ([1, 2], pa.uint64()),
+        ([date(2022, 1, 1), date(1999, 12, 31)], pa.date32()),
+        (
+            [pd.Timestamp("2022-01-01"), pd.Timestamp("2023-01-02")],
+            pa.timestamp(unit="ns"),
+        ),
+        ([Decimal("5"), Decimal("6.24")], pa.decimal128(10, 2)),
+        ([pd.Timedelta("1 day"), pd.Timedelta("20 days")], pa.duration("ns")),
+        ([time(12, 0), time(0, 12)], pa.time64("ns")),
+    ],
+)
+def test_set_index_pyarrow_dtype(data, arrow_dtype):
path: dask/dataframe/tests/test_pyarrow.py
text: @user1:
Currently this test module is used for testing things related to our automatic `pyarrow` string conversion (`dataframe.convert-string=True`). My sense is this test should live over in `dask/dataframe/tests/test_shuffle.py` alongside all the other `set_index`-related tests 

@author:
Yeah wasn't sure about this. Will move