# Buggy code
def transform(self, func, *args, **kwargs):
    """
    Call ``func`` producing the same type as `self` with transformed values
    and that has the same axis length as input.

    .. note:: unlike pandas, it is required for `func` to specify return type hint.

    Parameters
    ----------
    func : function or list
        A function or a list of functions to use for transforming the data.
    *args
        Positional arguments to pass to `func`.
    **kwargs
        Keyword arguments to pass to `func`.

    Returns
    -------
    An instance of the same type with `self` that must have the same length as input.

    See Also
    --------
    Series.apply : Invoke function on Series.

    Examples
    --------

    >>> s = ks.Series(range(3))
    >>> s
    0    0
    1    1
    2    2
    Name: 0, dtype: int64

    >>> def sqrt(x) -> float:
    ...    return np.sqrt(x)
    >>> s.transform(sqrt)
    0    0.000000
    1    1.000000
    2    1.414214
    Name: 0, dtype: float32

    Even though the resulting instance must have the same length as the
    input, it is possible to provide several input functions:

    >>> def exp(x) -> float:
    ...    return np.exp(x)
    >>> s.transform([sqrt, exp])
           sqrt       exp
    0  0.000000  1.000000
    1  1.000000  2.718282
    2  1.414214  7.389056

    """
    if isinstance(func, list):
        applied = []
        for f in func:
            applied.append(self.apply(f).rename(f.__name__))

        sdf = self._kdf._sdf.select(
            self._internal.index_columns + [c._scol for c in applied]
        )

        internal = self.to_dataframe()._internal.copy(
            sdf=sdf, data_columns=[c.name for c in applied]
        )

        return DataFrame(internal)
    else:
        return self.apply(func, args=args, **kwargs)


# 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 transform(self, func, *args, **kwargs):
    """
    Call ``func`` producing the same type as `self` with transformed values
    and that has the same axis length as input.

    .. note:: unlike pandas, it is required for `func` to specify return type hint.

    Parameters
    ----------
    func : function or list
        A function or a list of functions to use for transforming the data.
    *args
        Positional arguments to pass to `func`.
    **kwargs
        Keyword arguments to pass to `func`.

    Returns
    -------
    An instance of the same type with `self` that must have the same length as input.

    See Also
    --------
    Series.apply : Invoke function on Series.

    Examples
    --------

    >>> s = ks.Series(range(3))
    >>> s
    0    0
    1    1
    2    2
    Name: 0, dtype: int64

    >>> def sqrt(x) -> float:
    ...    return np.sqrt(x)
    >>> s.transform(sqrt)
    0    0.000000
    1    1.000000
    2    1.414214
    Name: 0, dtype: float32

    Even though the resulting instance must have the same length as the
    input, it is possible to provide several input functions:

    >>> def exp(x) -> float:
    ...    return np.exp(x)
    >>> s.transform([sqrt, exp])
           sqrt       exp
    0  0.000000  1.000000
    1  1.000000  2.718282
    2  1.414214  7.389056

    """
    if isinstance(func, list):
        applied = []
        for f in func:
            applied.append(self.apply(f).rename(f.__name__))

        sdf = self._kdf._sdf.select(
            self._internal.index_scols + [c._scol for c in applied]
        )

        internal = self.to_dataframe()._internal.copy(
            sdf=sdf, data_columns=[c.name for c in applied]
        )

        return DataFrame(internal)
    else:
        return self.apply(func, args=args, **kwargs)