Repository: dask/dask

Issue:
Allow users to disable repartitioning in `DataFrame.set_index`
`DataFrame.set_index` triggers expensive computations both to figure out new divisions and to repartition the dataframe according to them. Given how expensive such an (approximate) distributed sort is, users should have the chance to disable this behavior with a `repartition=True` keyword and instead simply set the index on the indivdual existing partitions. This would allow users to use methods like `.loc` without the unnecessary sorting overhead.


Discussion:
I would call the keyword `sort=False` personally.  Repartition doesn't seem right to me.  That operation usually implies that we maintain the current order, but shift around the boundaries.  The bigger difference here, I think, is that we're no longer sorting the dataset.

So to be clear:

```python
df = df.set_index("column-name", sort=False)
df.divisions
```
```
[None, None, None, None, ...]
```

Patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -5073,21 +5073,23 @@ def set_index(
         npartitions: int | Literal["auto"] | None = None,
         divisions: Sequence | None = None,
         inplace: bool = False,
+        sort: bool = True,
         **kwargs,
     ):
         """Set the DataFrame index (row labels) using an existing column.
 
-        This realigns the dataset to be sorted by a new column. This can have a
+        If ``sort=False``, this function operates exactly like ``pandas.set_index``
+        and sets the index on the DataFrame. If ``sort=True`` (default),
+        this function also sorts the DataFrame by the new index. This can have a
         significant impact on performance, because joins, groupbys, lookups, etc.
         are all much faster on that column. However, this performance increase
         comes with a cost, sorting a parallel dataset requires expensive shuffles.
         Often we ``set_index`` once directly after data ingest and filtering and
         then perform many cheap computations off of the sorted dataset.
 
-        This function operates exactly like ``pandas.set_index`` except with
-        different performance costs (dask dataframe ``set_index`` is much more expensive).
-        Under normal operation this function does an initial pass over the index column
-        to compute approximate quantiles to serve as future divisions. It then passes
+        With ``sort=True``, this function is much more expensive. Under normal
+        operation this function does an initial pass over the index column to
+        compute approximate quantiles to serve as future divisions. It then passes
         over the data a second time, splitting up each input partition into several
         pieces and sharing those pieces to all of the output partitions now in
         sorted order.
@@ -5127,6 +5129,10 @@ def set_index(
         inplace: bool, optional
             Modifying the DataFrame in place is not supported by Dask.
             Defaults to False.
+        sort: bool, optional
+            If ``True``, sort the DataFrame by the new index. Otherwise
+            set the index on the individual existing partitions.
+            Defaults to ``True``.
         shuffle: string, 'disk' or 'tasks', optional
             Either ``'disk'`` for single-node operation or ``'tasks'`` for
             distributed operation.  Will be inferred by your current scheduler.
@@ -5261,6 +5267,7 @@ def set_index(
                 drop=drop,
                 npartitions=npartitions,
                 divisions=divisions,
+                sort=sort,
                 **kwargs,
             )
 
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -211,9 +211,15 @@ def set_index(
     upsample: float = 1.0,
     divisions: Sequence | None = None,
     partition_size: float = 128e6,
+    sort: bool = True,
     **kwargs,
 ) -> DataFrame:
     """See _Frame.set_index for docstring"""
+    if not sort:
+        return df.map_partitions(
+            M.set_index, index, align_dataframes=False, drop=drop, **kwargs
+        ).clear_divisions()
+
     if npartitions == "auto":
         repartition = True
         npartitions = max(100, df.npartitions)


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
@@ -304,6 +304,34 @@ def test_set_index_3(shuffle_method):
     assert ddf2.npartitions == ddf.npartitions
 
 
+@pytest.mark.parametrize("drop", (True, False))
+@pytest.mark.parametrize("append", (True, False))
+def test_set_index_no_sort(drop, append):
+    """
+    GH10333 - Allow setting index on existing partitions without
+    computing new divisions and repartitioning.
+    """
+    df = pd.DataFrame({"col1": [2, 4, 1, 3, 5], "col2": [1, 2, 3, 4, 5]})
+    ddf = dd.from_pandas(df, npartitions=2)
+
+    assert ddf.npartitions > 1
+
+    # Default is sort=True
+    # Index in ddf will be same values, but sorted
+    df_result = df.set_index("col1")
+    ddf_result = ddf.set_index("col1")
+    assert ddf_result.known_divisions
+    assert_eq(ddf_result, df_result.sort_index(), sort_results=False)
+
+    # Unknown divisions and index remains unsorted when sort is False
+    # and thus equal to pandas set_index, adding extra kwargs also supported by
+    # pandas set_index to ensure they're forwarded.
+    df_result = df.set_index("col1", drop=drop, append=append)
+    ddf_result = ddf.set_index("col1", sort=False, drop=drop, append=append)
+    assert not ddf_result.known_divisions
+    assert_eq(ddf_result, df_result, sort_results=False)
+
+
 def test_shuffle_sort(shuffle_method):
     df = pd.DataFrame({"x": [1, 2, 3, 2, 1], "y": [9, 8, 7, 1, 5]})
     ddf = dd.from_pandas(df, npartitions=3)