def get_columns(X):
    """
    Select columns to work with

    :param X: imported training data
    :return: all columns, numerical columns, categorical columns
    """
    numerical_cols = [col for col in X.columns if X[col].dtype in ['int64', 'float64']]
    categorical_cols = [col for col in X.columns if X[col].dtype == 'object']
    columns = numerical_cols + categorical_cols

    return columns, numerical_cols, categorical_cols