Como obter os cabeçalhos da coluna Pandas DataFrame como uma lista

Pandas é um pacote open-source para análise de dados em Python. pandas.DataFrame é a estrutura de dados primária de Pandas. É uma estrutura de dados tabular bidimensional com eixos (linhas e colunas) rotulados.

Um caso de uso generalizado é obter uma lista de cabeçalhos de coluna de um objeto DataFrame.

Vamos reutilizar o objeto DataFrame, que definimos abaixo, em todos os outros exemplos de código deste tutorial.

>>> import pandas
>>> cities = {
...   'name': ['New York', 'Los Angeles', 'Chicago'],
...   'population': [8601186, 4057841, 2679044],
...   'state': ['NY', 'CA', 'IL'],
... }
>>> data_frame = pandas.DataFrame(cities)

Uma maneira de obter uma lista de nomes de colunas DataFrame é iterar sobre um objeto DataFrame em si. O iterador DataFrame retorna os nomes das colunas na ordem de definição.

>>> for column in data_frame:
...   print(column)
...
name
population
state

Quando há a necessidade de converter um iterável em uma lista, você pode chamar a função list embutida do Python.

>>> list(data_frame)
['name', 'population', 'state']

No entanto, o desempenho deste método é lento.

>>> from timeit import timeit
>>> timeit(lambda: list(data_frame))
7.818843764999997

Podemos também atravessar mais profundamente em um objeto DataFrame para acessar suas colunas a partir de uma propriedade DataFrame.columns.

>>> list(data_frame.columns)
['name', 'population', 'state']

Caso contrário, podemos utilizar a função DataFrame.columns.tolist() para alcançar a mesma coisa.

>>> data_frame.columns.tolist()
['name', 'population', 'state']

A performance de ambos os métodos não é muito melhor.

>>> timeit(lambda: list(data_frame.columns))
7.143133517000024
>>> timeit(lambda: data_frame.columns.tolist())
6.064925153999866

As coisas mudam muito quando atravessamos ainda mais a propriedade DataFrame.columns.values. Similarmente, como com o objeto DataFrame e a propriedade DataFrame.columns, podemos utilizá-la para obter uma seqüência de nomes de colunas DataFrame.

>>> list(data_frame.columns.values)
['name', 'population', 'state']

A performance desta abordagem é 5 a 6 vezes melhor quando comparada com os métodos anteriores.

>>> timeit(lambda: list(data_frame.columns.values))
1.301724927000123

Ainda assim, o melhor tempo de execução pode ser alcançado se utilizarmos o método DataFrame.columns.values.tolist() embutido.

>>> data_frame.columns.values.tolist()
['name', 'population', 'state']
>>> timeit(lambda: data_frame.columns.values.tolist())
0.6860591469999235

Como podemos ver, a performance desta abordagem é mais de dez vezes melhor do que se tivéssemos iterado diretamente sobre o objeto DataFrame. A maioria dos engenheiros estarão curiosos sobre as razões por trás de tal discrepância na performance.

A resposta se esconde no tipo de dados da propriedade DataFrame.columns.values. É um array NumPy. NumPy é um pacote Python para computação científica, e os mantenedores o otimizam bastante para a performance.

Pandas é construído em cima do NumPy e fornece abstrações convenientes de alto nível. Assim, executar operações diretas em estruturas de dados NumPy de nível inferior será quase sempre mais rápido do que executar operações similares em estruturas de dados Pandas de nível superior.

Artigo relacionado - Pandas DataFrame

  • Remover Colunas por índice em Pandas DataFrame
  • Como converter Pandas Dataframe em Numpy Array
  • Artigo relacionado - Pandas DataFrame Column

  • Criar Pandas Dataframe a partir de uma lista
  • Obter dummies em Pandas