Add variable 'level_param' to mimic pandas exception exactly

databricks · Oct 2, 2019 · 91160ea · 91160ea
1 parent 7cd8889
commit 91160ea
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/databricks/koalas/series.py b/databricks/koalas/series.py
@@ -1570,6 +1570,7 @@ def drop(self,
             if not isinstance(index, (str, tuple, list)):
                 raise ValueError("'index' type should be one of str, list, tuple")
             if level is None:
+                level_param = False
                 level = 0
             if level >= len(self._internal.index_scols):
                 raise ValueError("'level' should be less than the number of indexes")
@@ -1594,7 +1595,7 @@ def drop(self,
                     index_scols = [self._internal.index_scols[lvl] == idx
                                    for lvl, idx in enumerate(idxes, level)]
                 except IndexError:
-                    if level == 0:
+                    if not level_param:
                         raise KeyError("Key length ({}) exceeds index depth ({})"
                                        .format(len(self._internal.index_scols), len(idxes)))
                     else: