# Buggy code
def fillna(self, value=None, method=None, axis=None, inplace=False, limit=None):
    """Fill NA/NaN values.

    .. note:: the current implementation of 'method' parameter in fillna uses Spark's Window
        without specifying partition specification. This leads to move all data into
        single partition in single machine and could cause serious
        performance degradation. Avoid this method against very large dataset.

    Parameters
    ----------
    value : scalar, dict, Series
        Value to use to fill holes. alternately a dict/Series of values
        specifying which value to use for each column.
        DataFrame is not supported.
    method : {'backfill', 'bfill', 'pad', 'ffill', None}, default None
        Method to use for filling holes in reindexed Series pad / ffill: propagate last valid
        observation forward to next valid backfill / bfill:
        use NEXT valid observation to fill gap
    axis : {0 or `index`}
        1 and `columns` are not supported.
    inplace : boolean, default False
        Fill in place (do not create a new object)
    limit : int, default None
        If method is specified, this is the maximum number of consecutive NaN values to
        forward/backward fill. In other words, if there is a gap with more than this number of
        consecutive NaNs, it will only be partially filled. If method is not specified,
        this is the maximum number of entries along the entire axis where NaNs will be filled.
        Must be greater than 0 if not None

    Returns
    -------
    DataFrame
        DataFrame with NA entries filled.

    Examples
    --------
    >>> df = ks.DataFrame({
    ...     'A': [None, 3, None, None],
    ...     'B': [2, 4, None, 3],
    ...     'C': [None, None, None, 1],
    ...     'D': [0, 1, 5, 4]
    ...     },
    ...     columns=['A', 'B', 'C', 'D'])
    >>> df
         A    B    C  D
    0  NaN  2.0  NaN  0
    1  3.0  4.0  NaN  1
    2  NaN  NaN  NaN  5
    3  NaN  3.0  1.0  4

    Replace all NaN elements with 0s.

    >>> df.fillna(0)
         A    B    C  D
    0  0.0  2.0  0.0  0
    1  3.0  4.0  0.0  1
    2  0.0  0.0  0.0  5
    3  0.0  3.0  1.0  4

    We can also propagate non-null values forward or backward.

    >>> df.fillna(method='ffill')
         A    B    C  D
    0  NaN  2.0  NaN  0
    1  3.0  4.0  NaN  1
    2  3.0  4.0  NaN  5
    3  3.0  3.0  1.0  4

    Replace all NaN elements in column 'A', 'B', 'C', and 'D', with 0, 1,
    2, and 3 respectively.

    >>> values = {'A': 0, 'B': 1, 'C': 2, 'D': 3}
    >>> df.fillna(value=values)
         A    B    C  D
    0  0.0  2.0  2.0  0
    1  3.0  4.0  2.0  1
    2  0.0  1.0  2.0  5
    3  0.0  3.0  1.0  4
    """
    if axis is None:
        axis = 0
    if not (axis == 0 or axis == "index"):
        raise NotImplementedError(
            "fillna currently only works for axis=0 or axis='index'"
        )
    if (value is None) and (method is None):
        raise ValueError("Must specify a fill 'value' or 'method'.")

    sdf = self._sdf
    if value is not None:
        if not isinstance(value, (float, int, str, bool, dict, pd.Series)):
            raise TypeError("Unsupported type %s" % type(value))
        if isinstance(value, pd.Series):
            value = value.to_dict()
        if isinstance(value, dict):
            for v in value.values():
                if not isinstance(v, (float, int, str, bool)):
                    raise TypeError("Unsupported type %s" % type(v))
        if limit is not None:
            raise ValueError("limit parameter for value is not support now")
        sdf = sdf.fillna(value)
    else:
        for data_column in self._internal.data_columns:
            if method in ["ffill", "pad"]:
                func = F.last
                end = Window.currentRow - 1
                if limit is not None:
                    begin = Window.currentRow - limit
                else:
                    begin = Window.unboundedPreceding
            elif method in ["bfill", "backfill"]:
                func = F.first
                begin = Window.currentRow + 1
                if limit is not None:
                    end = Window.currentRow + limit
                else:
                    end = Window.unboundedFollowing
            else:
                raise ValueError("Expecting pad, ffill, backfill or bfill.")
            window = Window.orderBy(self._internal.index_columns).rowsBetween(
                begin, end
            )
            sdf = sdf.withColumn(
                data_column,
                F.when(
                    F.col(data_column).isNull(),
                    func(F.col(data_column), True).over(window),
                ).otherwise(F.col(data_column)),
            )
    internal = self._internal.copy(sdf=sdf)
    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 fillna(self, value=None, method=None, axis=None, inplace=False, limit=None):
    """Fill NA/NaN values.

    .. note:: the current implementation of 'method' parameter in fillna uses Spark's Window
        without specifying partition specification. This leads to move all data into
        single partition in single machine and could cause serious
        performance degradation. Avoid this method against very large dataset.

    Parameters
    ----------
    value : scalar, dict, Series
        Value to use to fill holes. alternately a dict/Series of values
        specifying which value to use for each column.
        DataFrame is not supported.
    method : {'backfill', 'bfill', 'pad', 'ffill', None}, default None
        Method to use for filling holes in reindexed Series pad / ffill: propagate last valid
        observation forward to next valid backfill / bfill:
        use NEXT valid observation to fill gap
    axis : {0 or `index`}
        1 and `columns` are not supported.
    inplace : boolean, default False
        Fill in place (do not create a new object)
    limit : int, default None
        If method is specified, this is the maximum number of consecutive NaN values to
        forward/backward fill. In other words, if there is a gap with more than this number of
        consecutive NaNs, it will only be partially filled. If method is not specified,
        this is the maximum number of entries along the entire axis where NaNs will be filled.
        Must be greater than 0 if not None

    Returns
    -------
    DataFrame
        DataFrame with NA entries filled.

    Examples
    --------
    >>> df = ks.DataFrame({
    ...     'A': [None, 3, None, None],
    ...     'B': [2, 4, None, 3],
    ...     'C': [None, None, None, 1],
    ...     'D': [0, 1, 5, 4]
    ...     },
    ...     columns=['A', 'B', 'C', 'D'])
    >>> df
         A    B    C  D
    0  NaN  2.0  NaN  0
    1  3.0  4.0  NaN  1
    2  NaN  NaN  NaN  5
    3  NaN  3.0  1.0  4

    Replace all NaN elements with 0s.

    >>> df.fillna(0)
         A    B    C  D
    0  0.0  2.0  0.0  0
    1  3.0  4.0  0.0  1
    2  0.0  0.0  0.0  5
    3  0.0  3.0  1.0  4

    We can also propagate non-null values forward or backward.

    >>> df.fillna(method='ffill')
         A    B    C  D
    0  NaN  2.0  NaN  0
    1  3.0  4.0  NaN  1
    2  3.0  4.0  NaN  5
    3  3.0  3.0  1.0  4

    Replace all NaN elements in column 'A', 'B', 'C', and 'D', with 0, 1,
    2, and 3 respectively.

    >>> values = {'A': 0, 'B': 1, 'C': 2, 'D': 3}
    >>> df.fillna(value=values)
         A    B    C  D
    0  0.0  2.0  2.0  0
    1  3.0  4.0  2.0  1
    2  0.0  1.0  2.0  5
    3  0.0  3.0  1.0  4
    """
    if axis is None:
        axis = 0
    if not (axis == 0 or axis == "index"):
        raise NotImplementedError(
            "fillna currently only works for axis=0 or axis='index'"
        )
    if (value is None) and (method is None):
        raise ValueError("Must specify a fill 'value' or 'method'.")

    sdf = self._sdf
    if value is not None:
        if not isinstance(value, (float, int, str, bool, dict, pd.Series)):
            raise TypeError("Unsupported type %s" % type(value))
        if isinstance(value, pd.Series):
            value = value.to_dict()
        if isinstance(value, dict):
            for v in value.values():
                if not isinstance(v, (float, int, str, bool)):
                    raise TypeError("Unsupported type %s" % type(v))
        if limit is not None:
            raise ValueError("limit parameter for value is not support now")
        sdf = sdf.fillna(value)
    else:
        for data_column in self._internal.data_columns:
            if method in ["ffill", "pad"]:
                func = F.last
                end = Window.currentRow - 1
                if limit is not None:
                    begin = Window.currentRow - limit
                else:
                    begin = Window.unboundedPreceding
            elif method in ["bfill", "backfill"]:
                func = F.first
                begin = Window.currentRow + 1
                if limit is not None:
                    end = Window.currentRow + limit
                else:
                    end = Window.unboundedFollowing
            else:
                raise ValueError("Expecting pad, ffill, backfill or bfill.")
            window = Window.orderBy(self._internal.index_scols).rowsBetween(begin, end)
            sdf = sdf.withColumn(
                data_column,
                F.when(
                    scol_for(sdf, data_column).isNull(),
                    func(scol_for(sdf, data_column), True).over(window),
                ).otherwise(scol_for(sdf, data_column)),
            )
    internal = self._internal.copy(sdf=sdf)
    if inplace:
        self._internal = internal
    else:
        return DataFrame(internal)