Repository: dask/dask

Issue:
`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
```


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 

Resolving patch:
diff --git a/dask/dataframe/_dtypes.py b/dask/dataframe/_dtypes.py
--- a/dask/dataframe/_dtypes.py
+++ b/dask/dataframe/_dtypes.py
@@ -1,3 +1,6 @@
+from datetime import date, time
+from decimal import Decimal
+
 import pandas as pd
 
 from dask.dataframe._compat import PANDAS_GT_150
@@ -23,7 +26,35 @@ def _(dtype):
 
     @make_array_nonempty.register(pd.ArrowDtype)
     def _make_array_nonempty_pyarrow_dtype(dtype):
-        return dtype.empty(2)
+        import pyarrow as pa
+
+        if pa.types.is_integer(dtype.pyarrow_dtype):
+            data = [1, 2]
+        elif pa.types.is_floating(dtype.pyarrow_dtype):
+            data = [1.5, 2.5]
+        elif pa.types.is_boolean(dtype.pyarrow_dtype):
+            data = [True, False]
+        elif pa.types.is_string(dtype.pyarrow_dtype) or pa.types.is_large_string(
+            dtype.pyarrow_dtype
+        ):
+            data = ["a", "b"]
+        elif pa.types.is_timestamp(dtype.pyarrow_dtype):
+            data = [pd.Timestamp("1970-01-01"), pd.Timestamp("1970-01-02")]
+        elif pa.types.is_date(dtype.pyarrow_dtype):
+            data = [date(1970, 1, 1), date(1970, 1, 2)]
+        elif pa.types.is_binary(dtype.pyarrow_dtype) or pa.types.is_large_binary(
+            dtype.pyarrow_dtype
+        ):
+            data = [b"a", b"b"]
+        elif pa.types.is_decimal(dtype.pyarrow_dtype):
+            data = [Decimal("1"), Decimal("0.0")]
+        elif pa.types.is_duration(dtype.pyarrow_dtype):
+            data = [pd.Timedelta("1 day"), pd.Timedelta("2 days")]
+        elif pa.types.is_time(dtype.pyarrow_dtype):
+            data = [time(12, 0), time(0, 12)]
+        else:
+            data = dtype.empty(2)
+        return pd.array(data, dtype=dtype)
 
 
 @make_scalar.register(str)
diff --git a/dask/dataframe/backends.py b/dask/dataframe/backends.py
--- a/dask/dataframe/backends.py
+++ b/dask/dataframe/backends.py
@@ -325,14 +325,6 @@ 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 typ is pd.Index:
-        if idx.dtype == bool:
-            # pd 1.5 introduce bool dtypes and respect non-uniqueness
-            return pd.Index([True, False], name=idx.name)
-        else:
-            # for pd 1.5 in the case of bool index this would be cast as [True, True]
-            # breaking uniqueness
-            return pd.Index(["a", "b"], name=idx.name, dtype=idx.dtype)
     elif typ is pd.DatetimeIndex:
         start = "1970-01-01"
         # Need a non-monotonic decreasing index to avoid issues with
@@ -380,6 +372,18 @@ def _nonempty_index(idx):
             return pd.MultiIndex(levels=levels, codes=codes, names=idx.names)
         except TypeError:  # older pandas versions
             return pd.MultiIndex(levels=levels, labels=codes, names=idx.names)
+    elif typ is pd.Index:
+        if type(idx.dtype) in make_array_nonempty._lookup:
+            return pd.Index(
+                make_array_nonempty(idx.dtype), dtype=idx.dtype, name=idx.name
+            )
+        elif idx.dtype == bool:
+            # pd 1.5 introduce bool dtypes and respect non-uniqueness
+            return pd.Index([True, False], name=idx.name)
+        else:
+            # for pd 1.5 in the case of bool index this would be cast as [True, True]
+            # breaking uniqueness
+            return pd.Index(["a", "b"], name=idx.name, dtype=idx.dtype)
 
     raise TypeError(f"Don't know how to handle index of type {typename(type(idx))}")
 


Test patch:
diff --git a/dask/dataframe/tests/test_shuffle.py b/dask/dataframe/tests/test_shuffle.py
--- a/dask/dataframe/tests/test_shuffle.py
+++ b/dask/dataframe/tests/test_shuffle.py
@@ -7,6 +7,8 @@
 import tempfile
 from concurrent.futures import ProcessPoolExecutor
 from copy import copy
+from datetime import date, time
+from decimal import Decimal
 from functools import partial
 from unittest import mock
 
@@ -20,6 +22,7 @@
 from dask.dataframe._compat import (
     PANDAS_GT_140,
     PANDAS_GT_150,
+    PANDAS_GT_200,
     assert_categorical_equal,
     tm,
 )
@@ -35,6 +38,11 @@
 from dask.dataframe.utils import assert_eq, make_meta
 from dask.optimization import cull
 
+try:
+    import pyarrow as pa
+except ImportError:
+    pa = None
+
 dsk = {
     ("x", 0): pd.DataFrame({"a": [1, 2, 3], "b": [1, 4, 7]}, index=[0, 1, 3]),
     ("x", 1): pd.DataFrame({"a": [4, 5, 6], "b": [2, 5, 8]}, index=[5, 6, 8]),
@@ -1548,3 +1556,36 @@ def test_calculate_divisions(pdf, expected):
     assert mins == expected[1]
     assert maxes == expected[2]
     assert presorted == expected[3]
+
+
+@pytest.mark.skipif(pa is None, reason="Need pyarrow")
+@pytest.mark.skipif(not PANDAS_GT_200, reason="dtype support not good before 2.0")
+@pytest.mark.parametrize(
+    "data, dtype",
+    [
+        (["a", "b"], "string[pyarrow]"),
+        ([b"a", b"b"], "binary[pyarrow]"),
+        # 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], "int64[pyarrow]"),
+        ([1, 2], "float64[pyarrow]"),
+        ([1, 2], "uint64[pyarrow]"),
+        ([date(2022, 1, 1), date(1999, 12, 31)], "date32[pyarrow]"),
+        (
+            [pd.Timestamp("2022-01-01"), pd.Timestamp("2023-01-02")],
+            "timestamp[ns][pyarrow]",
+        ),
+        ([Decimal("5"), Decimal("6.24")], "decimal128"),
+        ([pd.Timedelta("1 day"), pd.Timedelta("20 days")], "duration[ns][pyarrow]"),
+        ([time(12, 0), time(0, 12)], "time64[ns][pyarrow]"),
+    ],
+)
+def test_set_index_pyarrow_dtype(data, dtype):
+    if dtype == "decimal128":
+        dtype = pd.ArrowDtype(pa.decimal128(10, 2))
+    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)