使用 mssql-python 搭配 Polars

Polars 是一个用 Rust 编写的高性能 DataFrame 库，提供了一种快速、内存高效的替代 pandas 方案。 Polars 与 mssql-python 驱动结合使用后，您可以：

• 直接将SQL查询结果加载到Polars DataFrames中。
• 使用 Apache Arrow 进行从 Microsoft SQL 零复制的数据传输。
• 高效地将 Polars 数据帧写回 Microsoft SQL。
• 利用懒惰的评估构建高性能数据管道。

本文中的示例查询 AdventureWorks 示例数据库。 如果你还没有，可以参考 AdventureWorks的样本数据库。

将数据读取到Polars DataFrame中

你可以通过两种方式将 Microsoft SQL 数据加载到 Polars：通过标准光标方法逐行转换，或通过 Apache Arrow 实现零副本传输。 “零复制”意味着数据保持在一个单一的内存缓冲区，驱动、Arrow 和 Polars 都直接读取，这样就不会有行被重复到中间的 Python 对象中。 由于这种方式具有较高的效率，对于大多数工作负载，请使用 Arrow 方法。

对DataFrame的基本查询

这种方法使用标准光标获取所有行，并手动构建 Polars 数据框。 它接受参数 化查询 以实现安全值替换。 它在没有 PyArrow 的情况下工作，但对于大型结果集来说速度较慢，因为每个值都经过 Python。

import polars as pl
import mssql_python

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


def query_to_polars(cursor, query: str, params: dict = None) -> pl.DataFrame:
    """Execute query and return results as Polars DataFrame."""
    cursor.execute(query, params or {})

    columns = [col[0] for col in cursor.description]
    rows = cursor.fetchall()

    data = {col: [row[i] for row in rows] for i, col in enumerate(columns)}
    return pl.DataFrame(data)

# Usage: %(color)s is a parameterized placeholder. The driver safely substitutes
# the value from the dict, which prevents SQL injection.
df = query_to_polars(cursor, "SELECT TOP 5 Name, ListPrice FROM Production.Product WHERE Color = %(color)s", {"color": "Black"})
print(df)

使用Arrow进行零拷贝传输（推荐）

将 Microsoft SQL 数据加载到 Polars 的最高效方式是通过 Apache Arrow。 mssql-python 驱动程序的 arrow() 方法返回一个 pyarrow.Table，Polars 可以以零拷贝开销直接使用它。

def query_to_polars_arrow(cursor, query: str, params: dict = None) -> pl.DataFrame:
    """Execute query and load results through Arrow for best performance."""
    cursor.execute(query, params or {})
    arrow_table = cursor.arrow()
    return pl.from_arrow(arrow_table)

# Usage
df = query_to_polars_arrow(cursor, "SELECT ProductID, Name, ListPrice FROM Production.Product")
print(df)

使用 Arrow 批次流式传输大型数据集

对于内存无法容纳的数据集，可以用 arrow_reader() 流式批处理数据。 每一批都是一个 pyarrow.RecordBatch，Polars 可以独立处理它，因此内存使用量与 batch_size 成正比，而不是与整个结果集成正比。

def process_large_query(cursor, query: str, params: dict = None, batch_size: int = 50000) -> pl.DataFrame:
    """Process large query results as streaming Arrow batches."""
    cursor.execute(query, params or {})
    reader = cursor.arrow_reader(batch_size=batch_size)

    results = []
    for batch in reader:
        chunk_df = pl.from_arrow(batch)
        # Process each chunk
        results.append(chunk_df)

    return pl.concat(results) if results else pl.DataFrame()

# Usage
df = process_large_query(cursor, "SELECT * FROM Production.TransactionHistory")

使用 LazyFrames 进行延迟执行

Polars LazyFrames 允许你构建一系列操作链（过滤、分组、排序），而无需立即运行它们。 Polars 在执行前会优化整个链条，这比单独应用步骤更快。

def query_to_lazy(cursor, query: str, params: dict = None) -> pl.LazyFrame:
    """Execute query and return a Polars LazyFrame."""
    cursor.execute(query, params or {})
    arrow_table = cursor.arrow()
    return pl.from_arrow(arrow_table).lazy()

# Build a query plan without executing immediately
lf = query_to_lazy(cursor, "SELECT SalesOrderID, CustomerID, TotalDue, OrderDate FROM Sales.SalesOrderHeader")
result = (
    lf.filter(pl.col("TotalDue") > 100)
    .group_by("CustomerID")
    .agg([
        pl.col("TotalDue").sum().alias("TotalSpent"),
        pl.col("SalesOrderID").count().alias("OrderCount")
    ])
    .sort("TotalSpent", descending=True)
    .collect()  # Execute the optimized plan
)
print(result)

将 Polars DataFrames 写入 Microsoft SQL

为标识符加引号以防止 SQL 注入

在SQL中，表和列名不能作为查询参数传递。 当你构建带有动态标识符的SQL语句时，请用方括号包裹每个名称，并剔除任何嵌入 ] 字符，以防止SQL注入。

def quote_id(identifier: str) -> str:
    """Quote a Microsoft SQL identifier to prevent SQL injection.
    
    Wraps the name in square brackets and escapes any embedded ] characters.
    Raises ValueError if the identifier is empty or contains null bytes.
    """
    if not identifier or "\x00" in identifier:
        raise ValueError(f"Invalid identifier: {identifier!r}")
    escaped = identifier.replace("]", "]]")
    return f"[{escaped}]"

本节中的辅助函数在生成的 SQL 中对所有表名和列名都使用 quote_id()。

插入DataFrame行

逐行处理方法使用 iter_rows(named=True) 对 DataFrame 进行迭代，并对每一行执行一次 INSERT。 这种方法简单但对大批量来说较慢，因为每行都需要往返服务器。

def polars_to_sql(cursor, conn, df: pl.DataFrame, table: str) -> int:
    """Write Polars DataFrame to Microsoft SQL table."""
    columns = df.columns
    placeholders = ", ".join([f"%({col})s" for col in columns])
    col_list = ", ".join([quote_id(col) for col in columns])
    query = f"INSERT INTO {quote_id(table)} ({col_list}) VALUES ({placeholders})"

    rows_inserted = 0
    for row in df.iter_rows(named=True):
        params = {k: (None if v is None else v) for k, v in row.items()}
        cursor.execute(query, params)
        rows_inserted += 1

    conn.commit()
    return rows_inserted

# Usage
cursor.execute("CREATE TABLE #PolarsInsert (Name NVARCHAR(100), Price DECIMAL(10,2), CategoryID INT)")
df = pl.DataFrame({
    "Name": ["Product A", "Product B"],
    "Price": [29.99, 49.99],
    "CategoryID": [1, 2]
})
rows = polars_to_sql(cursor, conn, df, "#PolarsInsert")
print(f"Inserted {rows} rows")

批量插入（推荐用于大型数据帧）

对于大型数据框，请使用驱动程序的 bulkcopy() 方法，通过 TDS（表格数据流）协议批量发送多行数据；TDS 是 Microsoft SQL 使用的原生网络协议。 这种方法减少了往返次数，且比逐行插入更快。

def polars_to_sql_bulk(conn, df: pl.DataFrame, table: str) -> int:
    """Bulk insert Polars DataFrame using BCP for best performance."""
    rows = [tuple(None if v is None else v for v in row) for row in df.iter_rows()]

    cursor = conn.cursor()
    result = cursor.bulkcopy(table, rows)
    conn.commit()
    return result["rows_copied"]

# Usage
cursor.execute("CREATE TABLE ##PolarsBulk (Name NVARCHAR(50), Price FLOAT, CategoryID INT)")
conn.commit()
df = pl.DataFrame({
    "Name": ["Product A", "Product B", "Product C"],
    "Price": [29.99, 49.99, 19.99],
    "CategoryID": [1, 2, 1]
})
rows = polars_to_sql_bulk(conn, df, "##PolarsBulk")
print(f"Bulk inserted {rows} rows")

数据分析模式

以下示例展示了结合 Microsoft SQL 查询与 Polars 变换的常见分析任务。

聚合查询

本示例按子类别分组产品，计算 SQL 中的数量和价格统计，然后将摘要加载到 Polars 数据框架中：

def get_sales_summary(cursor) -> pl.DataFrame:
    """Get sales summary by subcategory."""
    cursor.execute("""
        SELECT
            sc.Name AS SubcategoryName,
            COUNT(*) AS ProductCount,
            AVG(p.ListPrice) AS AvgPrice,
            MIN(p.ListPrice) AS MinPrice,
            MAX(p.ListPrice) AS MaxPrice
        FROM Production.Product p
        JOIN Production.ProductSubcategory sc ON p.ProductSubcategoryID = sc.ProductSubcategoryID
        GROUP BY sc.Name
        ORDER BY ProductCount DESC
    """)
    return pl.from_arrow(cursor.arrow())

df = get_sales_summary(cursor)
print(df)

时序分析

从 Microsoft SQL 加载时间序列数据，并使用 Polars 表达式添加计算列，例如滚动平均值。

def get_daily_sales(cursor, start_date: str, end_date: str) -> pl.DataFrame:
    """Get daily sales and compute rolling statistics."""
    cursor.execute("""
        SELECT
            CAST(OrderDate AS DATE) AS Date,
            COUNT(*) AS OrderCount,
            SUM(TotalDue) AS Revenue
        FROM Sales.SalesOrderHeader
        WHERE OrderDate BETWEEN %(start)s AND %(end)s
        GROUP BY CAST(OrderDate AS DATE)
        ORDER BY Date
    """, {"start": start_date, "end": end_date})

    df = pl.from_arrow(cursor.arrow())

    # Add rolling 7-day average
    df = df.with_columns(
        pl.col("Revenue").rolling_mean(window_size=7).alias("RollingAvg")
    )
    return df

sales_df = get_daily_sales(cursor, "2013-01-01", "2013-12-31")
print(sales_df)

将SQL数据与本地文件连接

你可以通过在 Polars 中与本地 CSV 文件连接 Microsoft SQL 数据来丰富它们。 将每个源加载到数据框中，并在内存中进行联接。

# Load SQL data via Arrow
cursor.execute("SELECT c.CustomerID, p.FirstName, p.LastName FROM Sales.Customer c JOIN Person.Person p ON c.PersonID = p.BusinessEntityID")
customers = pl.from_arrow(cursor.arrow())

# Load local CSV
orders = pl.read_csv("orders_export.csv")

# Join in Polars
result = customers.join(orders, on="CustomerID", how="inner")
print(result)

ETL 模式

通过结合 Microsoft SQL 查询与 Polars 转换，构建提取、转换和加载（ETL）流水线。 Polars 表达式处理转换步骤，而 bulkcopy() 负责加载。

提取、转换和加载

本示例通过Arrow提取活跃客户数据，应用带有Polars表达式的业务分段逻辑，并使用批量复制加载结果。

def etl_pipeline(source_cursor, dest_conn):
    """ETL pipeline using Polars transformations."""

    # Extract: derive a per-customer summary from order history via Arrow
    source_cursor.execute("""
        SELECT
            CustomerID,
            COUNT(*) AS OrderCount,
            SUM(TotalDue) AS TotalSpent
        FROM Sales.SalesOrderHeader
        WHERE OrderDate > DATEADD(YEAR, -1, (SELECT MAX(OrderDate) FROM Sales.SalesOrderHeader))
        GROUP BY CustomerID
    """)
    df = pl.from_arrow(source_cursor.arrow())
    df = df.with_columns(pl.col("TotalSpent").cast(pl.Float64))

    # Transform with Polars expressions
    df = df.with_columns([
        pl.when(pl.col("TotalSpent") > 1000).then(pl.lit("Platinum"))
          .when(pl.col("TotalSpent") > 500).then(pl.lit("Gold"))
          .when(pl.col("TotalSpent") > 100).then(pl.lit("Silver"))
          .otherwise(pl.lit("Bronze"))
          .alias("CustomerSegment"),
        (pl.col("TotalSpent") / pl.col("OrderCount").clip(lower_bound=1))
          .alias("AvgOrderValue"),
        (pl.col("TotalSpent") > 500).alias("IsHighValue")
    ])

    # Load via bulk copy into the destination table
    dest_cursor = dest_conn.cursor()
    dest_cursor.execute("""
        CREATE TABLE ##CustomerAnalytics (
            CustomerID INT,
            CustomerSegment NVARCHAR(20),
            AvgOrderValue FLOAT,
            IsHighValue BIT
        )
    """)
    dest_conn.commit()

    load_df = df.select(["CustomerID", "CustomerSegment", "AvgOrderValue", "IsHighValue"])
    polars_to_sql_bulk(dest_conn, load_df, "##CustomerAnalytics")

    return len(df)

性能提示

以下建议能帮助你充分利用mssql-python和Polars的组合。

让 Microsoft SQL 来处理繁重的工作

Microsoft SQL 在聚合、过滤和连接方面比直接拉取原始数据再用本地 Python 处理更快。 尽量让 Microsoft SQL 来完成繁重工作，只在网络中移动你需要的数据，使用 Polars 进行分析和转换，这些在 Python 中更方便。

# Avoid: pulling all rows over the wire to aggregate locally in Polars
df = query_to_polars_arrow(cursor, "SELECT * FROM Production.Product WHERE Color IS NOT NULL")  # transfers entire table
summary = df.group_by("Color").agg(pl.col("ListPrice").sum())  # aggregation that SQL can do faster

# Better: push the aggregation into SQL and transfer only the summary
df = query_to_polars_arrow(cursor, """
    SELECT Color AS Category, SUM(ListPrice) AS TotalAmount
    FROM Production.Product
    WHERE Color IS NOT NULL
    GROUP BY Color
""")

所有读取操作都使用Arrow

基于箭头的传输避免了创建中间的 Python 对象，从而减少内存使用并提升吞吐量。 对于超过几行的任何结果集，相比手动逐行转换，更推荐使用 cursor.arrow()。

# Suboptimal: Row-by-row conversion
cursor.execute("SELECT * FROM Production.TransactionHistory")
columns = [col[0] for col in cursor.description]
rows = cursor.fetchall()
df = pl.DataFrame({col: [row[i] for row in rows] for i, col in enumerate(columns)})

# Better: Arrow-based transfer
cursor.execute("SELECT * FROM Production.TransactionHistory")
df = pl.from_arrow(cursor.arrow())

相关内容

• Apache Arrow 集成
• Pandas整合
• 批量复制操作
• 获取数据
• 性能优化