使用 mssql-python 配合 Apache Arrow

mssql-python 驱动提供 Apache Arrow 提取方法，用于从 Microsoft SQL 和 Azure SQL 数据库中实现高性能列式数据检索。

Apache Arrow 是一个跨语言的内存列式数据开发平台。 驱动程序将 ODBC 结果集直接转换为 C++ 的 Arrow 格式，绕过 Python 对象创建以提升性能。

Arrow 集成可实现：

• 零拷贝数据传输至 Polars、pandas 和 DuckDB。 “零复制”是指数据始终保存在单个内存缓冲区中，由驱动程序写入，使用这些数据的库可直接从中读取，因此无需将各行数据复制到中间的 Python 对象中。
• 通过 RecordBatchReader 流式传输结果集，而无需将所有结果加载到内存中。
• 列式数据格式非常适合分析和机器学习工作负载。
• 相比逐行 Python 对象创建，内存占用更低。

游标方法

若要使用 Arrow 获取方法，需要 pyarrow 包。 使用 pip install pyarrow 安装它。 如果 pyarrow 未安装，调用任意箭头方法都会触发一个 ImportError。

mssql-python 驱动在光标对象中添加了三种用于 Arrow 数据访问的方法。 这三种方法都会将 ODBC 结果集转换为驱动的 C++ 层中的 Arrow 格式，避免创建中间的 Python 对象。

• arrow() 将整个结果集返回为一个内存表。 用法最为简单。
• arrow_batch() 一次返回一批行，使你能够手动控制循环。
• arrow_reader() 返回一个自动生成批次的迭代器。 最适合流式传输大型结果集。

使用 cursor.arrow(batch_size=8192)

将整个结果集作为单个 pyarrow.Table 获取。 这种方法最简单，当整个结果集能放入内存时效果良好。

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

注释

如果你的连接字符串使用 Authentication=ActiveDirectoryDefault，驱动程序会使用 DefaultAzureCredential，该机制会按顺序尝试多个凭据提供程序。 第一次连接可能很慢，因为SDK会在链路上走动，直到找到可用的提供者。 在生产环境中，如果你知道环境使用的是哪种凭据类型，可以直接指定它（例如，对于托管标识可指定 ActiveDirectoryMSI），以避免遍历凭据链。 有关详细信息，请参阅 Microsoft Entra 身份验证。

使用 cursor.arrow_batch(batch_size=8192)

获取一个最多包含 batch_size 行的 pyarrow.RecordBatch 在需要精细控制每次提取行数的自定义批处理循环中，请使用此方法。

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

使用 cursor.arrow_reader(batch_size=8192)

返回一个读取器，该读取器会持续返回 RecordBatch 对象，直到结果集耗尽。 对于大型结果集，这种方法是内存效率最高的选择。

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

读取器会通过该连接以流式方式返回结果，因此，当某个读取器尚未读完且仍处于打开状态时，该连接无法执行另一条语句。 尝试其中一个会因 Connection is busy with results for another command 错误而失败。

有三种情况会释放该读取器：将其迭代至末尾、关闭父游标，或关闭读取器本身。 如果在结果集耗尽之前停止读取并继续使用游标，请关闭读取器。 关闭它也会重置父光标，这样你可以在上面运行另一个语句。

将读取器用作上下文管理器，这样即使异常中断循环，它也会被关闭：

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

你也可以直接打电话 reader.close() 。 多次调用它是安全的，并且 reader.closed 属性会报告你是否已将其关闭。

常见模式

Arrow 表可直接集成流行的 Python 数据库。 以下示例展示了如何在不复制数据的情况下，将 Arrow 数据传递给 pandas、Polars、DuckDB 和文件格式。

将结果加载到 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())

将结果加载到 Polars 中

import polars as pl

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

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

DuckDB 查询结果

DuckDB 可以直接在 SQL 中查询 Arrow 表，无需复制数据。 当您需要对已经是Arrow格式的结果集进行SQL式分析时，这项功能非常有用。

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

将大型结果集流式传输到 Parquet

对于大型结果集，可以直接将 Arrow 批处理流到一个 Parquet 文件，而无需将整个数据集加载到内存中。 ParquetWriter 会采用增量方式写入每一批。

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

导出到其他格式

PyArrow 内置了 CSV 和 Arrow IPC 文件格式（也称为 Feather V2）的写入工具。 Arrow IPC 文件准确保存 Arrow 类型，且读取速度快。

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)

将Arrow数据加载到SQL Server

该cursor.bulkcopy_arrow()方法可直接将 Arrow 数据写入表中，而无需先将其转换为 Python 行元组。 source 参数接受以下任一值：

• 一个 pyarrow.Table。

• 一个 pyarrow.RecordBatch。

• 一个 pyarrow.RecordBatchReader，包括由 cursor.arrow_reader() 返回的读取器。

• 任何通过 __arrow_c_stream__ 或 __arrow_c_array__ 公开 Arrow C 数据接口的对象。

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

箭头空值被写成 SQL 的 NULL 值。

将结果集流式传输到另一个表

由于 bulkcopy_arrow() 接受读者，你可以在不同表之间移动一个大型结果集，而无需在内存中实现它：

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

将箭头类型与目标列匹配

Arrow 编写器要求每个 Arrow 列类型与其目标 SQL 列类型兼容。 它不会在不同类别之间进行转换，因此在写入任何行之前，不匹配就会引发 ValueError：

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

反向使用 数据类型映射 中的映射来选择 Arrow 类型。 货币、 十进制和 数字 列需要 decimal128，而非 float64。 使用 cursor.arrow() 读回的数据已带有正确的数据类型，因此，从 SQL Server 读取的表可以直接加载到对应的表中，而无需进行转换。

按名称排列的地图列

当箭头列顺序与目标表不匹配时，将目标列名按箭头列顺序传递 column_mappings ：

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

该方法接受与 cursor.bulkcopy()相同的选项，包括 batch_size、 timeout、 keep_identitytable_lockkeep_nulls、 和 。 有关这些选项的更多信息，请参见 批量副本。

注释

传递箭源到 cursor.bulkcopy() ，会升 TypeError 起并引导你前往 cursor.bulkcopy_arrow()。

数据类型映射

Arrow 抓取方法将 Microsoft SQL 类型映射到 C++ 级别的 Arrow 类型。

Microsoft SQL 类型	箭头类型
int, smallint, tinyint, bigint	int32、int16、int8、int64
float、real	float64、float32
十进制， 数字	decimal128
比特	bool
char, varchar, nchar, nvarchar	utf8
文本， ntext	large_utf8
二进制， 变数	binary、large_binary
日期	date32
time	time64[us]
datetime，datetime2，smalldatetime	timestamp[us]
datetimeoffset	timestamp[us, tz=UTC]
uniqueidentifier	utf8 （大写字符串）
xml	utf8

注释

驱动程序将 datetimeoffset 类型转换为UTC，因为箭头列需要固定的时区。 该驱动程序在转换过程中将 Microsoft SQL 中每个单元格的时区信息规范化为 UTC。

sql_variant 类型不受 Arrow 提取方法支持，并会引发“不支持的数据类型”异常。 对于返回fetchone()列的查询，请使用标准 fetchmany()、 fetchall()或 sql_variant 。

性能注意事项

Arrow 提取方法在分析和大批量数据操作中速度更快，而标准游标方法则更适合结果集较小的事务型模式。

何时使用 Arrow 而非标准 fetch

Scenario	建议的方法
获取数行进行显示	fetchone() / fetchall()
将数据加载到 Pandas 或 Polars	cursor.arrow()
将大型数据集分块处理	cursor.arrow_reader()
单行查找或小型结果集	fetchone() / fetchval()
分析管道或聚合管道	cursor.arrow() + Polars/DuckDB
将结果写入 Parquet 或 Arrow IPC	cursor.arrow_reader() + PyArrow I/O

大型数据集的内存管理

对于可能超出可用内存的结果集，请使用 arrow_reader()，并设置合理的 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")

调整批量大小

该 batch_size 参数控制每批获取的行数。 最佳大小取决于你的行宽度和可用内存。 包含 nvarchar(max) 或 varbinary(max) 等较大列的宽行更适合较小的批处理大小，而窄行则更适合较大的批处理大小。

• 默认（8192）：大多数工作负载的平衡良好。
• 较小（1000-5000）：用于宽大表格和大列。
• 更大（50000-100000）：用于窄表或吞吐量比内存更重要时。

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

相关内容

• 在 Polars 中使用 mssql-python
• 使用 mssql-python 搭配 pandas
• 使用 mssql-python 检索数据
• mssql-python 的数据类型映射
• mssql-python 应用的性能调优