Remove InternalFrame.index_map. (#1901)

Removes `InternalFrame.index_map` which is not suitable in many cases or can easily be replaced with `zip(column_names, index_names)`.

Author: ueshin

databricks/koalas/frame.py

@@ -3432,11 +3432,17 @@ def rename(index):
                     return ("level_{}".format(index),)
 
         if level is None:
-            new_index_map = [
-                (column, name if name is not None else rename(i))
-                for i, (column, name) in enumerate(self._internal.index_map.items())
+            new_column_labels = [
+                name if name is not None else rename(i)
+                for i, name in enumerate(self._internal.index_names)
             ]
-            index_map = {}  # type: Dict
+            new_data_spark_columns = [
+                scol.alias(name_like_string(label))
+                for scol, label in zip(self._internal.index_spark_columns, new_column_labels)
+            ]
+
+            index_spark_column_names = []
+            index_names = []
         else:
             if is_list_like(level):
                 level = list(level)
@@ -3478,35 +3484,29 @@ def rename(index):
                 raise ValueError("Level should be all int or all string.")
             idx.sort()
 
-            new_index_map = []
-            index_map_items = list(self._internal.index_map.items())
-            new_index_map_items = index_map_items.copy()
-            for i in idx:
-                info = index_map_items[i]
-                index_column, index_name = info
-                new_index_map.append(
-                    (index_column, index_name if index_name is not None else rename(i))
-                )
-                new_index_map_items.remove(info)
+            new_column_labels = []
+            new_data_spark_columns = []
 
-            index_map = OrderedDict(new_index_map_items)
+            index_spark_column_names = self._internal.index_spark_column_names.copy()
+            index_spark_columns = self._internal.index_spark_columns.copy()
+            index_names = self._internal.index_names.copy()
 
-        if drop:
-            new_index_map = []
+            for i in idx[::-1]:
+                index_spark_column_names.pop(i)
 
-        for _, name in new_index_map:
-            if name in self._internal.column_labels:
-                raise ValueError("cannot insert {}, already exists".format(name_like_string(name)))
+                name = index_names.pop(i)
+                new_column_labels.insert(0, name if name is not None else rename(i))
 
-        sdf = self._internal.spark_frame
-        new_data_scols = [
-            scol_for(sdf, column).alias(name_like_string(name)) for column, name in new_index_map
-        ]
+                scol = index_spark_columns.pop(i)
+                new_data_spark_columns.insert(0, scol.alias(name_like_string(name)))
 
-        index_scols = [scol_for(sdf, column) for column in index_map]
-        sdf = sdf.select(
-            index_scols + new_data_scols + self._internal.data_spark_columns + list(HIDDEN_COLUMNS)
-        )
+        if drop:
+            new_data_spark_columns = []
+            new_column_labels = []
+
+        for label in new_column_labels:
+            if label in self._internal.column_labels:
+                raise ValueError("cannot insert {}, already exists".format(name_like_string(label)))
 
         if self._internal.column_labels_level > 1:
             column_depth = len(self._internal.column_labels[0])
@@ -3516,28 +3516,22 @@ def rename(index):
                         column_depth, col_level + 1
                     )
                 )
-            if any(col_level + len(name) > column_depth for _, name in new_index_map):
+            if any(col_level + len(label) > column_depth for label in new_column_labels):
                 raise ValueError("Item must have length equal to number of levels.")
-            column_labels = [
+            new_column_labels = [
                 tuple(
                     ([col_fill] * col_level)
-                    + list(name)
-                    + ([col_fill] * (column_depth - (len(name) + col_level)))
+                    + list(label)
+                    + ([col_fill] * (column_depth - (len(label) + col_level)))
                 )
-                for _, name in new_index_map
-            ] + self._internal.column_labels
-        else:
-            column_labels = [name for _, name in new_index_map] + self._internal.column_labels
+                for label in new_column_labels
+            ]
 
         internal = self._internal.copy(
-            spark_frame=sdf,
-            index_spark_column_names=list(index_map.keys()),
-            index_names=list(index_map.values()),
-            column_labels=column_labels,
-            data_spark_columns=(
-                [scol_for(sdf, name_like_string(name)) for _, name in new_index_map]
-                + [scol_for(sdf, col) for col in self._internal.data_spark_column_names]
-            ),
+            index_spark_column_names=index_spark_column_names,
+            index_names=index_names,
+            column_labels=new_column_labels + self._internal.column_labels,
+            data_spark_columns=new_data_spark_columns + self._internal.data_spark_columns,
         )
 
         if inplace:
@@ -5957,11 +5951,10 @@ def droplevel(self, level, axis=0) -> "DataFrame":
             if not isinstance(level, (tuple, list)):  # huh?
                 level = [level]
 
-            spark_frame = self._internal.spark_frame
-            index_map = self._internal.index_map.copy()
             index_names = self.index.names
-            nlevels = self.index.nlevels
-            int_levels = list()
+            nlevels = self._internal.index_level
+
+            int_level = set()
             for n in level:
                 if isinstance(n, int):
                     if n < 0:
@@ -5981,22 +5974,27 @@ def droplevel(self, level, axis=0) -> "DataFrame":
                     if n not in index_names:
                         raise KeyError("Level {} not found".format(n))
                     n = index_names.index(n)
-                int_levels.append(n)
+                int_level.add(n)
 
-            if len(int_levels) >= nlevels:
+            if len(level) >= nlevels:
                 raise ValueError(
                     "Cannot remove {} levels from an index with {} levels: "
-                    "at least one level must be left.".format(len(int_levels), nlevels)
+                    "at least one level must be left.".format(len(level), nlevels)
                 )
 
-            for int_level in int_levels:
-                index_spark_column = self._internal.index_spark_column_names[int_level]
-                spark_frame = spark_frame.drop(index_spark_column)
-                index_map.pop(index_spark_column)
+            index_spark_column_names, index_names = zip(
+                *[
+                    item
+                    for i, item in enumerate(
+                        zip(self._internal.index_spark_column_names, self._internal.index_names)
+                    )
+                    if i not in int_level
+                ]
+            )
+
             internal = self._internal.copy(
-                spark_frame=spark_frame,
-                index_spark_column_names=list(index_map.keys()),
-                index_names=list(index_map.values()),
+                index_spark_column_names=list(index_spark_column_names),
+                index_names=list(index_names),
             )
             return DataFrame(internal)
         else:
@@ -6845,33 +6843,38 @@ def to_list(os: Optional[Union[Any, List[Any], Tuple, List[Tuple]]]) -> List[Tup
             if right_index:
                 if how in ("inner", "left"):
                     exprs.extend(left_index_scols)
-                    index_map = self._internal.index_map
+                    index_spark_column_names = self._internal.index_spark_column_names
+                    index_names = self._internal.index_names
                 elif how == "right":
                     exprs.extend(right_index_scols)
-                    index_map = right._internal.index_map
+                    index_spark_column_names = right._internal.index_spark_column_names
+                    index_names = right._internal.index_names
                 else:
-                    index_map = OrderedDict()
-                    for (col, name), left_scol, right_scol in zip(
-                        self._internal.index_map.items(), left_index_scols, right_index_scols
+                    index_spark_column_names = self._internal.index_spark_column_names
+                    index_names = self._internal.index_names
+                    for col, left_scol, right_scol in zip(
+                        index_spark_column_names, left_index_scols, right_index_scols
                     ):
                         scol = F.when(left_scol.isNotNull(), left_scol).otherwise(right_scol)
                         exprs.append(scol.alias(col))
-                        index_map[col] = name
             else:
                 exprs.extend(right_index_scols)
-                index_map = right._internal.index_map
+                index_spark_column_names = right._internal.index_spark_column_names
+                index_names = right._internal.index_names
         elif right_index:
             exprs.extend(left_index_scols)
-            index_map = self._internal.index_map
+            index_spark_column_names = self._internal.index_spark_column_names
+            index_names = self._internal.index_names
         else:
-            index_map = OrderedDict()
+            index_spark_column_names = None
+            index_names = None
 
         selected_columns = joined_table.select(*exprs)
 
         internal = InternalFrame(
             spark_frame=selected_columns,
-            index_spark_column_names=list(index_map.keys()) if index_map else None,
-            index_names=list(index_map.values()) if index_map else None,
+            index_spark_column_names=index_spark_column_names,
+            index_names=index_names,
             column_labels=column_labels,
             data_spark_columns=[scol_for(selected_columns, col) for col in data_columns],
         )

databricks/koalas/indexes.py

@@ -1233,14 +1233,25 @@ def droplevel(self, level) -> "Index":
         if not is_list_like(level):
             level = [level]
 
+        int_level = set()
         for n in level:
             if isinstance(n, int):
-                if n > nlevels - 1:
+                if n < 0:
+                    n = n + nlevels
+                    if n < 0:
+                        raise IndexError(
+                            "Too many levels: Index has only {} levels, "
+                            "{} is not a valid level number".format(nlevels, (n - nlevels))
+                        )
+                if n >= nlevels:
                     raise IndexError(
                         "Too many levels: Index has only {} levels, not {}".format(nlevels, n + 1)
                     )
-            elif n not in names:
-                raise KeyError("Level {} not found".format(n))
+            else:
+                if n not in names:
+                    raise KeyError("Level {} not found".format(n))
+                n = names.index(n)
+            int_level.add(n)
 
         if len(level) >= nlevels:
             raise ValueError(
@@ -1249,19 +1260,24 @@ def droplevel(self, level) -> "Index":
                 "left.".format(len(level), nlevels)
             )
 
-        int_level = set(n if isinstance(n, int) else names.index(n) for n in level)
         index_spark_column_names, index_names = zip(
-            *[item for i, item in enumerate(self._internal.index_map.items()) if i not in int_level]
+            *[
+                item
+                for i, item in enumerate(
+                    zip(self._internal.index_spark_column_names, self._internal.index_names)
+                )
+                if i not in int_level
+            ]
         )
 
         sdf = self._internal.spark_frame
         sdf = sdf.select(*index_spark_column_names)
-        result = InternalFrame(
+        internal = InternalFrame(
             spark_frame=sdf,
             index_spark_column_names=list(index_spark_column_names),
             index_names=list(index_names),
         )
-        return DataFrame(result).index
+        return DataFrame(internal).index
 
     def symmetric_difference(self, other, result_name=None, sort=None) -> "Index":
         """
@@ -2632,7 +2648,7 @@ def swaplevel(self, i=-2, j=-1) -> "MultiIndex":
                     "%s is not a valid level number" % (len(self.names), index)
                 )
 
-        index_map = list(self._internal.index_map.items())
+        index_map = list(zip(self._internal.index_spark_column_names, self._internal.index_names))
         index_map[i], index_map[j], = index_map[j], index_map[i]
         index_spark_column_names, index_names = zip(*index_map)
         internal = self._kdf._internal.copy(
@@ -3011,22 +3027,24 @@ def drop(self, codes, level=None) -> "MultiIndex":
         elif isinstance(level, int):
             scol = index_scols[level]
         else:
-            spark_column_name = None
-            for index_spark_column_name, index_name in self._internal.index_map.items():
+            scol = None
+            for index_spark_column, index_name in zip(
+                self._internal.index_spark_columns, self._internal.index_names
+            ):
                 if not isinstance(level, tuple):
                     level = (level,)
                 if level == index_name:
-                    if spark_column_name is not None:
+                    if scol is not None:
                         raise ValueError(
                             "The name {} occurs multiple times, use a level number".format(
                                 name_like_string(level)
                             )
                         )
-                    spark_column_name = index_spark_column_name
-            if spark_column_name is None:
+                    scol = index_spark_column
+            if scol is None:
                 raise KeyError("Level {} not found".format(name_like_string(level)))
-            scol = scol_for(sdf, spark_column_name)
         sdf = sdf[~scol.isin(codes)]
+
         return MultiIndex(
             DataFrame(
                 InternalFrame(
@@ -3220,7 +3238,6 @@ def insert(self, loc: int, item) -> Index:
                 )
 
         index_name = self._internal.index_spark_column_names
-        sdf = self._internal.spark_frame
         sdf_before = self.to_frame(name=index_name)[:loc].to_spark()
         sdf_middle = Index([item]).to_frame(name=index_name).to_spark()
         sdf_after = self.to_frame(name=index_name)[loc:].to_spark()
@@ -3277,8 +3294,6 @@ def intersection(self, other) -> "MultiIndex":
         MultiIndex([('c', 'z')],
                    )
         """
-        keep_name = True
-
         if isinstance(other, Series) or not is_list_like(other):
             raise TypeError("other must be a MultiIndex or a list of tuples")
         elif isinstance(other, DataFrame):