Como checar qual(is) coluna(s) é(são) chave primária em um DataFrame do pandas?

Conjunto de chaves em linha sobre tecido escuro | Foto de regularguy.eth na Unsplash

Sabe quando você precisa saber se existem valores duplicados em uma ou mais colunas, para usá-la(s) como chave primária em um banco de dados? O pandas simplifica nossa vida nessa tarefa, e poucas pessoas sabem desse detalhe.

df = pd.DataFrame({'mes': [1, 4, 7, 10],
...                'ano': [2012, 2014, 2013, 2014],
...                'venda': [55, 40, 84, 31]})
>>> df
    mes   ano    venda
0      1  2012    55
1      4  2014    40
2      7  2013    84
3     10  2014    31

Imagine que você tenha um DataFrame semelhante ao acima, e que ele tenha milhões de registros (o exemplo é meramente didático). Eu poderia configurar a coluna “ano” como chave primária em um banco de dados?

No exemplo é fácil ver que não, pois o ano de 2014 aparece duas vezes (linhas 1 e 3). Mas e se eu tivesse milhões de registros? Vamos lá!

Caminho n° 1

Aqui, vamos tratar a forma que a maioria das pessoas sabe: checar por valores duplicados.

>>> df["ano"].duplicated().sum()
1

Ou seja, existe um valor na coluna ano que é duplicado, mas este método não nos informa qual valor está duplicado. Precisamos de um passo adicional para saber qual item é duplicado na nossa base.

>>> df[df["ano"].duplicated()]
  mes venda ano
3 10  31    2014

Aqui ficou claro que o registro na linha 3, ano 2014, é duplicado.

Caso eu queira verificar mais de uma coluna, basta adicioná-las ao filtro, lembrando que para duas ou mais colunas, devemos usar [[ ]] .

>>> df[["ano", "mes"]].duplicated().sum()
0

Não temos nenhum valor duplicado para a dupla de valores ano e mês. O ano 2014, que era duplicado quando usamos somente a coluna ano, deixou de ser duplicado quando adicionamos a coluna mes. Cada ano de 2014, antes duplicado, agora não o é, porque possui meses diferentes:

• ano 2014 mes 4
• ano 2014 mes 10

Conclusão: a dupla de colunas de ano e mês pode ser usada como chave primária composta em um banco de dados relacional.

Caminho n° 2

Usando o mesmo exemplo podemos configurar as colunas como índices dos nossos DataFrames, desta forma:

>>> df.set_index("ano")
    mes venda
ano
2012 1 55
2014 4 40
2013 5 84
2014 10 31

Percebam que agora o ano já não é mais uma coluna de nosso DataFrame. Ele representa os índices das linhas. Tal transformação facilita nossa vida quando temos que buscar por um ano específico, ou trabalhar com um intervalo de anos. Mas vamos voltar para o foco: chaves primárias.

Nesta transformação, nada aconteceu. Por quê? Se olharmos a documentação do método set_index, vamos encontrar um parâmetro chamado verify_integrity, que vem configurado como False, por padrão. O que esse parâmetro faz? Justamente o nosso objetivo: verificar se as colunas podem ser usadas como chave primária.

verify_integrity bool, default False

Check the new index for duplicates. Otherwise defer the check until necessary. Setting to False will improve the performance of this method.

Se configurarmos o parâmetro para True , ele passa a gerar um erro quando existem itens duplicados no índice.

>>> df.set_index("ano", verify_integrity=True)
Traceback (most recent call last):
 File “<stdin>”, line 1, in <module>
 File “C:\ProgramData\Anaconda3\lib\site-packages\pandas\core\frame.py”, line 2840, in set_index
 raise ValueError(‘Index has duplicate keys: %s’ % duplicates)
ValueError: Index has duplicate keys: [2014]

O erro aqui deixa bastante claro de que há duplicadas, e nos informa, inclusive, qual valor está duplicado!

Para testar mais de uma coluna, assim como fizemos anteriormente, podemos passar uma lista como argumento do set_index:

>>> df.set_index(["ano", "mes"], verify_integrity=True)
            venda
year month
2012 1        55
2014 4        40
2013 5        84
2014 10       31

Nenhum erro! Conclusão: o par de colunas ano e mês pode ser usado como chave primária composta em um banco de dados relacional.

Resumo

Vimos duas alternativas para verificar valores duplicados em colunas, usando dois métodos: duplicated() e set_index().

O primeiro apresenta um uso comum maior, porém requer um passo a mais para saber qual valor é duplicado, informando, em um primeiro momento, somente se existem ou não valores duplicados.

O segundo método já é mais direto e gera um erro, informando qual valor é duplicado, caso exista.