如何用Polars流式读取数据库并控制读取批次大小？

问题背景 #

我正在学习用Polars替代Pandas，但遇到了一个问题：如何让Polars流式读取数据库，或者控制每次读取的数据量？我的数据量很大，无法一次性加载到内存，必须分批次读取。尝试用pl.read_database的batch_size参数时，发现它还是会把全部数据加载到内存里；另外用connectorx配合pl.read_database_uri也没法实现流式读取，请问有没有解决办法？

我的环境配置：

[tool.poetry.dependencies]
python = "^3.11"
connectorx = "^0.3.2"
polars = {extras = ["connectorx", "pyarrow", "xlsx2csv"], version = "^0.19.3"}
sqlalchemy = "^2.0.21"
pymysql = "^1.1.0"
notebook = "^7.0.4"

我已经实现了一种可行的流式读取方式（基于SQLAlchemy）：

# 用SQLAlchemy创建engine
with engine.connect() as conn:
    query = text('select * from my_table')
    # 开启流式结果，逐批获取数据
    with conn.execution_options(stream_results=True).execute(query) as cursor:
        while True:
             res = cursor.fetchmany(chunk_size)
             if not res:
                 break
             df = pl.DataFrame(res)
             # 在这里处理当前批次的df
             process_batch(df)

但我想知道怎么用pl.read_database或pl.read_database_uri结合batch_size参数实现类似的流式读取？

我尝试过以下代码，但它会一次性加载全部数据，达不到流式效果：

cursor = conn.cursor()

query = 'select * from users'

df = pl.read_database(query=query, connection=cursor, batch_size=1000)
for i in df.groupby(by='age'):
    print(type(i))

核心问题说明 #

pl.read_database的batch_size参数并不是用来实现流式迭代的——它只是控制connectorx在读取数据时的内部批次大小，最终还是会把所有批次合并成一个完整的DataFrame返回，所以无法实现分批次加载到内存的效果。

解决方案1：SQLAlchemy流式查询+Polars逐批转换 #

你已有的SQLAlchemy流式代码是可行的，这里优化成更贴合Polars的写法：

from sqlalchemy import create_engine, text
import polars as pl

engine = create_engine("mysql+pymysql://user:password@host:port/db_name")
chunk_size = 1000

with engine.connect() as conn:
    result = conn.execution_options(stream_results=True).execute(text("SELECT * FROM my_table"))
    # 逐批获取并转换为Polars DataFrame
    while batch := result.fetchmany(chunk_size):
        # 直接用查询结果的字段名作为Schema，避免列名丢失
        df_batch = pl.DataFrame(batch, schema=result.keys())
        # 处理当前批次数据（示例：分组统计、写入文件等）
        print(f"处理批次，行数：{len(df_batch)}")
        batch_stats = df_batch.group_by("age").agg(pl.count("id").alias("count"))
        print(batch_stats)

解决方案2：用pl.scan_database实现惰性加载（推荐） #

Polars 0.19+版本支持pl.scan_database，它采用惰性查询模式，不会立即加载全量数据，而是按需获取。结合fetch_size和iter_batches可以轻松实现流式处理：

import polars as pl
from sqlalchemy import create_engine

engine = create_engine("mysql+pymysql://user:password@host:port/db_name")

# 创建惰性查询计划，fetch_size控制每次从数据库拉取的行数
lazy_df = pl.scan_database(
    query="SELECT * FROM my_table",
    connection=engine,
    fetch_size=1000
)

# 方式1：直接将数据分批写入文件（无需手动迭代）
lazy_df.sink_parquet("output.parquet", batch_size=1000)

# 方式2：手动迭代批次处理
for batch in lazy_df.iter_batches(batch_size=1000):
    print(f"批次行数：{len(batch)}")
    # 自定义批次处理逻辑
    process_batch(batch)

注意：scan_database依赖SQLAlchemy，不支持connectorx后端，需确保Polars已安装sqlalchemy相关依赖。

关于connectorx的限制 #

pl.read_database_uri基于connectorx实现时，目前不支持流式读取——connectorx会一次性把全量数据加载到内存再转换为Polars DataFrame，不适合超大规模数据场景。如果必须使用connectorx，只能通过SQL分页手动分批次查询，但这种方式效率较低，且可能存在数据一致性问题（查询过程中数据更新会导致重复或遗漏）：

import polars as pl

batch_size = 1000
offset = 0
# 先获取总行数
total_rows = pl.read_database_uri(
    "SELECT COUNT(*) FROM my_table",
    uri="mysql://user:password@host:port/db_name"
).item()

while offset < total_rows:
    query = f"SELECT * FROM my_table LIMIT {batch_size} OFFSET {offset}"
    df_batch = pl.read_database_uri(query, uri="mysql://user:password@host:port/db_name")
    process_batch(df_batch)
    offset += batch_size

内容的提问来源于stack exchange，提问作者Jerry wang