# Buggy code
def set_index(self, keys, drop=True, append=False, inplace=False):
    """Set the DataFrame index (row labels) using one or more existing columns.

    Set the DataFrame index (row labels) using one or more existing
    columns or arrays (of the correct length). The index can replace the
    existing index or expand on it.

    Parameters
    ----------
    keys : label or array-like or list of labels/arrays
        This parameter can be either a single column key, a single array of
        the same length as the calling DataFrame, or a list containing an
        arbitrary combination of column keys and arrays. Here, "array"
        encompasses :class:`Series`, :class:`Index` and ``np.ndarray``.
    drop : bool, default True
        Delete columns to be used as the new index.
    append : bool, default False
        Whether to append columns to existing index.
    inplace : bool, default False
        Modify the DataFrame in place (do not create a new object).

    Returns
    -------
    DataFrame
        Changed row labels.

    See Also
    --------
    DataFrame.reset_index : Opposite of set_index.

    Examples
    --------
    >>> df = ks.DataFrame({'month': [1, 4, 7, 10],
    ...                    'year': [2012, 2014, 2013, 2014],
    ...                    'sale': [55, 40, 84, 31]},
    ...                   columns=['month', 'year', 'sale'])
    >>> df
       month  year  sale
    0      1  2012    55
    1      4  2014    40
    2      7  2013    84
    3     10  2014    31

    Set the index to become the 'month' column:

    >>> df.set_index('month')  # doctest: +NORMALIZE_WHITESPACE
           year  sale
    month
    1      2012    55
    4      2014    40
    7      2013    84
    10     2014    31

    Create a MultiIndex using columns 'year' and 'month':

    >>> df.set_index(['year', 'month'])  # doctest: +NORMALIZE_WHITESPACE
                sale
    year  month
    2012  1     55
    2014  4     40
    2013  7     84
    2014  10    31
    """
    if isinstance(keys, str):
        keys = [keys]
    else:
        keys = list(keys)
    for key in keys:
        if key not in self.columns:
            raise KeyError(key)

    if drop:
        data_columns = [
            column for column in self._internal.data_columns if column not in keys
        ]
    else:
        data_columns = self._internal.data_columns
    if append:
        index_map = self._internal.index_map + [(column, column) for column in keys]
    else:
        index_map = [(column, column) for column in keys]

    index_columns = set(column for column, _ in index_map)
    columns = [column for column, _ in index_map] + [
        column for column in data_columns if column not in index_columns
    ]

    # Sync Spark's columns as well.
    sdf = self._sdf.select(["`{}`".format(name) for name in columns])

    internal = _InternalFrame(sdf=sdf, index_map=index_map, data_columns=data_columns)

    if inplace:
        self._internal = internal
    else:
        return DataFrame(internal)


# Traceback (KeyError)
df = ks.read_parquet("/mnt/ivan/sample.parquet")
df['a.b']

# KeyError: 'a.b'
# ---------------------------------------------------------------------------
# KeyError                                  Traceback (most recent call last)
# <command-3188040896828876> in <module>()
#       1 df = ks.read_parquet("/mnt/ivan/sample.parquet")
# ----> 2 df['a.b']
#
# /databricks/python/lib/python3.5/site-packages/databricks/koalas/frame.py in __getitem__(self, key)
#     448
#     449     def __getitem__(self, key):
# --> 450         return self._pd_getitem(key)
#     451
#     452     def __setitem__(self, key, value):
#
# /databricks/python/lib/python3.5/site-packages/databricks/koalas/frame.py in _pd_getitem(self, key)
#     427                 return Series(self._sdf.__getitem__(key), self, self._metadata.index_info)
#     428             except AnalysisException:
# --> 429                 raise KeyError(key)
#     430         if np.isscalar(key) or isinstance(key, (tuple, string_types)):
#     431             raise NotImplementedError(key)
#
# KeyError: 'a.b'

# Fixed code
def set_index(self, keys, drop=True, append=False, inplace=False):
    """Set the DataFrame index (row labels) using one or more existing columns.

    Set the DataFrame index (row labels) using one or more existing
    columns or arrays (of the correct length). The index can replace the
    existing index or expand on it.

    Parameters
    ----------
    keys : label or array-like or list of labels/arrays
        This parameter can be either a single column key, a single array of
        the same length as the calling DataFrame, or a list containing an
        arbitrary combination of column keys and arrays. Here, "array"
        encompasses :class:`Series`, :class:`Index` and ``np.ndarray``.
    drop : bool, default True
        Delete columns to be used as the new index.
    append : bool, default False
        Whether to append columns to existing index.
    inplace : bool, default False
        Modify the DataFrame in place (do not create a new object).

    Returns
    -------
    DataFrame
        Changed row labels.

    See Also
    --------
    DataFrame.reset_index : Opposite of set_index.

    Examples
    --------
    >>> df = ks.DataFrame({'month': [1, 4, 7, 10],
    ...                    'year': [2012, 2014, 2013, 2014],
    ...                    'sale': [55, 40, 84, 31]},
    ...                   columns=['month', 'year', 'sale'])
    >>> df
       month  year  sale
    0      1  2012    55
    1      4  2014    40
    2      7  2013    84
    3     10  2014    31

    Set the index to become the 'month' column:

    >>> df.set_index('month')  # doctest: +NORMALIZE_WHITESPACE
           year  sale
    month
    1      2012    55
    4      2014    40
    7      2013    84
    10     2014    31

    Create a MultiIndex using columns 'year' and 'month':

    >>> df.set_index(['year', 'month'])  # doctest: +NORMALIZE_WHITESPACE
                sale
    year  month
    2012  1     55
    2014  4     40
    2013  7     84
    2014  10    31
    """
    if isinstance(keys, str):
        keys = [keys]
    else:
        keys = list(keys)
    for key in keys:
        if key not in self.columns:
            raise KeyError(key)

    if drop:
        data_columns = [
            column for column in self._internal.data_columns if column not in keys
        ]
    else:
        data_columns = self._internal.data_columns
    if append:
        index_map = self._internal.index_map + [(column, column) for column in keys]
    else:
        index_map = [(column, column) for column in keys]

    index_columns = set(column for column, _ in index_map)
    columns = [column for column, _ in index_map] + [
        column for column in data_columns if column not in index_columns
    ]

    # Sync Spark's columns as well.
    sdf = self._sdf.select([self._internal.scol_for(name) for name in columns])

    internal = _InternalFrame(sdf=sdf, index_map=index_map, data_columns=data_columns)

    if inplace:
        self._internal = internal
    else:
        return DataFrame(internal)