Add some exception cases

databricks · Oct 2, 2019 · 7cd8889 · 7cd8889
1 parent 3ec0a43
commit 7cd8889
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/databricks/koalas/series.py b/databricks/koalas/series.py
@@ -1590,8 +1590,15 @@ def drop(self,
 
             drop_index_scols = []
             for idxes in index:
-                index_scols = [self._internal.index_scols[lvl] == idx
-                               for lvl, idx in enumerate(idxes, level)]
+                try:
+                    index_scols = [self._internal.index_scols[lvl] == idx
+                                   for lvl, idx in enumerate(idxes, level)]
+                except IndexError:
+                    if level == 0:
+                        raise KeyError("Key length ({}) exceeds index depth ({})"
+                                       .format(len(self._internal.index_scols), len(idxes)))
+                    else:
+                        return self
                 drop_index_scols.append(reduce(lambda x, y: x & y, index_scols))
 
             sdf = self._internal.sdf.where(~reduce(lambda x, y: x | y, drop_index_scols))

diff --git a/databricks/koalas/tests/test_series.py b/databricks/koalas/tests/test_series.py
@@ -688,3 +688,7 @@ def test_drop(self):
         msg = "Cannot specify both 'labels' and 'index'"
         with self.assertRaisesRegex(ValueError, msg):
             kser.drop('lama', index='cow')
+        msg = r"'Key length \(2\) exceeds index depth \(3\)'"
+        with self.assertRaisesRegex(KeyError, msg):
+            kser.drop(('lama', 'speed', 'x'))
+        self.assert_eq(kser.drop(('lama', 'speed', 'x'), level=1), kser)