From 7cd88898b90b4391a251226def69b95e7b196842 Mon Sep 17 00:00:00 2001 From: itholic Date: Wed, 2 Oct 2019 14:16:14 +0900 Subject: [PATCH] Add some exception cases --- databricks/koalas/series.py | 11 +++++++++-- databricks/koalas/tests/test_series.py | 4 ++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/databricks/koalas/series.py b/databricks/koalas/series.py index cb898efd8c..86bb594e00 100644 --- 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 index 333a73e420..f612c5c7c4 100644 --- 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)