Repository: dask/dask

Issue:
set_index() on list with one item
Dask is raising NotImplementedError when passing a single-item list to set_index()

```python
import pandas as pd
import dask.dataframe as dd

df = pd.DataFrame({
    'id'  : ['1', '1', '2', '2', '3'],
    'col1': ['a', 'b', 'a', 'b', 'a']
})

ddf = dd.from_pandas(df, 1)

# Works
ddf = ddf.set_index("id")
ddf = ddf.reset_index()

ddf = ddf.set_index(["id"])

NotImplementedError: Dask dataframe does not yet support multi-indexes.
You tried to index with this index: ['id']
Indexes must be single columns only.
```

The user is providing a single column on which to index, but Dask is assuming that the list contains multiple values.


Discussion:
It would be great to fix this.  Do you think that this is something that
you could do Wes?

On Mon, Dec 30, 2019 at 2:07 PM Wes Roach <notifications@github.com> wrote:

> Dask is raising NotImplementedError when passing a single-item list to
> set_index()
>
> import pandas as pdimport dask.dataframe as dd
>
> df = pd.DataFrame({
>     'id'  : ['1', '1', '2', '2', '3'],
>     'col1': ['a', 'b', 'a', 'b', 'a']
> })
>
> ddf = dd.from_pandas(df, 1)
> # Works
> ddf = ddf.set_index("id")
> ddf = ddf.reset_index()
>
> ddf = ddf.set_index(["id"])
> NotImplementedError: Dask dataframe does not yet support multi-indexes.
> You tried to index with this index: ['id']
> Indexes must be single columns only.
>
> The user is providing a single column on which to index, but Dask is
> assuming that the list contains multiple values.
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/5755?email_source=notifications&email_token=AACKZTE6IGXKEL4FYY4J5WTQ3JWK3A5CNFSM4KBQSGGKYY3PNVWWK3TUL52HS4DFUVEXG43VMWVGG33NNVSW45C7NFSM4IDNFPTA>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTG6DWYR4B5MZIQ6OVTQ3JWK3ANCNFSM4KBQSGGA>
> .
>

Yeah, I'll take a look sometime this week.

Resolving patch:
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -34,11 +34,19 @@ def set_index(
     if isinstance(index, Series) and index._name == df.index._name:
         return df
     if isinstance(index, (DataFrame, tuple, list)):
-        raise NotImplementedError(
-            "Dask dataframe does not yet support multi-indexes.\n"
-            "You tried to index with this index: %s\n"
-            "Indexes must be single columns only." % str(index)
-        )
+        # Accept ["a"], but not [["a"]]
+        if (
+            isinstance(index, list)
+            and len(index) == 1
+            and not isinstance(index[0], list)  # if index = [["a"]], leave it that way
+        ):
+            index = index[0]
+        else:
+            raise NotImplementedError(
+                "Dask dataframe does not yet support multi-indexes.\n"
+                "You tried to index with this index: %s\n"
+                "Indexes must be single columns only." % str(index)
+            )
 
     if npartitions == "auto":
         repartition = True


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
@@ -556,6 +556,10 @@ def test_set_index():
     assert d4.index.name == "b"
     assert_eq(d4, full.set_index("b"))
 
+    d5 = d.set_index(["b"])
+    assert d5.index.name == "b"
+    assert_eq(d5, full.set_index(["b"]))
+
 
 def test_set_index_interpolate():
     df = pd.DataFrame({"x": [4, 1, 1, 3, 3], "y": [1.0, 1, 1, 1, 2]})
@@ -643,6 +647,14 @@ def test_set_index_raises_error_on_bad_input():
         ddf.set_index(["a", "b"])
     assert msg in str(err.value)
 
+    with pytest.raises(NotImplementedError) as err:
+        ddf.set_index([["a", "b"]])
+    assert msg in str(err.value)
+
+    with pytest.raises(NotImplementedError) as err:
+        ddf.set_index([["a"]])
+    assert msg in str(err.value)
+
 
 def test_set_index_sorted_true():
     df = pd.DataFrame({"x": [1, 2, 3, 4], "y": [10, 20, 20, 40], "z": [4, 3, 2, 1]})