Como usar embeddings de texto de forma eficiente com Parquet e Polars

A melhor forma de usar embeddings de texto de maneira portátil é com Parquet e Polars

• Embeddings de texto são vetores gerados por grandes modelos de linguagem, uma forma de representar numericamente palavras, frases e documentos
• Em fevereiro de 2025, foram gerados embeddings para um total de 32.254 cartas de "Magic: The Gathering"
• Isso permite analisar matematicamente a similaridade com base no design e nas propriedades mecânicas das cartas
• Também é possível visualizar os embeddings gerados em 2D por meio de redução de dimensionalidade com UMAP
• O modelo de embedding usado foi gte-modernbert-base, e o processo detalhado está organizado no repositório no GitHub
• Esse conjunto de dados de embeddings está disponível no Hugging Face

Repensando a necessidade de bancos de dados vetoriais

• Em geral, bancos de dados vetoriais (faiss, qdrant, Pinecone) são usados para armazenar e buscar embeddings
• Porém, bancos de dados vetoriais exigem configuração complexa, e serviços em nuvem podem ter custo elevado
• Para dados em pequena escala (na casa das dezenas de milhares), é possível fazer buscas rápidas de similaridade com numpy mesmo sem um banco de dados vetorial
• Usando a operação de dot product do numpy, é possível calcular similaridade cosseno de forma simples, com média de 1,08 ms para 32.254 embeddings

def fast_dot_product(query, matrix, k=3):  
    dot_products = query @ matrix.T  
  
    idx = np.argpartition(dot_products, -k)[-k:]  
    idx = idx[np.argsort(dot_products[idx])[::-1]]  
  
    score = dot_products[idx]  
  
    return idx, score  

• Ao usar um banco de dados vetorial, há grande chance de ficar preso a bibliotecas e serviços específicos
• Se os embeddings forem gerados em um servidor com GPU e depois baixados localmente, é preciso um método eficiente de armazenamento e transferência de dados

As piores formas de armazenar embeddings

• Arquivo CSV
  • Ao armazenar dados de ponto flutuante (float32) como texto, o tamanho aumenta mais de 6 vezes
  • Até no tutorial oficial da OpenAI, o uso de CSV é recomendado apenas para conjuntos de dados pequenos
  • Ao salvar com .savetxt() do numpy, o tamanho do arquivo sobe para 631,5 MB
• Arquivo pickle
  • Pode ser salvo e carregado rapidamente, mas traz riscos de segurança e baixa compatibilidade entre versões
  • O tamanho do arquivo é 94,49 MB, igual ao tamanho original em memória, mas com baixa portabilidade

Formas aceitáveis, mas não ideais, de armazenamento

• Formato .npy do numpy
  • Com a configuração allow_pickle=False, é possível impedir o uso de pickle no armazenamento
  • O tamanho do arquivo e a velocidade são os mesmos do método com pickle, mas é difícil armazenar metadados individuais junto
• Problemas de uma estrutura de armazenamento separada dos metadados
  • Ao salvar como array numpy (.npy), as informações das cartas (nome, texto etc.) ficam separadas dos embeddings
  • Quando os dados mudam (adição/remoção), fica difícil manter a correspondência entre metadados e embeddings
  • Em bancos de dados vetoriais, metadados e vetores são armazenados juntos e há recursos de filtragem

A melhor forma de armazenar embeddings: Parquet + polars

Introdução ao formato de arquivo Parquet

• Apache Parquet é um formato de armazenamento de dados em colunas, no qual é possível definir claramente o tipo de cada coluna
• Como ele pode armazenar dados em formato de lista (arrays float32), é adequado para guardar embeddings
• Oferece desempenho de leitura e gravação superior ao CSV, e permite carregar seletivamente apenas parte dos dados
• Também oferece compressão, embora embeddings tenham baixa redundância e, portanto, pouco ganho de compressão

Uso de arquivos Parquet em Python

• Salvando e carregando arquivos Parquet com pandas:

  df = pd.read_parquet("mtg-embeddings.parquet", columns=["name", "embedding"])  
  df  
  

  • O pandas não lida de forma eficiente com dados aninhados (listas), convertendo-os em objetos numpy object
  • Ao converter para array numpy, é necessária uma operação adicional (np.vstack()), o que pode degradar o desempenho
• Salvando e carregando arquivos Parquet com polars:

  df = pl.read_parquet("mtg-embeddings.parquet", columns=["name", "embedding"])  
  df  
  

  • O polars preserva os arrays float32 como estão, e ao chamar to_numpy() já é possível obter imediatamente um array numpy 2D
  • Com a configuração allow_copy=False, é possível evitar cópias desnecessárias de dados

  embeddings = df["embedding"].to_numpy(allow_copy=False)  
  

• Ao adicionar novos embeddings, também é fácil salvar tudo apenas incluindo uma nova coluna

  df = df.with_columns(embedding=embeddings)  
  df.write_parquet("mtg-embeddings.parquet")  
  

Busca por similaridade e filtragem com Parquet + polars

• É possível primeiro filtrar apenas os dados que atendem a certas condições e depois executar a busca por similaridade
• Exemplo: encontrar cartas semelhantes a uma carta específica (query_embed), mas buscando apenas cartas do tipo 'Sorcery' e com a cor 'Black' incluída

  df_filter = df.filter(  
      pl.col("type").str.contains("Sorcery"),  
      pl.col("manaCost").str.contains("B"),  
  )  
  
  embeddings_filter = df_filter["embedding"].to_numpy(allow_copy=False)  
  idx, _ = fast_dot_product(query_embed, embeddings_filter, k=4)  
  related_cards = df_filter[idx]  
  

• O tempo médio de execução é de 1,48 ms, 37% mais lento que a busca no conjunto completo, mas ainda assim rápido

Alternativas para processar grandes volumes de dados vetoriais

• O método com Parquet e dot product consegue lidar tranquilamente com centenas de milhares de embeddings
• Para conjuntos de dados ainda maiores, pode ser necessário usar um banco de dados vetorial
• Como alternativa, é possível usar o sqlite-vec, baseado em SQLite, para realizar buscas e filtragens vetoriais adicionais

Conclusão

• Bancos de dados vetoriais não são obrigatórios
• A combinação Parquet + polars é uma alternativa poderosa para armazenar, buscar e filtrar embeddings com eficiência
• Especialmente em projetos menores, usar arquivos Parquet pode ser mais rápido e mais econômico
• Dependendo do projeto, é importante escolher a solução adequada entre Parquet e um banco de dados vetorial
• O código e os dados podem ser consultados no repositório no GitHub

[GN⁺]

Comentários do Hacker News

• O problema do Parquet é que ele é estático. Não é adequado quando há necessidade de gravações e atualizações contínuas. Ainda assim, tive bons resultados usando arquivos Parquet com DuckDB e armazenamento de objetos. O tempo de carregamento é rápido

  • Se você hospedar seu próprio modelo de embeddings, pode transmitir arrays compactados numpy float32 como bytes e depois decodificá-los novamente em arrays do numpy
  • Pessoalmente, prefiro usar SQLite com a extensão usearch. Uso vetores binários e depois reordeno os 100 melhores em float32. Isso leva cerca de 2 ms para aproximadamente 20.000 itens, o que é mais rápido que o LanceDB. Em coleções maiores, o Lance pode vencer. Mas, no meu caso de uso, funciona bem porque cada usuário tem um arquivo SQLite dedicado
  • Para portabilidade, existe o Litestream

• Artigo realmente excelente. Gosto do seu trabalho há muito tempo. Para quem está se aprofundando em implementações com SQLite, vale acrescentar que o DuckDB começou a incluir alguns recursos de similaridade vetorial para leitura de Parquet e lida perfeitamente com esse caso de uso

• Ainda não gosto de dataframes, mas o Polars é muito melhor que o pandas

  • Eu estava fazendo cálculos de séries temporais, basicamente alguns ajustes simples em preços de ações
  • Fiquei surpreso com o fato de que o código realmente era legível e testável
  • A execução foi tão rápida que parecia estar quebrada

• Dê uma olhada no usearch da Unum. Ele ganha de qualquer coisa e é muito fácil de usar. Faz exatamente o que você precisa

• Se quiser experimentar, dá para carregar de forma preguiçosa a partir do HF e aplicar filtros

  • O Polars é excelente de usar e eu o recomendo fortemente. Ele é ótimo para saturar a CPU em um único nó e, se você precisar distribuir o trabalho, pode aplicar POLARS_MAX_THREADS ao Ray Actor e ajustá-lo de acordo com o nível de saturação de um único nó

• Há muitas descobertas excelentes

  • Fico me perguntando se é melhor passar dados estruturados para uma API de embeddings ou se é melhor passar dados não estruturados. Se você perguntar ao ChatGPT, ele diz que é melhor enviar dados não estruturados
  • Meu caso de uso é para jsonresume. Estou enviando a versão JSON completa como string para gerar embeddings, mas também estou experimentando usar um modelo que primeiro traduz o resume.json para uma versão de texto completo antes de gerar os embeddings. Os resultados parecem melhores, mas não vi opiniões concretas sobre isso
  • A razão pela qual dados não estruturados podem ser melhores é que eles incluem significado textual/semântico por estarem em linguagem natural

• Há um truque elegante na documentação do Vespa que converte vetores em binário e depois usa uma representação hexadecimal

  • Esse truque pode ser usado para reduzir o tamanho do payload. O Vespa oferece suporte a esse formato e isso é especialmente útil quando o mesmo vetor é referenciado várias vezes em um documento. Em casos como ColBERT ou ColPaLi (quando há vários vetores de embedding), é possível reduzir bastante o tamanho dos vetores armazenados em disco

• Polars + Parquet é excelente em portabilidade e desempenho. Este post focou na portabilidade em Python, mas o Polars tem uma API Rust fácil de usar que permite embutir o motor em vários lugares

• Sou um grande fã do Polars, mas não tinha pensado em usá-lo para armazenar embeddings (eu estava experimentando com sqlite-vec). Parece uma ideia realmente interessante

• Outra biblioteca que eu recomendaria é o lancedb, por seu ótimo desempenho e recursos como indexação de texto completo e versionamento de alterações

  • É um banco de dados vetorial e mais complexo, mas pode ser usado sem criar índices e também oferece excelente suporte Arrow com cópia zero para polars e pandas