Testando Comando COPY no Postgres 16 e Inserindo 1 Milhão de Dados em 13s

Introdução

Eu estava vendo algumas notícias e vi algo bem absurdo de atualização. Basicamente que o Postgres 16 tinha atualizado e que a performance em tempo do comando COPY melhorou em 300%.

Fiquei bem interessado, pois eu costumo usar ele no trabalho. No caso, tem vezes que preciso inserir em torno de 100.000 ou até 1.000.000 de dados. Se eu fosse um único comando SQL com vários INSERT iria demorar muito tempo, então eu aprendi a usar o COPY que performa muito mais para inserir dados.

Caso não conheça o que é o COPY, ele é um comando que o Postgres tem para copiar dados de um arquivo para uma tabela específica. Você passa a tabela, algumas configurações de formato do arquivo, colunas, delimitadores e depois executa o comando. Aqui a documentação oficial dele COPY Postgres.

Eu estou deixando aqui o Notebook com meu teste caso tenha interesse, mas vou descrever tudo que fiz dentro dele.

Configuração do Ambiente

Bom, primeiro eu configurei o postgres para executar no Google Colab. Eu achei essa configuração nesse aqui: Leveraging Google Colab to run Postgres: A Comprehensive Guide. Admito que não entendi muito bem o que ele está fazendo, mas algo que melhorei foi adicionar esse comando final aqui > /dev/null 2>&1 para não aparecer tantos logs no notebook.

!sudo apt update > /dev/null 2>&1
!sudo apt install dirmngr ca-certificates software-properties-common gnupg gnupg2 apt-transport-https curl -y  > /dev/null 2>&1
!curl -fSsL https://www.postgresql.org/media/keys/ACCC4CF8.asc | gpg --dearmor | sudo tee /usr/share/keyrings/postgresql.gpg > /dev/null
!echo 'deb [arch=amd64,arm64,ppc64el signed-by=/usr/share/keyrings/postgresql.gpg] http://apt.postgresql.org/pub/repos/apt/ jammy-pgdg main' | sudo tee /etc/apt/sources.list.d/pgdg.list
!sudo apt update  > /dev/null 2>&1
!sudo apt install postgresql-client-16 postgresql-16 -y  > /dev/null 2>&1

Depois eu parei o servidor e reiniciei. Fiz isso para caso eu execute várias vezes esse notebook eu ter como parar ele se eu quiser.

!sudo service postgresql stop
!sudo service postgresql start

Essa parte foi algo que fiz para eliminar alguas regras que crie no processo. Eu elimino minhas roles e users antes de iniciar. Isso foi bastante útil enquanto eu estava construindo esse notebook.

!sudo -u postgres psql -U postgres -c "DROP DATABASE IF EXISTS test_copy";
!sudo -u postgres psql -U postgres -c "DROP ROLE IF EXISTS root;"
!sudo -u postgres psql -U postgres -c "DROP ROLE IF EXISTS myuser;"

Depois eu criei as configurações principais, que são o usuário root e o usuário myuser com senha supersecret123.

# Setup a password `myuser` for username `supersecret123`

!sudo -u postgres psql -U postgres -c "CREATE ROLE root WITH SUPERUSER;"
!sudo -u postgres psql -U postgres -c "ALTER ROLE root WITH LOGIN;"
!sudo -u postgres psql -U postgres -c "CREATE ROLE myuser WITH PASSWORD 'supersecret123';"
!sudo -u postgres psql -c "ALTER ROLE myuser WITH LOGIN;"

Em seguida, criei o banco de dados test_copy e adicionar o myuser para ter o máximo de permissões possíveis para executar as operações necessárias. A permissão mais importante é a pg_read_server_files, pois ela que cuida da parte do usuário poder execuar o comando COPY.

!sudo -u postgres psql -c "CREATE DATABASE test_copy;"
!sudo -u postgres psql -d test_copy -c "GRANT ALL PRIVILEGES ON SCHEMA public TO myuser;"
!sudo -u postgres psql -c "GRANT ALL PRIVILEGES ON DATABASE test_copy TO myuser;"
!sudo -u postgres psql -c "GRANT pg_read_server_files TO myuser;"

Fiz o restart de tudo só para não ter problemas de atualização de permissões depois.

!service postgresql restart

A partir daqui já é a parte do SQLAlchemy. Instalei ele junto com o driver psycopg2 para poder fazer os comandos SQL.

!pip install psycopg2-binary sqlalchemy -q

import psycopg2
from sqlalchemy import create_engine, text

Coloquei em destaque a connection string, pois levei um tempo para notar que só precisava colocar o host como localhost, mas deixo abaixo caso queira modificar para outro formato.

# postgresql+psycopg2://<user>:<password>@<host>:<port>/<database>
connection_string = 'postgresql+psycopg2://myuser:supersecret123@localhost:5432/test_copy'

Depois eu criei essa função execute_query. Ela basicamente faz o seguinte, eu passo uma query, que pode ou não mudar algo no banco de dados. Se for para alterar, é necessário passar op=write, caso contrário op=read. Isso me ajudou bastante na reulização do meu código, mas não acho que é uma boa implementação, pois faço muitas restrições. Contudo para o que preciso é o suficiente.

from typing import Literal, List, Optional

def execute_query(query: str, op: Literal['read', 'write']) -> Optional[List[tuple]]:
  if(not op in ['read', 'write']):
    raise ValueError('op must be *read* or *write*')

  try:
    engine = create_engine(connection_string)

    with engine.connect() as connection:
      result = connection.execute(text(query))

      if op == 'read':
        return result.fetchall()
      elif op == 'write':
        connection.commit()
  except Exception as e:
    print(e)
  finally:
    engine.dispose()
    return None

Minha primeira operação no banco de dados test_copy é criar uma tabela users. Essa tabela é apenas fictícia, pois quero inserir o máximo de informação possível nela. Escolhi os campos abaixo.

create_table_query = '''
CREATE TABLE users (
  id SERIAL PRIMARY KEY,
  username VARCHAR(50) NOT NULL,
  email VARCHAR(100) NOT NULL,
  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
  age INTEGER NOT NULL,
  is_active BOOLEAN DEFAULT TRUE,
  bio TEXT,
  last_login TIMESTAMP
);
'''
execute_query(create_table_query, 'write')

Depois dei mais permissões para meu usuário, por segurança.

!sudo -u postgres psql -d test_copy -c "GRANT ALL PRIVILEGES ON TABLE users TO myuser;"

Usei essa lib chamada Faker. Ela gera dados randomicos de de várias formas. Apenas escolhi a coluna e gerei os dados.

!pip install faker -q

Espera bastante aqui. Se quiser testar um caso menor, mude n para 1000.

import csv
from faker import Faker
from tqdm import tqdm

csv_file_path = "fake_users.csv"
n = 1000000

fake = Faker()

with open(csv_file_path, mode='w', newline='', encoding='utf-8') as file:
  writer = csv.writer(file)

  writer.writerow(["id", "username", "email", "created_at", "age", "is_active", "bio", "last_login"])

  for i in tqdm(range(n)):
    writer.writerow([
        i + 1,
        fake.user_name(),
        fake.email(),
        fake.date_time_this_year().isoformat(),
        fake.random_int(min=18, max=80),
        fake.boolean(),
        fake.text(max_nb_chars=200),
        fake.date_time_this_year().isoformat()
    ])

Antes de fazer a query, eu deleto os dados anteriores. Acho importante, pois não quero acumular nada na tabela users

execute_query('delete from users', 'write')

Nessa parte eu utilizei o urllib para reutilizar o que eu tinha na minha string de conexão.

import urllib.parse as up

url = up.urlparse(connection_string)
dbname = url.path[1:]
user = url.username
password = url.password
host = url.hostname
port = url.port

Nessa parte eu fiz a operação que eu estava precisando. Se notar, tem uma variável chamada sql onde passo algumas configurações.

Não consegui reutlizar minha função anterior nessa parte, mas achai uma discussão bem interessante no Stack Overflow sobre como fazer o COPY com python. Vou deixar ela aqui: copy data from csv to postgresql using python.

from time import time

start_time = None
end_time = None

with psycopg2.connect(dbname=dbname, user=user, password=password, host=host, port=port) as conn:
    with conn.cursor() as cursor:
        with open(csv_file_path, 'r') as f:
            sql = "COPY users FROM STDIN DELIMITER ',' CSV HEADER"
            try:
                start_time = time()
                cursor.copy_expert(sql, f)
                end_time = time()
                conn.commit()
            except psycopg2.Error as e:
                print("Error executing COPY command:", e)

if(start_time and end_time):
  print(f"Tempo de execução: {end_time - start_time} segundos")

No meu caso foi bem rápido. Fiz um teste para 1.000.000 de dados e ele inseriu os dados em torno de 13s! Demorou mais tempo para gerar os dados do que para inserir, basicamente!

Caso queira vizualizar, basta usar o pandas:

import pandas as pd
with psycopg2.connect(dbname=dbname, user=user, password=password, host=host, port=port) as conn:
  df = pd.read_sql_query("SELECT * FROM users", conn)
df.head()

Conclusão

É bem impressionante como o COPY melhorou. 13s é uma performance muito boa. Mas lembre-se que o COPY é bom para inserir dados, apenas. Se for consultar dados ou fazer outras operações, prinpalmente para essa quantidade, o ideal é colocar filtros no SELECT para não vir tudo de uma vez.

A transformação para dataframe ajuda bastante também nas consultas. Se o dataframa estiver carregado no notebook as consultas não demoram tanto. Só leva um tempo para inserir no dataframe em si.