loc of bad element raises KeyError

mrocklin · Jul 20, 2015 · 7ba29bd · 7ba29bd
1 parent 319e1e8
commit 7ba29bd
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
@@ -228,6 +228,8 @@ def _loc_series(self, ind):
     def _loc_element(self, ind):
         name = 'loc-element' + next(tokens)
         part = _partition_of_index_value(self.divisions, ind)
+        if ind < self.divisions[0] or ind > self.divisions[-1]:
+            raise KeyError('the label [%s] is not in the index' % str(ind))
         dsk = {(name, 0): (lambda df: df.loc[ind], (self._name, part))}
         return type(self)(merge(self.dask, dsk), name,
                           self.column_info, [ind, ind])

diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
@@ -324,6 +324,10 @@ def test_loc():
     assert eq(d.loc[:8], full.loc[:8])
     assert eq(d.loc[3:], full.loc[3:])
 
+    assert raises(KeyError, lambda: d.loc[1000])
+    assert eq(d.loc[1000:], full.loc[1000:])
+    assert eq(d.loc[-2000:-1000], full.loc[-2000:-1000])
+
 
 
 def test_loc_with_text_dates():