Repository: dask/dask

Issue:
`set_index` does not support passing in a dask Index
**What happened**:

Passing a dask Index to `DataFrame.set_index` fails while trying to generate `meta` while inferring the new output divisions.

**What you expected to happen**:

New output divisions should not even need to be calculated. The dask Index object being passed in already has divisions; we should just use those. Furthermore, since the dask Index has known divisions, we should also infer that `sorted=True`.

Basically, `set_index` with another Index should be very cheap—just a `map_partitions` call internally—as long as the other Index has known divisions and the same number of partitions as the current one.

**Minimal Complete Verifiable Example**:

```python
import dask.dataframe as ddf
df = ddf.demo.make_timeseries()
new_index = df.index.dt.floor("D")
df.set_index(new_index)
# TypeError: Expected meta to specify type Series, got type pandas.core.indexes.datetimes.DatetimeIndex
```
<details><summary>Full traceback</summary>

```python
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-19-2be06237e532> in <module>
----> 1 df.set_index(df.index.dt.floor("D"))

~/dev/dask/dask/dataframe/core.py in set_index(***failed resolving arguments***)
   4359             from .shuffle import set_index
   4360 
-> 4361             return set_index(
   4362                 self,
   4363                 other,

~/dev/dask/dask/dataframe/shuffle.py in set_index(df, index, npartitions, shuffle, compute, drop, upsample, divisions, partition_size, **kwargs)
    189 
    190     if divisions is None:
--> 191         divisions, mins, maxes = _calculate_divisions(
    192             df, index2, repartition, npartitions, upsample, partition_size
    193         )

~/dev/dask/dask/dataframe/shuffle.py in _calculate_divisions(df, partition_col, repartition, npartitions, upsample, partition_size)
     35     """
     36     sizes = df.map_partitions(sizeof) if repartition else []
---> 37     divisions = partition_col._repartition_quantiles(npartitions, upsample=upsample)
     38     mins = partition_col.map_partitions(M.min)
     39     maxes = partition_col.map_partitions(M.max)

~/dev/dask/dask/dataframe/core.py in _repartition_quantiles(self, npartitions, upsample)
   3356         from .partitionquantiles import partition_quantiles
   3357 
-> 3358         return partition_quantiles(self, npartitions, upsample=upsample)
   3359 
   3360     def __getitem__(self, key):

~/dev/dask/dask/dataframe/partitionquantiles.py in partition_quantiles(df, npartitions, upsample, random_state)
    493     dsk = merge(df.dask, dtype_dsk, val_dsk, merge_dsk, last_dsk)
    494     new_divisions = [0.0, 1.0]
--> 495     return return_type(dsk, name3, df._meta, new_divisions)

~/dev/dask/dask/dataframe/core.py in __init__(self, dsk, name, meta, divisions)
    318         meta = make_meta(meta)
    319         if not self._is_partition_type(meta):
--> 320             raise TypeError(
    321                 f"Expected meta to specify type {type(self).__name__}, got type "
    322                 f"{typename(type(meta))}"

TypeError: Expected meta to specify type Series, got type pandas.core.indexes.datetimes.DatetimeIndex
```

</details>

**Environment**:

- Dask version: 5b38dda3ff4c2432baa6af8b52ba1572a539548e
- Python version: 3.8.8
- Operating System: macOS
- Install method (conda, pip, source): source


Discussion:
Yeah that makes sense to me. Do you think we need to do any validation?
I think we'd just validate that the new index has same number of partitions as old index. If it has (valid) known divisions, then we can go down the `sorted=True, divisions=new_index.divisions` fastpath. If not, we convert it to a Series and go down the normal unknown divisions path?
sounds like a good plan to me. 
FYI -- I'm going to take look at this over the coming days


Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -4594,8 +4594,22 @@ def set_index(
         pre_sorted = sorted
         del sorted
 
+        if isinstance(other, Series) and other._name == self.index._name:
+            # Index is equal to current index, no-op
+            return self
+
         if divisions is not None:
             check_divisions(divisions)
+        elif (
+            isinstance(other, Index)
+            and other.known_divisions
+            and other.npartitions == self.npartitions
+        ):
+            # If the index has the same number of partitions and known
+            # divisions, then we can treat it as pre-sorted with known
+            # divisions
+            pre_sorted = True
+            divisions = other.divisions
 
         if pre_sorted:
             from .shuffle import set_sorted_index
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -1017,12 +1017,20 @@ def compute_and_set_divisions(df, **kwargs):
 
 
 def set_sorted_index(df, index, drop=True, divisions=None, **kwargs):
-    if not isinstance(index, Series):
-        meta = df._meta.set_index(index, drop=drop)
-    else:
+    if isinstance(index, _Frame):
         meta = df._meta.set_index(index._meta, drop=drop)
+    else:
+        meta = df._meta.set_index(index, drop=drop)
 
-    result = map_partitions(M.set_index, df, index, drop=drop, meta=meta)
+    result = map_partitions(
+        M.set_index,
+        df,
+        index,
+        drop=drop,
+        meta=meta,
+        align_dataframes=False,
+        transform_divisions=False,
+    )
 
     if not divisions:
         return compute_and_set_divisions(result, **kwargs)


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
@@ -188,17 +188,12 @@ def test_set_index_general(npartitions, shuffle_method):
     ddf = dd.from_pandas(df, npartitions=npartitions)
 
     assert_eq(df.set_index("x"), ddf.set_index("x", shuffle=shuffle_method))
-
     assert_eq(df.set_index("y"), ddf.set_index("y", shuffle=shuffle_method))
-
     assert_eq(df.set_index(df.x), ddf.set_index(ddf.x, shuffle=shuffle_method))
-
     assert_eq(
         df.set_index(df.x + df.y), ddf.set_index(ddf.x + ddf.y, shuffle=shuffle_method)
     )
-
     assert_eq(df.set_index(df.x + 1), ddf.set_index(ddf.x + 1, shuffle=shuffle_method))
-
     assert_eq(df.set_index(df.index), ddf.set_index(ddf.index, shuffle=shuffle_method))
 
 
@@ -1230,6 +1225,63 @@ def test_set_index_nan_partition():
     assert_eq(a, a)
 
 
+def test_set_index_with_dask_dt_index():
+    values = {
+        "x": [1, 2, 3, 4] * 3,
+        "y": [10, 20, 30] * 4,
+        "name": ["Alice", "Bob"] * 6,
+    }
+    date_index = pd.date_range(
+        start="2022-02-22", freq="16h", periods=12
+    ) - pd.Timedelta(seconds=30)
+    df = pd.DataFrame(values, index=date_index)
+    ddf = dd.from_pandas(df, npartitions=3)
+
+    # specify a different date index entirely
+    day_index = ddf.index.dt.floor("D")
+    day_df = ddf.set_index(day_index)
+    expected = dd.from_pandas(
+        pd.DataFrame(values, index=date_index.floor("D")), npartitions=3
+    )
+    assert_eq(day_df, expected)
+
+    # specify an index with shifted dates
+    one_day = pd.Timedelta(days=1)
+    next_day_df = ddf.set_index(ddf.index + one_day)
+    expected = dd.from_pandas(
+        pd.DataFrame(values, index=date_index + one_day), npartitions=3
+    )
+    assert_eq(next_day_df, expected)
+
+    # try a different index type
+    no_dates = dd.from_pandas(pd.DataFrame(values), npartitions=3)
+    range_df = ddf.set_index(no_dates.index)
+    expected = dd.from_pandas(pd.DataFrame(values), npartitions=3)
+    assert_eq(range_df, expected)
+
+
+def test_set_index_with_series_uses_fastpath():
+    dates = pd.date_range(start="2022-02-22", freq="16h", periods=12) - pd.Timedelta(
+        seconds=30
+    )
+    one_day = pd.Timedelta(days=1)
+    df = pd.DataFrame(
+        {
+            "x": [1, 2, 3, 4] * 3,
+            "y": [10, 20, 30] * 4,
+            "name": ["Alice", "Bob"] * 6,
+            "d1": dates + one_day,
+            "d2": dates + one_day * 5,
+        },
+        index=dates,
+    )
+    ddf = dd.from_pandas(df, npartitions=3)
+
+    res = ddf.set_index(ddf.d2 + one_day)
+    expected = df.set_index(df.d2 + one_day)
+    assert_eq(res, expected)
+
+
 @pytest.mark.parametrize("ascending", [True, False])
 @pytest.mark.parametrize("by", ["a", "b"])
 @pytest.mark.parametrize("nelem", [10, 500])