快速入门:Apache Arrow 搭配 Python 的 mssql-python 驱动

在本快速入门教程中,使用 mssql-python 驱动程序内置的 Arrow 获取方法,以 Apache Arrow 列式表的形式检索 SQL Server 数据。 Arrow 的列式内存格式可实现高性能分析、与 pandas、Polars 和 DuckDB 的零拷贝互操作,以及高效的 Parquet 文件 I/O,而无需逐行创建 Python 对象。

驱动程序 mssql-python 不需要 Windows 计算机上的任何外部依赖项。 驱动一次性安装了所有需要 pip 的东西,所以你可以用最新版本的驱动来运行新脚本,而不会破坏那些你没时间升级和测试的其他脚本。

mssql-python 文档 | mssql-python 源代码 | 包 (PyPI) | Uv

先决条件

安装一次性操作系统特定的先决条件。 Windows 用户可以跳过这一步。 完整平台详情请参见 “安装 mssql-python”。

apk add libtool krb5-libs krb5-dev

创建 SQL 数据库

在以下平台之一创建或连接SQL数据库:

创建项目并运行代码

  1. 创建新项目
  2. 添加依赖项
  3. 启动 Visual Studio Code
  4. 更新 pyproject.toml
  5. 更新 main.py
  6. 保存连接字符串
  7. 使用 uv 运行执行脚本

创建新项目

  1. 在开发目录中打开命令提示符。 如果没有,可以创建一个新的目录,比如 pythonscripts。 避免在OneDrive上放置文件夹,因为同步会影响管理虚拟环境。

  2. 通过使用 。 创建新uv

    uv init arrow-qs
    cd arrow-qs
    

添加依赖项

在同一目录中安装 mssql-pythonpython-dotenvpyarrowrich 包。

uv add mssql-python python-dotenv pyarrow rich

启动 Visual Studio Code

在同一目录中运行以下命令。

code .

更新 pyproject.toml

  1. pyproject.toml 文件包含了你项目的元数据。 在喜欢的编辑器中打开该文件。

  2. 查看文件的内容。 它应类似于此示例。 请注意,对于 mssql-python 的 Python 版本和依赖项,应使用 >= 来定义最低版本。 如果更喜欢确切的版本,请将 >= 版本号之前的版本号更改为 ==。 然后,每个包的解析版本存储在 uv.lock 中。 锁文件确保开发者使用一致的包版本。 提交 pyproject.tomluv.lock 两者,并在 CI 中运行经组织批准的依赖项扫描器。 不要直接编辑 uv.lock 文件。

    [project]
    name = "arrow-qs"
    version = "0.1.0"
    description = "Add your description here"
    readme = "README.md"
    requires-python = ">=3.11"
    dependencies = [
        "mssql-python>=1.5.0",
        "pyarrow>=19.0.0",
        "python-dotenv>=1.1.1",
        "rich>=14.1.0",
    ]
    
  3. 更新说明以更具描述性。

    description = "Fetch SQL Server data as Apache Arrow tables using mssql-python"
    
  4. 保存并关闭该文件。

更新 main.py

  1. 打开名为main.py的文件。 它应类似于此示例。

    def main():
        print("Hello from arrow-qs!")
    
    if __name__ == "__main__":
        main()
    
  2. main.py 的所有内容替换为以下代码。

    """Fetch SQL Server data as Apache Arrow tables using mssql-python."""
    
    from os import getenv
    
    import pyarrow as pa
    import pyarrow.parquet as pq
    from dotenv import load_dotenv
    from mssql_python import connect, Connection
    from rich.console import Console
    from rich.table import Table
    
    console = Console()
    
    
    def get_connection() -> Connection:
        """Create a connection using the connection string from .env."""
        load_dotenv()
        conn_str = getenv("SQL_CONNECTION_STRING")
        if not conn_str:
            raise ValueError("SQL_CONNECTION_STRING not set in .env file")
        return connect(conn_str)
    
    
    def fetch_arrow_table(conn: Connection) -> pa.Table:
        """Run a query and return the full result as an Arrow Table."""
        cursor = conn.cursor()
        cursor.execute("""
            SELECT
                p.ProductID,
                p.Name,
                p.ProductNumber,
                p.Color,
                p.StandardCost,
                p.ListPrice,
                p.Size,
                p.Weight,
                p.SellStartDate,
                pc.Name AS Category
            FROM SalesLT.Product AS p
            INNER JOIN SalesLT.ProductCategory AS pc
                ON p.ProductCategoryID = pc.ProductCategoryID
            ORDER BY p.ListPrice DESC
        """)
        arrow_table = cursor.arrow()
        cursor.close()
        return arrow_table
    
    
    def fetch_arrow_batches(conn: Connection) -> pa.Table:
        """Stream results one batch at a time using arrow_batch()."""
        cursor = conn.cursor()
        cursor.execute("""
            SELECT
                c.CustomerID,
                c.CompanyName,
                c.EmailAddress,
                COUNT(soh.SalesOrderID) AS OrderCount,
                SUM(soh.SubTotal + soh.TaxAmt + soh.Freight) AS TotalSpent
            FROM SalesLT.Customer AS c
            LEFT OUTER JOIN SalesLT.SalesOrderHeader AS soh
                ON c.CustomerID = soh.CustomerID
            GROUP BY
                c.CustomerID,
                c.CompanyName,
                c.EmailAddress
            ORDER BY TotalSpent DESC
        """)
        batches = []
        while True:
            batch = cursor.arrow_batch()
            if batch is None or batch.num_rows == 0:
                break
            batches.append(batch)
        cursor.close()
    
        if not batches:
            return pa.table({})
    
        return pa.Table.from_batches(batches)
    
    
    def fetch_with_reader(conn: Connection) -> pa.Table:
        """Use arrow_reader() to stream results as a RecordBatchReader."""
        cursor = conn.cursor()
        cursor.execute("""
            SELECT
                soh.SalesOrderID,
                soh.OrderDate,
                (soh.SubTotal + soh.TaxAmt + soh.Freight) AS TotalDue,
                c.CompanyName
            FROM SalesLT.SalesOrderHeader AS soh
            INNER JOIN SalesLT.Customer AS c
                ON soh.CustomerID = c.CustomerID
            ORDER BY soh.OrderDate DESC
        """)
        reader = cursor.arrow_reader()
        arrow_table = reader.read_all()
        cursor.close()
        return arrow_table
    
    
    def display_arrow_table(arrow_table: pa.Table, title: str, max_rows: int = 10) -> None:
        """Display an Arrow table using rich formatting."""
        rich_table = Table(title=title)
    
        for name in arrow_table.column_names:
            rich_table.add_column(name, style="bright_white")
    
        for i in range(min(max_rows, arrow_table.num_rows)):
            row = [str(arrow_table.column(col)[i].as_py()) for col in range(arrow_table.num_columns)]
            rich_table.add_row(*row)
    
        if arrow_table.num_rows > max_rows:
            rich_table.add_row(*[f"... ({arrow_table.num_rows - max_rows} more rows)" if col == 0 else "" for col in range(arrow_table.num_columns)])
    
        console.print(rich_table)
        console.print(f"\n[dim]Schema: {arrow_table.num_columns} columns, {arrow_table.num_rows} rows[/dim]\n")
    
    
    def save_to_parquet(arrow_table: pa.Table, file_path: str) -> None:
        """Save an Arrow table to a Parquet file."""
        pq.write_table(arrow_table, file_path)
        console.print(f"[green]Saved {arrow_table.num_rows} rows to {file_path}[/green]\n")
    
    
    def main() -> None:
        conn = get_connection()
    
        # 1. Fetch entire result as an Arrow Table with cursor.arrow()
        console.rule("[bold]cursor.arrow() - Full table fetch[/bold]")
        products = fetch_arrow_table(conn)
        display_arrow_table(products, "Products (Top 10 by List Price)")
    
        # 2. Stream results in batches with cursor.arrow_batch()
        console.rule("[bold]cursor.arrow_batch() - Batch streaming[/bold]")
        customers = fetch_arrow_batches(conn)
        display_arrow_table(customers, "Customers by Total Spent")
    
        # 3. Use RecordBatchReader with cursor.arrow_reader()
        console.rule("[bold]cursor.arrow_reader() - RecordBatchReader[/bold]")
        orders = fetch_with_reader(conn)
        display_arrow_table(orders, "Recent Orders")
    
        # 4. Save to Parquet
        console.rule("[bold]Save to Parquet[/bold]")
        save_to_parquet(products, "products.parquet")
    
        # 5. Read back from Parquet and verify
        loaded = pq.read_table("products.parquet")
        console.print(f"[green]Read back {loaded.num_rows} rows from products.parquet[/green]")
        console.print(f"[dim]Schema: {loaded.schema}[/dim]\n")
    
        conn.close()
    
    
    if __name__ == "__main__":
        main()
    

保存连接字符串

  1. 打开.gitignore文件,并为.env文件添加一个排除项。 文件应类似于此示例。 完成后,请务必保存并关闭它。

    # Python-generated files
    __pycache__/
    *.py[oc]
    build/
    dist/
    wheels/
    *.egg-info
    
    # Virtual environments
    .venv
    
    # Connection strings and secrets
    .env
    
    # Generated data files
    *.parquet
    
  2. 在当前目录中,创建一个名为 .env 的新文件。

  3. .env 文件中,为连接字符串 SQL_CONNECTION_STRING 添加一个条目。 将此处的示例替换为实际连接字符串值。

    SQL_CONNECTION_STRING="Server=<server_name>;Database=<database_name>;Encrypt=yes;TrustServerCertificate=no;Authentication=ActiveDirectoryInteractive"
    

    Important

    .env 保留在本地,不要纳入版本控制。 对于 CI 和部署环境,应从平台的机密存储中注入连接字符串或其组成部分所需的机密信息,而不是在机器之间复制 .env

    Tip

    你使用的连接字符串很大程度上取决于你连接的SQL数据库类型。 如果要连接到 Fabric 中的 Azure SQL 数据库SQL 数据库,请使用连接字符串选项卡中的 ODBC 连接字符串。可能需要根据方案调整身份验证类型。 有关连接字符串及其语法的详细信息,请参阅 连接字符串语法参考

使用 uv run 命令执行脚本

Tip

在 macOS 上, ActiveDirectoryInteractiveActiveDirectoryDefault 都可用于 Microsoft Entra 身份验证。 ActiveDirectoryInteractive 每次运行脚本时都会提示你登录。 为避免重复登录提示,请通过Azure CLI运行 az login,然后使用ActiveDirectoryDefault,该登录可重复使用缓存的凭证。

  • 在之前所在的终端窗口中,或打开同一目录的新终端窗口,运行以下命令。

    uv run main.py
    

    该脚本演示了三种 Arrow 获取方法:

    • cursor.arrow() 返回一个包含所有行的完整 pyarrow.Table。 最适合需要将整个数据集载入内存的中小型结果集。

    • cursor.arrow_batch() 每次返回一个 pyarrow.RecordBatch。 最适合大型结果集,比如你想以增量方式处理数据,但又不想把所有数据加载到内存里。

    • cursor.arrow_reader() 返回一个用于流式传输的 pyarrow.RecordBatchReader。 最适合用于管道式处理,或直接传递给可接收读取器的库。

    脚本还会将产品数据保存到Parquet文件中,并读取以验证往返。

代码的工作原理

  1. 连接:脚本从.env文件加载连接字符串,并使用mssql_python.connect()创建连接。

  2. 全表提取cursor.arrow()执行查询,并将整个结果集作为pyarrow.Table返回。 驱动通过 Arrow C 数据接口转换其 C++ 层的数据,绕过 Python 对象创建以提升性能。

  3. 批量流式传输cursor.arrow_batch() 每次调用返回一个 pyarrow.RecordBatch 。 循环收集批次直到没有更多行,然后将它们合并成一个单一表。 这种方法适用于大型数据集或你想独立处理每批数据时。

  4. RecordBatchReadercursor.arrow_reader() 返回一个 pyarrow.RecordBatchReader,一个标准的 Arrow 接口,许多库直接接受。 调用 reader.read_all() 会将整个流汇聚成一个表。

  5. Parquet I/Opyarrow.parquet.write_table() 将 Arrow 表保存为压缩的 Parquet 文件。 该格式保留列类型并支持高效的部分读取。

后续步骤

通过这些文章继续深入学习:

  • Arrow 集成 ,支持高级 Arrow 模式,包括批处理、内存管理和库互操作。
  • pandas 集成 ,将查询结果直接加载到 DataFrames 中。
  • Polars 集成,用于从 Arrow 原生查询构建 Polars 数据帧。