Bugfix/81 get missing column

sparsity/sparse_frame.py

@@ -806,7 +806,8 @@ def __getitem__(self, item):
             #  where it is used with Multiindex
             item = [item]
         if len(item) > 0:
-            return self.reindex_axis(item, axis=1)
+            indexer = self.loc._convert_to_indexer(item, axis=1)
+            return self._take(indexer, axis=1)
         else:
             data = np.empty(shape=(self.shape[0], 0))
             return SparseFrame(data, index=self.index,

sparsity/test/test_dask_sparse_frame.py

@@ -63,16 +63,27 @@ def test_todense_series():
     pdt.assert_series_equal(computed, data[0], check_dtype=False)
 
 
-@pytest.mark.parametrize('item', [
-    'X',
-    ['X', 'Y'],
+# noinspection PyStatementEffect
+@pytest.mark.parametrize('item, raises', [
+    ('X', False),
+    (['X', 'Y'], False),
+    ('A', True),
+    (['A'], True),
+    (['X', 'A'], True),
+    (['A', 'B'], True),
 ])
-def test_getitem(item):
+def test_getitem(item, raises):
     df = pd.DataFrame(np.random.rand(10, 3), columns=list('XYZ'),
                       index=list('ABCDEFGHIJ'))
     dsf = dsp.from_pandas(df, npartitions=2)
 
     correct_cols = item if isinstance(item, list) else [item]
+
+    if raises:
+        with pytest.raises(KeyError):
+            dsf[item]
+        return
+
     res = dsf[item]
     assert res.columns.tolist() == correct_cols
     res_computed = res.compute()

sparsity/test/test_sparse_frame.py

@@ -579,6 +579,7 @@ def test_npz_io_s3(complex_example):
         assert np.all(loaded.columns == sf.columns)
 
 
+# noinspection PyStatementEffect
 def test_getitem():
     id_ = np.identity(10)
     sf = SparseFrame(id_, columns=list('abcdefghij'))
@@ -614,6 +615,35 @@ def test_getitem_empty():
     assert sf[['a', 'b']].columns.tolist() == ['a', 'b']
 
 
+# noinspection PyStatementEffect
+def test_getitem_missing_col():
+    id_ = np.identity(10)
+    sf = SparseFrame(id_, columns=list('abcdefghij'))
+
+    with pytest.raises(ValueError):
+        sf[None]
+    with pytest.raises(KeyError):
+        sf['x']
+    with pytest.raises(KeyError):
+        sf[['x']]
+    with pytest.raises(KeyError):
+        sf[['a', 'x']]
+    with pytest.raises(KeyError):
+        sf[['y', 'x']]
+
+    idx = pd.Index(list('abx'))
+    with pytest.raises(KeyError):
+        sf[idx]
+    with pytest.raises(KeyError):
+        sf[idx.to_series()]
+    with pytest.raises(KeyError):
+        sf[idx.tolist()]
+    with pytest.raises(KeyError):
+        sf[tuple(idx)]
+    with pytest.raises(KeyError):
+        sf[idx.values]
+
+
 def test_vstack():
     frames = []
     data = []