Repository: dask/dask

Issue:
DataFrame.loc does not work with list of indices
 If `i = [3, 5, 1]` and `df` is a pandas data frame this returns the numbered rows in the defined order: `df.loc[i]`. 

This is the expected behavior for `dask.dataframe.DataFrame` but if `df` is such an object then `df.loc[i]` does not work. Note that you _can_ give it a slice: `df.loc[3:6]` for example.

On a side note, although using numbers in the example above, `loc` is for labels and as such we could expect partitions made of indices (or multindices) of strings, datetimes, and other objects...



Discussion:
This is reasonable to do and not particularly hard.  Just needs someone with a bit of time and expertise.

For short lists you could also use a combination of a list comprehension and `dd.concat`

```
dd.concat([df.loc[i] for i in L])
```

This would be inefficient for large `L` though

What about this one for larger lists?

`ddf.map_partitions(lambda x: x[x.index.isin(L)])`

We can use ``_partition_of_index_value`` to map input labels to partitions. Then simply use ``.loc[partitioned_list]`` for each.

- https://github.com/dask/dask/blob/master/dask/dataframe/indexing.py#L164

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -1424,7 +1424,7 @@ def first(self, offset):
 
         offset = pd.tseries.frequencies.to_offset(offset)
         date = self.divisions[0] + offset
-        end = self.loc._partition_of_index_value(date)
+        end = self.loc._get_partitions(date)
 
         include_right = offset.isAnchored() or not hasattr(offset, '_inc')
 
@@ -1449,7 +1449,7 @@ def last(self, offset):
 
         offset = pd.tseries.frequencies.to_offset(offset)
         date = self.divisions[-1] - offset
-        start = self.loc._partition_of_index_value(date)
+        start = self.loc._get_partitions(date)
 
         if start == 0:
             divs = self.divisions
diff --git a/dask/dataframe/indexing.py b/dask/dataframe/indexing.py
--- a/dask/dataframe/indexing.py
+++ b/dask/dataframe/indexing.py
@@ -1,6 +1,7 @@
 from __future__ import absolute_import, division, print_function
 
 from datetime import datetime
+from collections import defaultdict
 
 from toolz import merge
 import bisect
@@ -58,11 +59,18 @@ def _loc(self, iindexer, cindexer):
 
             if isinstance(iindexer, slice):
                 return self._loc_slice(iindexer, cindexer)
+            elif isinstance(iindexer, list):
+                return self._loc_list(iindexer, cindexer)
             else:
                 # element should raise KeyError
                 return self._loc_element(iindexer, cindexer)
         else:
-            if not isinstance(iindexer, slice):
+            if isinstance(iindexer, list):
+                # applying map_pattition to each partitions
+                # results in duplicated NaN rows
+                msg = 'Cannot index with list against unknown division'
+                raise KeyError(msg)
+            elif not isinstance(iindexer, slice):
                 iindexer = slice(iindexer, iindexer)
 
             meta = self._make_meta(iindexer, cindexer)
@@ -83,11 +91,27 @@ def _loc_series(self, iindexer, cindexer):
         return self.obj.map_partitions(methods.loc, iindexer, cindexer,
                                        token='loc-series', meta=meta)
 
-    def _loc_element(self, iindexer, cindexer):
-
+    def _loc_list(self, iindexer, cindexer):
         name = 'loc-%s' % tokenize(iindexer, self.obj)
+        parts = self._get_partitions(iindexer)
+        dsk = {}
+
+        divisions = []
+        items = sorted(parts.items())
+        for i, (div, indexer) in enumerate(items):
+            dsk[name, i] = (methods.loc, (self._name, div),
+                            indexer, cindexer)
+            # append minimum value as division
+            divisions.append(sorted(indexer)[0])
+        # append maximum value of the last division
+        divisions.append(sorted(items[-1][1])[-1])
+        meta = self._make_meta(iindexer, cindexer)
+        return new_dd_object(merge(self.obj.dask, dsk), name,
+                             meta=meta, divisions=divisions)
 
-        part = self._partition_of_index_value(iindexer)
+    def _loc_element(self, iindexer, cindexer):
+        name = 'loc-%s' % tokenize(iindexer, self.obj)
+        part = self._get_partitions(iindexer)
 
         if iindexer < self.obj.divisions[0] or iindexer > self.obj.divisions[-1]:
             raise KeyError('the label [%s] is not in the index' % str(iindexer))
@@ -99,8 +123,12 @@ def _loc_element(self, iindexer, cindexer):
         return new_dd_object(merge(self.obj.dask, dsk), name,
                              meta=meta, divisions=[iindexer, iindexer])
 
-    def _partition_of_index_value(self, key):
-        return _partition_of_index_value(self.obj.divisions, key)
+    def _get_partitions(self, keys):
+        if isinstance(keys, list):
+            return _partitions_of_index_values(self.obj.divisions, keys)
+        else:
+            # element
+            return _partition_of_index_value(self.obj.divisions, keys)
 
     def _coerce_loc_index(self, key):
         return _coerce_loc_index(self.obj.divisions, key)
@@ -112,11 +140,11 @@ def _loc_slice(self, iindexer, cindexer):
         assert iindexer.step in (None, 1)
 
         if iindexer.start is not None:
-            start = self._partition_of_index_value(iindexer.start)
+            start = self._get_partitions(iindexer.start)
         else:
             start = 0
         if iindexer.stop is not None:
-            stop = self._partition_of_index_value(iindexer.stop)
+            stop = self._get_partitions(iindexer.stop)
         else:
             stop = self.obj.npartitions - 1
 
@@ -181,6 +209,29 @@ def _partition_of_index_value(divisions, val):
     return min(len(divisions) - 2, max(0, i - 1))
 
 
+def _partitions_of_index_values(divisions, values):
+    """ Return defaultdict of division and values pairs
+    Each key corresponds to the division which values are index values belong
+    to the division.
+
+    >>> sorted(_partitions_of_index_values([0, 5, 10], [3]).items())
+    [(0, [3])]
+    >>> sorted(_partitions_of_index_values([0, 5, 10], [3, 8, 5]).items())
+    [(0, [3]), (1, [8, 5])]
+    """
+    if divisions[0] is None:
+        msg = "Can not use loc on DataFrame without known divisions"
+        raise ValueError(msg)
+
+    results = defaultdict(list)
+    values = pd.Index(values, dtype=object)
+    for val in values:
+        i = bisect.bisect_right(divisions, val)
+        div = min(len(divisions) - 2, max(0, i - 1))
+        results[div].append(val)
+    return results
+
+
 def _coerce_loc_index(divisions, o):
     """ Transform values to be comparable against divisions
 


Test patch:
diff --git a/dask/dataframe/tests/test_indexing.py b/dask/dataframe/tests/test_indexing.py
--- a/dask/dataframe/tests/test_indexing.py
+++ b/dask/dataframe/tests/test_indexing.py
@@ -32,11 +32,15 @@ def test_loc():
     assert_eq(d.loc[3:8], full.loc[3:8])
     assert_eq(d.loc[:8], full.loc[:8])
     assert_eq(d.loc[3:], full.loc[3:])
+    assert_eq(d.loc[[5]], full.loc[[5]])
+    assert_eq(d.loc[[3, 4, 1, 8]], full.loc[[3, 4, 1, 8]])
 
     assert_eq(d.a.loc[5], full.a.loc[5:5])
     assert_eq(d.a.loc[3:8], full.a.loc[3:8])
     assert_eq(d.a.loc[:8], full.a.loc[:8])
     assert_eq(d.a.loc[3:], full.a.loc[3:])
+    assert_eq(d.a.loc[[5]], full.a.loc[[5]])
+    assert_eq(d.a.loc[[3, 4, 1, 8]], full.a.loc[[3, 4, 1, 8]])
 
     pytest.raises(KeyError, lambda: d.loc[1000])
     assert_eq(d.loc[1000:], full.loc[1000:])
@@ -100,10 +104,12 @@ def test_loc2d():
     assert_eq(d.loc[3:8, 'a'], full.loc[3:8, 'a'])
     assert_eq(d.loc[:8, 'a'], full.loc[:8, 'a'])
     assert_eq(d.loc[3:, 'a'], full.loc[3:, 'a'])
+    assert_eq(d.loc[[8], 'a'], full.loc[[8], 'a'])
 
     assert_eq(d.loc[3:8, ['a']], full.loc[3:8, ['a']])
     assert_eq(d.loc[:8, ['a']], full.loc[:8, ['a']])
     assert_eq(d.loc[3:, ['a']], full.loc[3:, ['a']])
+    assert_eq(d.loc[[3, 4, 3], ['a']], full.loc[[3, 4, 3], ['a']])
 
     # 3d
     with tm.assertRaises(pd.core.indexing.IndexingError):
@@ -120,6 +126,23 @@ def test_loc2d():
         d.a.loc[d.a % 2 == 0, 3]
 
 
+def test_loc2d_with_known_divisions():
+    df = pd.DataFrame(np.random.randn(20, 5),
+                      index=list('abcdefghijklmnopqrst'),
+                      columns=list('ABCDE'))
+    ddf = dd.from_pandas(df, 3)
+
+    assert_eq(ddf.loc['a', 'A'], df.loc[['a'], 'A'])
+    assert_eq(ddf.loc['a', ['A']], df.loc[['a'], ['A']])
+    assert_eq(ddf.loc['a':'o', 'A'], df.loc['a':'o', 'A'])
+    assert_eq(ddf.loc['a':'o', ['A']], df.loc['a':'o', ['A']])
+    assert_eq(ddf.loc[['n'], ['A']], df.loc[['n'], ['A']])
+    assert_eq(ddf.loc[['a', 'c', 'n'], ['A']], df.loc[['a', 'c', 'n'], ['A']])
+    assert_eq(ddf.loc[['t', 'b'], ['A']], df.loc[['t', 'b'], ['A']])
+    assert_eq(ddf.loc[['r', 'r', 'c', 'g', 'h'], ['A']],
+              df.loc[['r', 'r', 'c', 'g', 'h'], ['A']])
+
+
 def test_loc2d_with_unknown_divisions():
     df = pd.DataFrame(np.random.randn(20, 5),
                       index=list('abcdefghijklmnopqrst'),