Usar mssql-python com Apache Arrow

O driver mssql-python fornece métodos de busca Apache Arrow para recuperação de dados colunares de alto desempenho a partir do Microsoft SQL e do Banco de Dados SQL do Azure.

Apache Arrow é uma plataforma de desenvolvimento multi-linguagens para dados colunares em memória. O driver converte conjuntos de resultados ODBC diretamente para o formato Arrow em C++, ignorando a criação de objetos em Python para melhorar o desempenho.

A integração com seta possibilita:

• Transferência de dados sem cópia para Polars, pandas e DuckDB. "Zero-copy" significa que os dados permanecem em um único buffer na memória, que o driver grava e as bibliotecas consumidoras leem diretamente, de modo que nenhuma linha é duplicada para objetos Python intermediários.
• Transmitir conjuntos de resultados via RecordBatchReader sem carregar tudo na memória.
• Formato de dados colunar ideal para cargas de análise e aprendizado de máquina.
• Uso de memória reduzido em comparação à criação de objetos Python linha por linha.

Métodos de cursor

O pyarrow pacote é obrigado a usar métodos de busca Arrow. Instale-o com pip install pyarrow. Se pyarrow não estiver instalado, chamar qualquer método Arrow gera um ImportError.

O driver mssql-python adiciona três métodos ao objeto cursor para acesso a dados do Arrow. Os três métodos convertem conjuntos de resultados ODBC para o formato Arrow na camada C++ do driver, o que evita criar objetos Python intermediários.

• arrow() retorna todo o conjunto de resultados como uma única tabela em memória. O mais simples de usar.
• arrow_batch() retorna um lote de linhas por vez, permitindo controle manual do loop.
• arrow_reader() retorna um iterador que gera lotes automaticamente. Ideal para transmitir grandes volumes de resultados.

Usando cursor.arrow(batch_size=8192)

Obtenha todo o conjunto de resultados como um único pyarrow.Table. Esse método é o mais simples e funciona bem quando o conjunto completo de resultados cabe na memória.

import mssql_python

conn = mssql_python.connect(connection_string)
cursor = conn.cursor()

cursor.execute("SELECT ProductID, Name, ListPrice FROM Production.Product")
table = cursor.arrow()

print(type(table))       # <class 'pyarrow.lib.Table'>
print(table.num_rows)    # Number of rows fetched
print(table.num_columns) # Number of columns
print(table.schema)      # Column names and Arrow types
print(table.to_pandas()) # Convert to pandas DataFrame

Usando cursor.arrow_batch(batch_size=8192)

Busque um único pyarrow.RecordBatch contendo até batch_size linhas. Use esse método para ciclos personalizados de processamento em lote, onde você precisa de controle detalhado sobre quantas linhas são buscadas ao mesmo tempo.

cursor.execute("SELECT * FROM Production.TransactionHistory")

while True:
    batch = cursor.arrow_batch(batch_size=10000)
    if batch.num_rows == 0:
        break
    # Process each batch
    print(f"Fetched {batch.num_rows} rows")

Usando cursor.arrow_reader(batch_size=8192)

Retorna um leitor que fornece objetos RecordBatch até que o conjunto de resultados se esgote. Esse método é a opção mais eficiente em memória para grandes conjuntos de resultados.

cursor.execute("SELECT * FROM Production.TransactionHistory")
reader = cursor.arrow_reader(batch_size=50000)

for batch in reader:
    # Process streaming batches without loading all data
    print(f"Batch: {batch.num_rows} rows")

O leitor transmite resultados pela conexão; assim, enquanto um leitor não lido estiver aberto, essa conexão não poderá iniciar outra instrução. Ao tentar um, ocorre uma falha com o erro Connection is busy with results for another command.

Três coisas liberam o leitor: iterar até o final, fechar o cursor pai ou fechar o leitor. Se você parar de ler antes que o conjunto de resultados se esgote e continuar usando o cursor, feche o leitor. Ao fechá-lo, o cursor pai também é redefinido, então você pode executar outra instrução nele.

Use o leitor como um gerenciador de contexto para que ele se feche mesmo que uma exceção interrompa o ciclo:

cursor.execute("SELECT * FROM Production.TransactionHistory")

rows_seen = 0
with cursor.arrow_reader(batch_size=50000) as reader:
    for batch in reader:
        rows_seen += batch.num_rows
        if rows_seen >= 100000:
            break

# The reader is closed here, and the cursor is ready for the next statement.
cursor.execute("SELECT COUNT(*) FROM Production.TransactionHistory")

Você também pode ligar reader.close() diretamente. Ligar mais de uma vez é seguro, e a reader.closed propriedade informa se você a fechou.

Padrões comuns

As tabelas de seta se integram diretamente com bibliotecas de dados populares em Python. Os exemplos a seguir mostram como passar dados do Arrow para pandas, Polars, DuckDB e formatos de arquivo sem copiar dados.

Carregue os resultados no pandas

cursor.execute("SELECT * FROM Production.Product")
table = cursor.arrow()

# Convert to pandas with zero-copy where possible
df = table.to_pandas()
print(df.head())

Carregar resultados no Polars

import polars as pl

cursor.execute("SELECT * FROM Production.Product")
table = cursor.arrow()

df = pl.from_arrow(table)
print(df)

Resultados da consulta com o DuckDB

O DuckDB pode consultar tabelas de seta diretamente em SQL sem copiar dados. Essa funcionalidade é útil quando você precisa de análise no estilo SQL em conjuntos de resultados que já estão no formato Arrow.

import duckdb

cursor.execute("SELECT * FROM Sales.SalesOrderHeader")
arrow_table = cursor.arrow()

# Query the Arrow table with DuckDB SQL
result = duckdb.sql("SELECT CustomerID, SUM(TotalDue) FROM arrow_table GROUP BY CustomerID")
print(result.fetchall())

Faça streaming de grandes conjuntos de resultados em Parquet

Para conjuntos de resultados grandes, transmita lotes do Arrow diretamente para um arquivo Parquet sem carregar todo o conjunto de dados na memória. O ParquetWriter grava cada lote incrementalmente.

import pyarrow.parquet as pq

cursor.execute("SELECT * FROM Production.TransactionHistory")
reader = cursor.arrow_reader(batch_size=100000)

# Write streaming batches to a Parquet file
writer = None
for batch in reader:
    if writer is None:
        writer = pq.ParquetWriter("output.parquet", batch.schema)
    writer.write_batch(batch)

if writer:
    writer.close()

Exportação para outros formatos

O PyArrow fornece gravadores integrados para CSV e o formato de arquivo IPC Arrow (também conhecido como Feather V2). Arquivos Arrow IPC preservam exatamente os tipos Arrow e são rápidos de ler novamente.

import pyarrow as pa
import pyarrow.csv as pcsv

cursor.execute("SELECT * FROM Production.Product")
table = cursor.arrow()

# Write to CSV
pcsv.write_csv(table, "products.csv")

# Write to an Arrow IPC file
with pa.ipc.new_file("products.arrow", table.schema) as writer:
    writer.write_table(table)

Carregar os dados do Arrow no SQL Server

O cursor.bulkcopy_arrow() método grava dados Arrow em uma tabela sem convertê-los para tuplas de linha em Python primeiro. O source argumento aceita qualquer um dos seguintes pontos:

• Um pyarrow.Table.

• Um pyarrow.RecordBatch.

• Um pyarrow.RecordBatchReader, incluindo o leitor retornado por cursor.arrow_reader().

• Qualquer objeto que exponha a interface de dados C do Arrow por meio de __arrow_c_stream__ ou __arrow_c_array__.

import mssql_python
import pyarrow as pa

conn = mssql_python.connect(connection_string)

# bulkcopy_arrow() opens its own connection, so commit the table creation first.
conn.autocommit = True
cursor = conn.cursor()

cursor.execute("""
    CREATE TABLE ##SensorArchive (
        SensorID int NOT NULL,
        Reading float NULL,
        Location nvarchar(50) NULL
    )
""")

table = pa.table({
    "SensorID": pa.array([1, 2, 3], type=pa.int32()),
    "Reading": pa.array([20.5, None, 22.1], type=pa.float64()),
    "Location": pa.array(["Plant A", "Plant B", None], type=pa.string()),
})

result = cursor.bulkcopy_arrow("##SensorArchive", table)
print(f"Copied {result['rows_copied']} rows in {result['batch_count']} batches")

Valores nulos de setas são escritos como valores NULL em SQL.

Transmita um conjunto de resultados para outra tabela

Como bulkcopy_arrow() aceita um leitor, você pode mover um grande conjunto de resultados entre tabelas sem materializá-lo na memória:

cursor.execute("""
    CREATE TABLE ##ProductArchive (
        ProductID int NOT NULL,
        Name nvarchar(50) NOT NULL,
        ListPrice money NOT NULL
    )
""")

cursor.execute("SELECT ProductID, Name, ListPrice FROM Production.Product")

with cursor.arrow_reader(batch_size=100000) as reader:
    result = cursor.bulkcopy_arrow("##ProductArchive", reader, batch_size=100000)

print(f"Copied {result['rows_copied']} rows")

Corresponder os tipos de seta às colunas de destino

O Arrow writer exige que cada tipo de coluna Arrow seja compatível com seu tipo de coluna SQL de destino. Ele não converte entre famílias, então um erro de incompatibilidade é gerado ValueError antes que qualquer linha seja gravada:

ValueError: Cannot map Arrow column 'ListPrice' (Float64) to SQL column 'ListPrice'
(Money): Usage Error: type combination is not supported by the Arrow row-major writer

Use os mapeamentos em Mapeamentos de tipos de dados na ordem inversa para escolher o tipo Seta. Colunas monetárias, decimais e numéricas precisam de decimal128, não float64. Os dados lidos de volta com cursor.arrow() já vêm com os tipos corretos, portanto uma tabela lida do SQL Server é carregada em uma tabela correspondente sem conversão.

Mapear colunas por nome

Quando a ordem das colunas Arrow não corresponder à da tabela de destino, passe column_mappings com os nomes das colunas de destino na ordem das colunas Arrow:

from decimal import Decimal

table = pa.table({
    "Name": pa.array(["Widget"], type=pa.string()),
    "ProductID": pa.array([9001], type=pa.int32()),
    "ListPrice": pa.array([Decimal("12.34")], type=pa.decimal128(19, 4)),
})

cursor.bulkcopy_arrow(
    "##ProductArchive",
    table,
    column_mappings=["Name", "ProductID", "ListPrice"],
)

O método aceita as mesmas opções que cursor.bulkcopy(), incluindo batch_size, timeout, keep_identity, table_lock, e keep_nulls. Para mais informações sobre essas opções, veja Cópia em massa.

Mapeamentos de tipo de dados

Os métodos de busca do Arrow mapeiam os tipos SQL da Microsoft para tipos do Arrow no nível de C++.

Considerações sobre desempenho

Métodos de busca por seta são os mais rápidos para análises e operações de dados em massa, enquanto métodos padrão de cursor são mais adequados para padrões transacionais com conjuntos de resultados pequenos.

Quando usar Arrow em vez do fetch padrão

Gerenciamento de memória para grandes conjuntos de dados

Para conjuntos de resultados que possam exceder a memória disponível, use arrow_reader() com um batch_size razoável.

cursor.execute("SELECT * FROM Production.TransactionHistory")

# Process in batches of 100K rows
reader = cursor.arrow_reader(batch_size=100000)
total_rows = 0

for batch in reader:
    # Work with each batch individually
    total_rows += batch.num_rows
    # batch goes out of scope and memory is freed

print(f"Processed {total_rows} rows")

Ajustar tamanho do lote

O batch_size parâmetro controla quantas linhas são buscadas em cada lote. O tamanho ideal depende da largura da sua linha e da memória disponível. Linhas mais largas com colunas grandes como nvarchar(max) ou varbinary(max) se beneficiam de tamanhos de lote menores, enquanto fileiras estreitas se beneficiam de tamanhos maiores.

• Padrão (8192): Bom equilíbrio para a maioria das cargas de trabalho.
• Menor (1000-5000): Use para tabelas amplas com colunas largas.
• Maior (50000-100000): Use para tabelas estreitas ou quando a taxa de transferência importa mais do que a memória.

# Narrow table with many rows - use larger batches
cursor.execute("SELECT ProductID, ListPrice FROM Production.Product")
table = cursor.arrow(batch_size=100000)

# Wide table with LOB columns - use smaller batches
cursor.execute("SELECT * FROM Production.Document")
table = cursor.arrow(batch_size=1000)

Conteúdo relacionado

• Usando mssql-python com Polars
• Use mssql-python com pandas
• Recuperar dados com mssql-python
• Mapeamentos de tipos de dados para mssql-python
• Ajuste de desempenho para aplicações mssql-python