Repository: dask/dask

Issue:
 AttributeError: 'tuple' object has no attribute 'groupby
When I create a `DataFrame`:

```python
df = dataframe.from_pandas(
    pd.DataFrame({"x": range(10), "y": ["a"] * 5 + ["b"] * 5, "z": 1}), 5)
```

And try to set index: with the same number of partitions

```
df.set_index("x", divisions=[0, 2, 4, 6, 8, 10], sorted=True).groupby("y").count().compute()
```

it works as expected. But if number of partitions differs:

```
df.set_index("x", divisions=[0, 2, 5, 8, 10, 12, 14], sorted=True).groupby("y").count().compute()
```

I get confusing exception shown below:

```
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-205-e8ff0304bc2c> in <module>()
----> 1 df.set_index("x", divisions=[0, 2, 5, 8, 10, 12, 14], sorted=True).groupby("y").count().compute()

/path/to/python3.6/site-packages/dask/base.py in compute(self, **kwargs)
     92             Extra keywords to forward to the scheduler ``get`` function.
     93         """
---> 94         (result,) = compute(self, traverse=False, **kwargs)
     95         return result
     96 

/path/to/python3.6/site-packages/dask/base.py in compute(*args, **kwargs)
    199     dsk = collections_to_dsk(variables, optimize_graph, **kwargs)
    200     keys = [var._keys() for var in variables]
--> 201     results = get(dsk, keys, **kwargs)
    202 
    203     results_iter = iter(results)

/path/to/python3.6/site-packages/dask/threaded.py in get(dsk, result, cache, num_workers, **kwargs)
     74     results = get_async(pool.apply_async, len(pool._pool), dsk, result,
     75                         cache=cache, get_id=_thread_get_id,
---> 76                         **kwargs)
     77 
     78     # Cleanup pools associated to dead threads

/path/to/python3.6/site-packages/dask/async.py in get_async(apply_async, num_workers, dsk, result, cache, get_id, raise_on_exception, rerun_exceptions_locally, callbacks, dumps, loads, **kwargs)
    498                     _execute_task(task, data)  # Re-execute locally
    499                 else:
--> 500                     raise(remote_exception(res, tb))
    501             state['cache'][key] = res
    502             finish_task(dsk, key, state, results, keyorder.get)

AttributeError: 'tuple' object has no attribute 'groupby'

Traceback
---------
  File "/path/to/python3.6/site-packages/dask/async.py", line 266, in execute_task
    result = _execute_task(task, data)
  File "/path/to/python3.6/site-packages/dask/async.py", line 247, in _execute_task
    return func(*args2)
  File "/path/to/python3.6/site-packages/dask/compatibility.py", line 47, in apply
    return func(*args, **kwargs)
  File "/path/to/python3.6/site-packages/dask/dataframe/groupby.py", line 175, in _apply_chunk
    g = _groupby_raise_unaligned(df, by=index)
  File "/path/to/python3.6/site-packages/dask/dataframe/groupby.py", line 133, in _groupby_raise_unaligned
    return df.groupby(**kwargs)
```
Versions

- Python                     3.6.0 
- pandas                    0.19.2
- dask                        0.14.0


Discussion:
Replicated.  Thanks for the excellent bug report @LostInOverflow .  

I'm actually looking at shuffling / groupbys / set_index now, so this is good timing.
Digging into this now.  It looks like this is only an issue in the `sorted=True` case.  In this case it looks like we are defining the `divisions=` keyword somewhat differently, which is causing some confusion.  

When using sorted=True we state that we know that the data is already sorted by the new index.  In these cases we can also state the known values of those divisions along the index using the divisions= keyword.  This is an entirely different meaning from how we usually use the divisions= keyword (which is usually "this is what I want my divisions to look like").  This double meaning of the same keyword is obviously an issue.  

Short term @LostInOverflow you probably want this:

```python
df.set_index('x', sorted=True).repartition(divisions=...)
```
Thank you so much @mrocklin!

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -2218,9 +2218,9 @@ def select_dtypes(self, include=None, exclude=None):
         cs = self._meta.select_dtypes(include=include, exclude=exclude).columns
         return self[list(cs)]
 
-    def set_index(self, other, drop=True, sorted=False, npartitions=None, **kwargs):
-        """
-        Set the DataFrame index (row labels) using an existing column
+    def set_index(self, other, drop=True, sorted=False, npartitions=None,
+                  divisions=None, **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
         significant impact on performance, because joins, groupbys, lookups, etc.
@@ -2260,9 +2260,12 @@ def set_index(self, other, drop=True, sorted=False, npartitions=None, **kwargs):
         divisions: list, optional
             Known values on which to separate index values of the partitions.
             See http://dask.pydata.org/en/latest/dataframe-design.html#partitions
-            Defaults to computing this with a single pass over the data
+            Defaults to computing this with a single pass over the data. Note
+            that if ``sorted=True``, specified divisions are assumed to match
+            the existing partitions in the data. If this is untrue, you should
+            leave divisions empty and call ``repartition`` after ``set_index``.
         compute: bool
-            Whether or not to trigger an immediate computation. Defaults to True.
+            Whether or not to trigger an immediate computation. Defaults to False.
 
         Examples
         --------
@@ -2279,24 +2282,20 @@ def set_index(self, other, drop=True, sorted=False, npartitions=None, **kwargs):
         >>> divisions = pd.date_range('2000', '2010', freq='1D')
         >>> df2 = df.set_index('timestamp', sorted=True, divisions=divisions)  # doctest: +SKIP
         """
-        if sorted:
+        pre_sorted = sorted
+        del sorted
+
+        if divisions is not None:
+            check_divisions(divisions)
+
+        if pre_sorted:
             from .shuffle import set_sorted_index
-            return set_sorted_index(self, other, drop=drop, **kwargs)
+            return set_sorted_index(self, other, drop=drop, divisions=divisions,
+                                    **kwargs)
         else:
             from .shuffle import set_index
             return set_index(self, other, drop=drop, npartitions=npartitions,
-                             **kwargs)
-
-    def set_partition(self, column, divisions, **kwargs):
-        """ Set explicit divisions for new column index
-
-        >>> df2 = df.set_partition('new-index-column', divisions=[10, 20, 50])  # doctest: +SKIP
-
-        See Also
-        --------
-        set_index
-        """
-        raise Exception("Deprecated, use set_index(..., divisions=...) instead")
+                             divisions=divisions, **kwargs)
 
     @derived_from(pd.DataFrame)
     def nlargest(self, n=5, columns=None, split_every=None):
@@ -3467,6 +3466,17 @@ def _take_last(a, skipna=True):
             return last_row.values[0]
 
 
+def check_divisions(divisions):
+    if not isinstance(divisions, (list, tuple)):
+        raise ValueError('New division must be list or tuple')
+    divisions = list(divisions)
+    if divisions != sorted(divisions):
+        raise ValueError('New division must be sorted')
+    if len(divisions[:-1]) != len(list(unique(divisions[:-1]))):
+        msg = 'New division must be unique, except for the last element'
+        raise ValueError(msg)
+
+
 def repartition_divisions(a, b, name, out1, out2, force=False):
     """ dask graph to repartition dataframe by new divisions
 
@@ -3499,20 +3509,12 @@ def repartition_divisions(a, b, name, out1, out2, force=False):
      ('c', 1): ('b', 2),
      ('c', 2): ('b', 3)}
     """
-    if not isinstance(b, (list, tuple)):
-        raise ValueError('New division must be list or tuple')
-    b = list(b)
+    check_divisions(b)
 
     if len(b) < 2:
         # minimum division is 2 elements, like [0, 0]
         raise ValueError('New division must be longer than 2 elements')
 
-    if b != sorted(b):
-        raise ValueError('New division must be sorted')
-    if len(b[:-1]) != len(list(unique(b[:-1]))):
-        msg = 'New division must be unique, except for the last element'
-        raise ValueError(msg)
-
     if force:
         if a[0] < b[0]:
             msg = ('left side of the new division must be equal or smaller '
diff --git a/dask/dataframe/multi.py b/dask/dataframe/multi.py
--- a/dask/dataframe/multi.py
+++ b/dask/dataframe/multi.py
@@ -101,8 +101,8 @@ def align_partitions(*dfs):
         raise ValueError("dfs contains no DataFrame and Series")
     if not all(df.known_divisions for df in dfs1):
         raise ValueError("Not all divisions are known, can't align "
-                         "partitions. Please use `set_index` or "
-                         "`set_partition` to set the index.")
+                         "partitions. Please use `set_index` "
+                         "to set the index.")
 
     divisions = list(unique(merge_sorted(*[df.divisions for df in dfs1])))
     dfs2 = [df.repartition(divisions, force=True)
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -459,6 +459,15 @@ def set_sorted_index(df, index, drop=True, divisions=None, **kwargs):
 
     if not divisions:
         divisions = compute_divisions(result, **kwargs)
+    elif len(divisions) != len(df.divisions):
+        msg = ("When doing `df.set_index(col, sorted=True, divisions=...)`, "
+               "divisions indicates known splits in the index column. In this "
+               "case divisions must be the same length as the existing "
+               "divisions in `df`\n\n"
+               "If the intent is to repartition into new divisions after "
+               "setting the index, you probably want:\n\n"
+               "`df.set_index(col, sorted=True).repartition(divisions=divisions)`")
+        raise ValueError(msg)
 
     result.divisions = tuple(divisions)
     return result


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -1,5 +1,4 @@
 import sys
-from copy import copy
 from operator import getitem, add
 from itertools import product
 
@@ -9,8 +8,6 @@
 import pytest
 
 import dask
-from dask.async import get_sync
-from dask import delayed
 from dask.utils import ignoring, put_lines
 import dask.dataframe as dd
 
@@ -174,77 +171,6 @@ def test_index_names():
     assert ddf.index.compute().name == 'x'
 
 
-def test_set_index():
-    dsk = {('x', 0): pd.DataFrame({'a': [1, 2, 3], 'b': [4, 2, 6]},
-                                  index=[0, 1, 3]),
-           ('x', 1): pd.DataFrame({'a': [4, 5, 6], 'b': [3, 5, 8]},
-                                  index=[5, 6, 8]),
-           ('x', 2): pd.DataFrame({'a': [7, 8, 9], 'b': [9, 1, 8]},
-                                  index=[9, 9, 9])}
-    d = dd.DataFrame(dsk, 'x', meta, [0, 4, 9, 9])
-    full = d.compute()
-
-    d2 = d.set_index('b', npartitions=3)
-    assert d2.npartitions == 3
-    assert d2.index.name == 'b'
-    assert_eq(d2, full.set_index('b'))
-
-    d3 = d.set_index(d.b, npartitions=3)
-    assert d3.npartitions == 3
-    assert d3.index.name == 'b'
-    assert_eq(d3, full.set_index(full.b))
-
-    d4 = d.set_index('b')
-    assert d4.index.name == 'b'
-    assert_eq(d4, full.set_index('b'))
-
-
-def test_set_index_interpolate():
-    df = pd.DataFrame({'x': [4, 1, 1, 3, 3], 'y': [1., 1, 1, 1, 2]})
-    d = dd.from_pandas(df, 2)
-
-    d1 = d.set_index('x', npartitions=3)
-    assert d1.npartitions == 3
-    assert set(d1.divisions) == set([1, 2, 3, 4])
-
-    d2 = d.set_index('y', npartitions=3)
-    assert d2.divisions[0] == 1.
-    assert 1. < d2.divisions[1] < d2.divisions[2] < 2.
-    assert d2.divisions[3] == 2.
-
-
-def test_set_index_interpolate_int():
-    L = sorted(list(range(0, 200, 10)) * 2)
-    df = pd.DataFrame({'x': 2 * L})
-    d = dd.from_pandas(df, 2)
-    d1 = d.set_index('x', npartitions=10)
-    assert all(np.issubdtype(type(x), np.integer) for x in d1.divisions)
-
-
-def test_set_index_timezone():
-    s_naive = pd.Series(pd.date_range('20130101', periods=3))
-    s_aware = pd.Series(pd.date_range('20130101', periods=3, tz='US/Eastern'))
-    df = pd.DataFrame({'tz': s_aware, 'notz': s_naive})
-    d = dd.from_pandas(df, 2)
-
-    d1 = d.set_index('notz', npartitions=2)
-    s1 = pd.DatetimeIndex(s_naive.values, dtype=s_naive.dtype)
-    assert d1.divisions[0] == s_naive[0] == s1[0]
-    assert d1.divisions[-1] == s_naive[2] == s1[2]
-
-    # We currently lose "freq".  Converting data with pandas-defined dtypes
-    # to numpy or pure Python can be lossy like this.
-    d2 = d.set_index('tz', npartitions=2)
-    s2 = pd.DatetimeIndex(s_aware, dtype=s_aware.dtype)
-    assert d2.divisions[0] == s2[0]
-    assert d2.divisions[-1] == s2[2]
-    assert d2.divisions[0].tz == s2[0].tz
-    assert d2.divisions[0].tz is not None
-    s2badtype = pd.DatetimeIndex(s_aware.values, dtype=s_naive.dtype)
-    with pytest.raises(TypeError):
-        d2.divisions[0] == s2badtype[0]
-
-
 @pytest.mark.parametrize(
     'npartitions',
     [1, pytest.mark.xfail(2, reason='pandas join removes freq')]
@@ -258,49 +184,6 @@ def test_timezone_freq(npartitions):
     assert pdf.tz[0].freq == ddf.compute().tz[0].freq == ddf.tz.compute()[0].freq
 
 
-@pytest.mark.parametrize('drop', [True, False])
-def test_set_index_drop(drop):
-
-    pdf = pd.DataFrame({'A': list('ABAABBABAA'),
-                        'B': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
-                        'C': [1, 2, 3, 2, 1, 3, 2, 4, 2, 3]})
-    ddf = dd.from_pandas(pdf, 3)
-
-    assert_eq(ddf.set_index('A', drop=drop),
-              pdf.set_index('A', drop=drop))
-    assert_eq(ddf.set_index('B', drop=drop),
-              pdf.set_index('B', drop=drop))
-    assert_eq(ddf.set_index('C', drop=drop),
-              pdf.set_index('C', drop=drop))
-    assert_eq(ddf.set_index(ddf.A, drop=drop),
-              pdf.set_index(pdf.A, drop=drop))
-    assert_eq(ddf.set_index(ddf.B, drop=drop),
-              pdf.set_index(pdf.B, drop=drop))
-    assert_eq(ddf.set_index(ddf.C, drop=drop),
-              pdf.set_index(pdf.C, drop=drop))
-
-    # numeric columns
-    pdf = pd.DataFrame({0: list('ABAABBABAA'),
-                        1: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
-                        2: [1, 2, 3, 2, 1, 3, 2, 4, 2, 3]})
-    ddf = dd.from_pandas(pdf, 3)
-    assert_eq(ddf.set_index(0, drop=drop),
-              pdf.set_index(0, drop=drop))
-    assert_eq(ddf.set_index(2, drop=drop),
-              pdf.set_index(2, drop=drop))
-
-
-def test_set_index_raises_error_on_bad_input():
-    df = pd.DataFrame({'a': [1, 2, 3, 4, 5, 6, 7],
-                       'b': [7, 6, 5, 4, 3, 2, 1]})
-    ddf = dd.from_pandas(df, 2)
-
-    msg = r"Dask dataframe does not yet support multi-indexes"
-    with pytest.raises(NotImplementedError) as err:
-        ddf.set_index(['a', 'b'])
-    assert msg in str(err.value)
-
-
 def test_rename_columns():
     # GH 819
     df = pd.DataFrame({'a': [1, 2, 3, 4, 5, 6, 7],
@@ -662,32 +545,6 @@ def test_drop_duplicates_subset():
                       ddf.drop_duplicates(subset=ss, **kwarg))
 
 
-def test_set_partition():
-    d2 = d.set_index('b', divisions=[0, 2, 9])
-    assert d2.divisions == (0, 2, 9)
-    expected = full.set_index('b')
-    assert_eq(d2, expected)
-
-
-def test_set_partition_compute():
-    d2 = d.set_index('b', divisions=[0, 2, 9], compute=False)
-    d3 = d.set_index('b', divisions=[0, 2, 9], compute=True)
-
-    assert_eq(d2, d3)
-    assert_eq(d2, full.set_index('b'))
-    assert_eq(d3, full.set_index('b'))
-    assert len(d2.dask) > len(d3.dask)
-
-    d4 = d.set_index(d.b, divisions=[0, 2, 9], compute=False)
-    d5 = d.set_index(d.b, divisions=[0, 2, 9], compute=True)
-    exp = full.copy()
-    exp.index = exp.b
-    assert_eq(d4, d5)
-    assert_eq(d4, exp)
-    assert_eq(d5, exp)
-    assert len(d4.dask) > len(d5.dask)
-
-
 def test_get_partition():
     pdf = pd.DataFrame(np.random.randn(10, 5), columns=list('abcde'))
     ddf = dd.from_pandas(pdf, 3)
@@ -1145,16 +1002,6 @@ def test_series_round():
     assert_eq(s.round(), ps.round())
 
 
-def test_set_partition_2():
-    df = pd.DataFrame({'x': [1, 2, 3, 4, 5, 6], 'y': list('abdabd')})
-    ddf = dd.from_pandas(df, 2)
-
-    result = ddf.set_index('y', divisions=['a', 'c', 'd'])
-    assert result.divisions == ('a', 'c', 'd')
-
-    assert list(result.compute(get=get_sync).index[-2:]) == ['d', 'd']
-
-
 @pytest.mark.slow
 def test_repartition():
     def _check_split_data(orig, d):
@@ -2195,46 +2042,6 @@ def iseven(x):
               b.y.groupby(iseven).sum())
 
 
-def test_set_index_sorted_true():
-    df = pd.DataFrame({'x': [1, 2, 3, 4],
-                       'y': [10, 20, 30, 40],
-                       'z': [4, 3, 2, 1]})
-    a = dd.from_pandas(df, 2, sort=False)
-    assert not a.known_divisions
-
-    b = a.set_index('x', sorted=True)
-    assert b.known_divisions
-    assert set(a.dask).issubset(set(b.dask))
-
-    for drop in [True, False]:
-        assert_eq(a.set_index('x', drop=drop),
-                  df.set_index('x', drop=drop))
-        assert_eq(a.set_index(a.x, sorted=True, drop=drop),
-                  df.set_index(df.x, drop=drop))
-        assert_eq(a.set_index(a.x + 1, sorted=True, drop=drop),
-                  df.set_index(df.x + 1, drop=drop))
-
-    with pytest.raises(ValueError):
-        a.set_index(a.z, sorted=True)
-
-
-def test_compute_divisions():
-    from dask.dataframe.shuffle import compute_divisions
-    df = pd.DataFrame({'x': [1, 2, 3, 4],
-                       'y': [10, 20, 30, 40],
-                       'z': [4, 3, 2, 1]},
-                      index=[1, 3, 10, 20])
-    a = dd.from_pandas(df, 2, sort=False)
-    assert not a.known_divisions
-
-    divisions = compute_divisions(a)
-    b = copy(a)
-    b.divisions = divisions
-
-    assert_eq(a, b, check_divisions=False)
-    assert b.known_divisions
-
-
 def test_methods_tokenize_differently():
     df = pd.DataFrame({'x': [1, 2, 3, 4]})
     df = dd.from_pandas(df, npartitions=1)
@@ -2242,13 +2049,6 @@ def test_methods_tokenize_differently():
             df.x.map_partitions(lambda x: pd.Series(x.max()))._name)
 
 
-def test_sorted_index_single_partition():
-    df = pd.DataFrame({'x': [1, 2, 3, 4], 'y': [1, 0, 1, 0]})
-    ddf = dd.from_pandas(df, npartitions=1)
-    assert_eq(ddf.set_index('x', sorted=True),
-              df.set_index('x'))
-
-
 def _assert_info(df, ddf, memory_usage=True):
     from io import StringIO
     assert isinstance(df, pd.DataFrame)
@@ -2470,20 +2270,6 @@ def test_getitem_multilevel():
     assert_eq(pdf[[('A', '0'), ('B', '1')]], ddf[[('A', '0'), ('B', '1')]])
 
 
-def test_set_index_sorted_min_max_same():
-    a = pd.DataFrame({'x': [1, 2, 3], 'y': [0, 0, 0]})
-    b = pd.DataFrame({'x': [1, 2, 3], 'y': [1, 1, 1]})
-
-    aa = delayed(a)
-    bb = delayed(b)
-
-    df = dd.from_delayed([aa, bb], meta=a)
-    assert not df.known_divisions
-
-    df2 = df.set_index('y', sorted=True)
-    assert df2.divisions == (0, 1, 1)
-
-
 def test_diff():
     df = pd.DataFrame(np.random.randn(100, 5), columns=list('abcde'))
     ddf = dd.from_pandas(df, 5)
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
@@ -4,8 +4,10 @@
 import pickle
 import numpy as np
 import string
+from copy import copy
 
 import dask
+from dask import delayed
 import dask.dataframe as dd
 from dask.threaded import get as threaded_get
 from dask.multiprocessing import get as mp_get
@@ -148,7 +150,7 @@ def test_partitioning_index_categorical_on_values():
 
 
 @pytest.mark.parametrize('npartitions', [1, 4, 7, pytest.mark.slow(23)])
-def test_set_partition_tasks(npartitions):
+def test_set_index_tasks(npartitions):
     df = pd.DataFrame({'x': np.random.random(100),
                        'y': np.random.random(100) // 0.2},
                       index=np.random.random(100))
@@ -188,7 +190,7 @@ def test_set_index_self_index(shuffle):
 
 
 @pytest.mark.parametrize('shuffle', ['tasks'])
-def test_set_partition_names(shuffle):
+def test_set_index_names(shuffle):
     df = pd.DataFrame({'x': np.random.random(100),
                        'y': np.random.random(100) // 0.2},
                       index=np.random.random(100))
@@ -206,7 +208,7 @@ def test_set_partition_names(shuffle):
 
 
 @pytest.mark.parametrize('shuffle', ['disk', 'tasks'])
-def test_set_partition_tasks_2(shuffle):
+def test_set_index_tasks_2(shuffle):
     df = dd.demo.make_timeseries(
         '2000', '2004', {'value': float, 'name': str, 'id': int},
         freq='2H', partition_freq='1M', seed=1)
@@ -216,7 +218,7 @@ def test_set_partition_tasks_2(shuffle):
 
 
 @pytest.mark.parametrize('shuffle', ['disk', 'tasks'])
-def test_set_partition_tasks_3(shuffle):
+def test_set_index_tasks_3(shuffle):
     df = pd.DataFrame(np.random.random((10, 2)), columns=['x', 'y'])
     ddf = dd.from_pandas(df, npartitions=5)
 
@@ -286,10 +288,73 @@ def throw(*args, **kwargs):
 
     with dask.set_options(get=throw):
         ddf2 = ddf.set_index('x', divisions=[1, 3, 5])
+    assert ddf2.divisions == (1, 3, 5)
 
     df2 = df.set_index('x')
     assert_eq(ddf2, df2)
 
+    # Divisions must be sorted
+    with pytest.raises(ValueError):
+        ddf.set_index('x', divisions=[3, 1, 5])
+
+
+def test_set_index_divisions_2():
+    df = pd.DataFrame({'x': [1, 2, 3, 4, 5, 6], 'y': list('abdabd')})
+    ddf = dd.from_pandas(df, 2)
+
+    result = ddf.set_index('y', divisions=['a', 'c', 'd'])
+    assert result.divisions == ('a', 'c', 'd')
+
+    assert list(result.compute(get=get_sync).index[-2:]) == ['d', 'd']
+
+
+def test_set_index_divisions_compute():
+    d2 = d.set_index('b', divisions=[0, 2, 9], compute=False)
+    d3 = d.set_index('b', divisions=[0, 2, 9], compute=True)
+
+    assert_eq(d2, d3)
+    assert_eq(d2, full.set_index('b'))
+    assert_eq(d3, full.set_index('b'))
+    assert len(d2.dask) > len(d3.dask)
+
+    d4 = d.set_index(d.b, divisions=[0, 2, 9], compute=False)
+    d5 = d.set_index(d.b, divisions=[0, 2, 9], compute=True)
+    exp = full.copy()
+    exp.index = exp.b
+    assert_eq(d4, d5)
+    assert_eq(d4, exp)
+    assert_eq(d5, exp)
+    assert len(d4.dask) > len(d5.dask)
+
+
+def test_set_index_divisions_sorted():
+    p1 = pd.DataFrame({'x': [10, 11, 12], 'y': ['a', 'a', 'a']})
+    p2 = pd.DataFrame({'x': [13, 14, 15], 'y': ['b', 'b', 'c']})
+    p3 = pd.DataFrame({'x': [16, 17, 18], 'y': ['d', 'e', 'e']})
+
+    ddf = dd.DataFrame({('x', 0): p1, ('x', 1): p2, ('x', 2): p3},
+                       'x', p1, [None, None, None, None])
+    df = ddf.compute()
+
+    def throw(*args, **kwargs):
+        raise Exception("Shouldn't have computed")
+
+    with dask.set_options(get=throw):
+        res = ddf.set_index('x', divisions=[10, 13, 16, 18], sorted=True)
+    assert_eq(res, df.set_index('x'))
+
+    with dask.set_options(get=throw):
+        res = ddf.set_index('y', divisions=['a', 'b', 'd', 'e'], sorted=True)
+    assert_eq(res, df.set_index('y'))
+
+    # with sorted=True, divisions must be same length as df.divisions
+    with pytest.raises(ValueError):
+        ddf.set_index('y', divisions=['a', 'b', 'c', 'd', 'e'], sorted=True)
+
+    # Divisions must be sorted
+    with pytest.raises(ValueError):
+        ddf.set_index('y', divisions=['a', 'b', 'd', 'c'], sorted=True)
+
 
 @pytest.mark.parametrize('shuffle', ['disk', 'tasks'])
 def test_set_index_reduces_partitions_small(shuffle):
@@ -333,6 +398,180 @@ def test_set_index_detects_sorted_data(shuffle):
     assert len(ddf2.dask) < ddf.npartitions * 4
 
 
+def test_set_index():
+    dsk = {('x', 0): pd.DataFrame({'a': [1, 2, 3], 'b': [4, 2, 6]},
+                                  index=[0, 1, 3]),
+           ('x', 1): pd.DataFrame({'a': [4, 5, 6], 'b': [3, 5, 8]},
+                                  index=[5, 6, 8]),
+           ('x', 2): pd.DataFrame({'a': [7, 8, 9], 'b': [9, 1, 8]},
+                                  index=[9, 9, 9])}
+    d = dd.DataFrame(dsk, 'x', meta, [0, 4, 9, 9])
+    full = d.compute()
+
+    d2 = d.set_index('b', npartitions=3)
+    assert d2.npartitions == 3
+    assert d2.index.name == 'b'
+    assert_eq(d2, full.set_index('b'))
+
+    d3 = d.set_index(d.b, npartitions=3)
+    assert d3.npartitions == 3
+    assert d3.index.name == 'b'
+    assert_eq(d3, full.set_index(full.b))
+
+    d4 = d.set_index('b')
+    assert d4.index.name == 'b'
+    assert_eq(d4, full.set_index('b'))
+
+
+def test_set_index_interpolate():
+    df = pd.DataFrame({'x': [4, 1, 1, 3, 3], 'y': [1., 1, 1, 1, 2]})
+    d = dd.from_pandas(df, 2)
+
+    d1 = d.set_index('x', npartitions=3)
+    assert d1.npartitions == 3
+    assert set(d1.divisions) == set([1, 2, 3, 4])
+
+    d2 = d.set_index('y', npartitions=3)
+    assert d2.divisions[0] == 1.
+    assert 1. < d2.divisions[1] < d2.divisions[2] < 2.
+    assert d2.divisions[3] == 2.
+
+
+def test_set_index_interpolate_int():
+    L = sorted(list(range(0, 200, 10)) * 2)
+    df = pd.DataFrame({'x': 2 * L})
+    d = dd.from_pandas(df, 2)
+    d1 = d.set_index('x', npartitions=10)
+    assert all(np.issubdtype(type(x), np.integer) for x in d1.divisions)
+
+
+def test_set_index_timezone():
+    s_naive = pd.Series(pd.date_range('20130101', periods=3))
+    s_aware = pd.Series(pd.date_range('20130101', periods=3, tz='US/Eastern'))
+    df = pd.DataFrame({'tz': s_aware, 'notz': s_naive})
+    d = dd.from_pandas(df, 2)
+
+    d1 = d.set_index('notz', npartitions=2)
+    s1 = pd.DatetimeIndex(s_naive.values, dtype=s_naive.dtype)
+    assert d1.divisions[0] == s_naive[0] == s1[0]
+    assert d1.divisions[-1] == s_naive[2] == s1[2]
+
+    # We currently lose "freq".  Converting data with pandas-defined dtypes
+    # to numpy or pure Python can be lossy like this.
+    d2 = d.set_index('tz', npartitions=2)
+    s2 = pd.DatetimeIndex(s_aware, dtype=s_aware.dtype)
+    assert d2.divisions[0] == s2[0]
+    assert d2.divisions[-1] == s2[2]
+    assert d2.divisions[0].tz == s2[0].tz
+    assert d2.divisions[0].tz is not None
+    s2badtype = pd.DatetimeIndex(s_aware.values, dtype=s_naive.dtype)
+    with pytest.raises(TypeError):
+        d2.divisions[0] == s2badtype[0]
+
+
+@pytest.mark.parametrize('drop', [True, False])
+def test_set_index_drop(drop):
+    pdf = pd.DataFrame({'A': list('ABAABBABAA'),
+                        'B': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+                        'C': [1, 2, 3, 2, 1, 3, 2, 4, 2, 3]})
+    ddf = dd.from_pandas(pdf, 3)
+
+    assert_eq(ddf.set_index('A', drop=drop),
+              pdf.set_index('A', drop=drop))
+    assert_eq(ddf.set_index('B', drop=drop),
+              pdf.set_index('B', drop=drop))
+    assert_eq(ddf.set_index('C', drop=drop),
+              pdf.set_index('C', drop=drop))
+    assert_eq(ddf.set_index(ddf.A, drop=drop),
+              pdf.set_index(pdf.A, drop=drop))
+    assert_eq(ddf.set_index(ddf.B, drop=drop),
+              pdf.set_index(pdf.B, drop=drop))
+    assert_eq(ddf.set_index(ddf.C, drop=drop),
+              pdf.set_index(pdf.C, drop=drop))
+
+    # numeric columns
+    pdf = pd.DataFrame({0: list('ABAABBABAA'),
+                        1: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+                        2: [1, 2, 3, 2, 1, 3, 2, 4, 2, 3]})
+    ddf = dd.from_pandas(pdf, 3)
+    assert_eq(ddf.set_index(0, drop=drop),
+              pdf.set_index(0, drop=drop))
+    assert_eq(ddf.set_index(2, drop=drop),
+              pdf.set_index(2, drop=drop))
+
+
+def test_set_index_raises_error_on_bad_input():
+    df = pd.DataFrame({'a': [1, 2, 3, 4, 5, 6, 7],
+                       'b': [7, 6, 5, 4, 3, 2, 1]})
+    ddf = dd.from_pandas(df, 2)
+
+    msg = r"Dask dataframe does not yet support multi-indexes"
+    with pytest.raises(NotImplementedError) as err:
+        ddf.set_index(['a', 'b'])
+    assert msg in str(err.value)
+
+
+def test_set_index_sorted_true():
+    df = pd.DataFrame({'x': [1, 2, 3, 4],
+                       'y': [10, 20, 30, 40],
+                       'z': [4, 3, 2, 1]})
+    a = dd.from_pandas(df, 2, sort=False)
+    assert not a.known_divisions
+
+    b = a.set_index('x', sorted=True)
+    assert b.known_divisions
+    assert set(a.dask).issubset(set(b.dask))
+
+    for drop in [True, False]:
+        assert_eq(a.set_index('x', drop=drop),
+                  df.set_index('x', drop=drop))
+        assert_eq(a.set_index(a.x, sorted=True, drop=drop),
+                  df.set_index(df.x, drop=drop))
+        assert_eq(a.set_index(a.x + 1, sorted=True, drop=drop),
+                  df.set_index(df.x + 1, drop=drop))
+
+    with pytest.raises(ValueError):
+        a.set_index(a.z, sorted=True)
+
+
+def test_set_index_sorted_single_partition():
+    df = pd.DataFrame({'x': [1, 2, 3, 4], 'y': [1, 0, 1, 0]})
+    ddf = dd.from_pandas(df, npartitions=1)
+    assert_eq(ddf.set_index('x', sorted=True),
+              df.set_index('x'))
+
+
+def test_set_index_sorted_min_max_same():
+    a = pd.DataFrame({'x': [1, 2, 3], 'y': [0, 0, 0]})
+    b = pd.DataFrame({'x': [1, 2, 3], 'y': [1, 1, 1]})
+
+    aa = delayed(a)
+    bb = delayed(b)
+
+    df = dd.from_delayed([aa, bb], meta=a)
+    assert not df.known_divisions
+
+    df2 = df.set_index('y', sorted=True)
+    assert df2.divisions == (0, 1, 1)
+
+
+def test_compute_divisions():
+    from dask.dataframe.shuffle import compute_divisions
+    df = pd.DataFrame({'x': [1, 2, 3, 4],
+                       'y': [10, 20, 30, 40],
+                       'z': [4, 3, 2, 1]},
+                      index=[1, 3, 10, 20])
+    a = dd.from_pandas(df, 2, sort=False)
+    assert not a.known_divisions
+
+    divisions = compute_divisions(a)
+    b = copy(a)
+    b.divisions = divisions
+
+    assert_eq(a, b, check_divisions=False)
+    assert b.known_divisions
+
+
 def test_temporary_directory():
     df = pd.DataFrame({'x': np.random.random(100),
                        'y': np.random.random(100),