Utilizar o mssql-python com o 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 Base de Dados SQL do Azure.

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

A integração com Arrow permite:

  • Transferência de dados sem cópia para Polars, pandas e DuckDB. "Zero-copy" significa que os dados permanecem num único buffer de memória, escrito pelo driver e lido diretamente pelas bibliotecas consumidoras, pelo que nenhuma linha é duplicada em objetos Python intermédios.
  • Transmitir conjuntos de resultados em fluxo através de RecordBatchReader sem carregar tudo na memória.
  • Formato de dados colunar ideal para cargas de trabalho de análise e aprendizagem automática.
  • Redução do uso de memória em comparação com a criação de objetos Python linha a linha.

Métodos de cursor

O pyarrow pacote é obrigatório para 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 o acesso aos 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 intermédios.

  • arrow() devolve todo o conjunto de resultados como uma tabela em memória. Mais simples de utilizar.
  • arrow_batch() devolve um lote de linhas de cada vez, o que lhe dá controlo manual sobre o ciclo.
  • arrow_reader() devolve um iterador que gera lotes automaticamente. Ideal para a transmissão em fluxo de grandes volumes de resultados.

Ao utilizar cursor.arrow(batch_size=8192)

Obtenha todo o conjunto de resultados num único pyarrow.Table. Este 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

Note

Se a sua cadeia de ligação usar Authentication=ActiveDirectoryDefault, o controlador usa DefaultAzureCredential, que tenta vários fornecedores de credenciais em sequência. A primeira conexão pode ser lenta porque o SDK percorre a cadeia até encontrar um provedor funcional. Em produção, se souber que tipo de credencial o seu ambiente utiliza, especifique-o diretamente (por exemplo, ActiveDirectoryMSI para identidade gerida) para evitar o chain walk. Para obter mais informações, consulte Autenticação do Microsoft Entra.

Ao utilizar cursor.arrow_batch(batch_size=8192)

Obtém um único pyarrow.RecordBatch com até batch_size linhas. Use este método para ciclos personalizados de processamento em lote, onde precisa de controlo detalhado sobre quantas linhas são obtidas de cada vez.

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")

Ao utilizar cursor.arrow_reader(batch_size=8192)

Devolve um leitor que devolve objetos RecordBatch até esgotar o conjunto de resultados. Este método é a opção mais eficiente em termos de memória para conjuntos de resultados grandes.

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 os resultados através da ligação, por isso, enquanto um leitor não lido estiver aberto, essa ligação não pode iniciar outra afirmação. Ao tentar um, ocorre um erro Connection is busy with results for another command.

Há três formas de libertar o leitor: iterá-lo até ao fim, fechar o cursor pai ou fechar o leitor. Se parar de ler antes de esgotar o conjunto de resultados e continuar a usar o cursor, feche o leitor. Fechar também reinicia o cursor pai, por isso podes executar outra instrução nele.

Use o leitor como gestor 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")

Também pode ligar reader.close() diretamente. Ligar mais do que uma vez é seguro, e a reader.closed propriedade informa se foi tu que a fechou.

Padrões comuns

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

Carregar resultados para o 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 para o 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 Arrow diretamente em SQL sem copiar dados. Esta funcionalidade é útil quando precisas de análise ao estilo SQL em conjuntos de resultados que já estão em 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())

Transmitir grandes conjuntos de resultados para Parquet

Para conjuntos de resultados grandes, transmita lotes do Arrow diretamente para um ficheiro Parquet sem carregar todo o conjunto de dados na memória. O ParquetWriter escreve 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 incorporados para CSV e o formato de ficheiro IPC Arrow (também conhecido como Feather V2). Os ficheiros Arrow IPC preservam exatamente os tipos de Arrow e são rápidos de ler.

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 dados do Arrow no SQL Server

O método cursor.bulkcopy_arrow() escreve dados Arrow numa tabela sem primeiro os converter em tuplos de linhas em Python. O source argumento aceita qualquer um dos seguintes:

  • Um pyarrow.Table.

  • Um pyarrow.RecordBatch.

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

  • Qualquer objeto que exponha a interface de dados Arrow C através de __arrow_c_stream__ ou de __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")

Os valores nulos das setas são escritos como valores SQL NULL.

Transmitir um conjunto de resultados para outra tabela

Como bulkcopy_arrow() aceita um leitor, pode mover um grande conjunto de resultados entre tabelas sem o materializar 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")

Associar tipos de seta às colunas de destino

O escritor Arrow exige que cada tipo de coluna Arrow seja compatível com o seu tipo de coluna SQL de destino. Não converte entre famílias, pelo que uma incompatibilidade gera ValueError antes de qualquer linha ser escrita:

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

Utilize os mapeamentos em Mapeamentos de tipos de dados na ordem inversa para escolher o tipo Arrow. As colunas monetárias, decimais e numéricas precisam de decimal128, não de float64. Os dados lidos com cursor.arrow() já têm os tipos corretos, pelo que uma tabela lida a partir do SQL Server é carregada para uma tabela correspondente sem conversão.

Colunas do mapa por nome

Quando a ordem das colunas Arrow não corresponde à tabela de destino, passe column_mappings com os nomes das colunas de destino pela 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, consulte Cópia em bloco.

Note

Passar uma fonte de Flecha para cursor.bulkcopy() eleva TypeError e direciona-te para cursor.bulkcopy_arrow().

Mapeamentos de tipo de dados

Os métodos de obtenção do Arrow mapeiam os tipos SQL da Microsoft para tipos Arrow ao nível de C++.

Tipo Microsoft SQL Tipo de flecha
int, smallint, tinyint, bigint int32, int16, int8, int64
float, real float64, float32
decimal, numérico decimal128
bit bool
char, varchar, nchar, nvarchar utf8
texto, ntext large_utf8
binário, varbinário binary, large_binary
date date32
time time64[us]
datetime, datetime2, smalldatetime timestamp[us]
datetimeoffset timestamp[us, tz=UTC]
uniqueidentifier utf8 (corda maiúscula)
xml utf8

Note

O driver converte o datetimeoffset tipo para UTC porque as colunas Arrow requerem um fuso horário fixo. O driver normaliza a informação por fuso horário por célula do Microsoft SQL para UTC durante a conversão.

O sql_variant tipo não é suportado pelos métodos Arrow fetch e gera uma exceção de tipo de dado não suportado. Use o padrão fetchone(), fetchmany(), ou fetchall() para consultas que devolvam sql_variant colunas.

Considerações sobre desempenho

Os métodos de busca por seta são os mais rápidos para análises e operações de dados em massa, enquanto os 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

Scenario Abordagem recomendada
Obtém algumas linhas para apresentação fetchone() / fetchall()
Carregar dados para pandas ou Polars cursor.arrow()
Processar grandes conjuntos de dados em blocos cursor.arrow_reader()
Consultas de linha única ou pequenos conjuntos de resultados fetchone() / fetchval()
Cadeias de análise ou de agregação cursor.arrow() + Polars/DuckDB
Escrever resultados em Parquet ou Arrow IPC cursor.arrow_reader() + PyArrow I/O

Gestão de memória para grandes conjuntos de dados

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

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 o tamanho do lote

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

  • Predefinição (8192): Bom equilíbrio para a maioria das cargas de trabalho.
  • Mais pequeno (1000-5000): Utilize para tabelas largas com colunas largas.
  • Maior (50000-100000): Uso para tabelas estreitas ou quando o throughput 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

  • Last updated on