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