# Buggy code
def to_numeric_df(kdf: "ks.DataFrame") -> Tuple[pyspark.sql.DataFrame, List[str]]:
    """
    Takes a dataframe and turns it into a dataframe containing a single numerical
    vector of doubles. This dataframe has a single field called '_1'.

    TODO: index is not preserved currently
    :param kdf: the koalas dataframe.
    :return: a pair of dataframe, list of strings (the name of the columns
             that were converted to numerical types)

    >>> to_numeric_df(ks.DataFrame({'A': [0, 1], 'B': [1, 0], 'C': ['x', 'y']}))
    (DataFrame[_correlation_output: vector], ['A', 'B'])
    """
    # TODO, it should be more robust.
    accepted_types = {
        np.dtype(dt)
        for dt in [
            np.int8,
            np.int16,
            np.int32,
            np.int64,
            np.float32,
            np.float64,
            np.bool_,
        ]
    }
    numeric_fields = [
        fname
        for fname in kdf._internal.data_columns
        if kdf[fname].dtype in accepted_types
    ]
    numeric_df = kdf._sdf.select(*numeric_fields)
    va = VectorAssembler(inputCols=numeric_fields, outputCol=CORRELATION_OUTPUT_COLUMN)
    v = va.transform(numeric_df).select(CORRELATION_OUTPUT_COLUMN)
    return v, numeric_fields


# 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 to_numeric_df(kdf: "ks.DataFrame") -> Tuple[pyspark.sql.DataFrame, List[str]]:
    """
    Takes a dataframe and turns it into a dataframe containing a single numerical
    vector of doubles. This dataframe has a single field called '_1'.

    TODO: index is not preserved currently
    :param kdf: the koalas dataframe.
    :return: a pair of dataframe, list of strings (the name of the columns
             that were converted to numerical types)

    >>> to_numeric_df(ks.DataFrame({'A': [0, 1], 'B': [1, 0], 'C': ['x', 'y']}))
    (DataFrame[_correlation_output: vector], ['A', 'B'])
    """
    # TODO, it should be more robust.
    accepted_types = {
        np.dtype(dt)
        for dt in [
            np.int8,
            np.int16,
            np.int32,
            np.int64,
            np.float32,
            np.float64,
            np.bool_,
        ]
    }
    numeric_fields = [
        fname
        for fname in kdf._internal.data_columns
        if kdf[fname].dtype in accepted_types
    ]
    numeric_df = kdf._sdf.select(
        *[kdf._internal.scol_for(fname) for fname in numeric_fields]
    )
    va = VectorAssembler(inputCols=numeric_fields, outputCol=CORRELATION_OUTPUT_COLUMN)
    v = va.transform(numeric_df).select(CORRELATION_OUTPUT_COLUMN)
    return v, numeric_fields