Polars 内存装不下时如何用 streaming 引擎执行查询并查看哪些算子落回内存?
【免费下载链接】polarsExtremely fast Query Engine for DataFrames, written in Rust项目地址: https://gitcode.com/GitHub_Trending/po/polars
数据集大到超出可用内存时,Polars 默认的collect()会把全部数据当作一个批次处理,查询在内存峰值点上要求所有数据都装得下。解决办法是把 lazy 查询交给 streaming 引擎按批次执行,再通过物理计划图确认哪些算子是以流式方式运行、哪些算子被回退(fall back)到内存引擎。适用前提:查询以 lazy API 构建(如pl.scan_csv返回的LazyFrame);若要用show_graph出图,需要在本机安装 Graphviz 并加入 PATH。
确认默认执行方式为什么会把内存吃满
先明确问题所在。Query execution 文档对默认collect的说明是:
With the default
collectmethod Polars processes all of your data as one batch. This means that all the data has to fit into your available memory at the point of peak memory usage in your query.
也就是说,如果查询中的某些步骤(例如聚合前的中间结果)需要超过可用内存,默认执行就会失败或产生严重内存压力。文档对更大的数据集给出的路径是:
If your data requires more memory than you have available Polars may be able to process the data in batches usingstreamingmode. To use streaming mode you simply pass the
engine="streaming"argument tocollect
注意这里文档用的是 "may be able to"——是否真的能按批次处理,取决于查询里的算子是否都有流式实现,这一点在后面的计划图中可以看到。
用collect(engine="streaming")执行查询
Streaming 文档给出的最小示例是:
import polars as pl q1 = ( pl.scan_csv("docs/assets/data/iris.csv") .filter(pl.col("sepal_length") > 5) .group_by("species") .agg(pl.col("sepal_width").mean()) ) df = q1.collect(engine="streaming")把示例中的docs/assets/data/iris.csv替换成你自己的数据文件,并在collect上传engine="streaming"即可,查询结构不需要其他改动。
关于默认值有一个版本上的重要变化,见 Version 2.0-rc 升级说明:从 2.0 起,lazy 查询的engine="auto"(即不传参)解析为 streaming 引擎,而不再是内存引擎;eager 的DataFrame操作不受影响,仍解析到内存引擎。sink_*写文件本来就派发到 streaming 引擎,不受该参数影响。
如果某条查询你明确要强制回内存引擎执行,文档给出的写法是:
>>> lf.collect(engine="in-memory") # per query或者进程级设置:
>>> pl.Config.set_engine_affinity("in-memory") # process-wide以及环境变量POLARS_ENGINE_AFFINITY=in-memory。这些写法同时适用于 SQL 路径。
查看哪些算子落回内存引擎
Streaming 文档明确说明:Polars 可以把很多操作按流式方式运行,但有些操作本质上不是流式的,或(暂时)还没有流式实现。这种情况下 Polars 会对这些算子回退到内存引擎,用户不需要感知这个回退,但排查内存或性能问题时值得去看:
To inspect the physical plan of streaming query, you can plot the physical graph. The legend shows how memory intensive the operation can be.
对应的操作是对查询调用show_graph,指定plan_stage="physical"和engine="streaming"(示例取自 streaming.py):
q1 = ( pl.scan_csv("docs/assets/data/iris.csv") .filter(pl.col("sepal_length") > 5) .group_by("species") .agg( mean_width=pl.col("sepal_width").mean(), mean_width2=pl.col("sepal_width").sum() / pl.col("sepal_length").count(), ) ) q1.show_graph(plan_stage="physical", engine="streaming")出图后按文档的描述来读:图例(legend)标注了每个算子可能有多大的内存开销,据此就能看出哪些节点是真正流式执行的、哪些会占用大块内存(即回退到内存引擎执行的部分)。如果不想弹窗查看,也可以像文档示例那样把图落到文件:
q1.show_graph( plan_stage="physical", engine="streaming", show=False, output_path="query_plan.png", )output_path替换为你想保存图片的路径即可。
还有一个版本相关的注意点:2.0 中explain()和show_graph()默认不再渲染流式计划,只有显式传入engine="streaming"时才会输出 streaming 物理计划(2.0 同时把show_graph的默认plan_stage从"ir"改为"physical")。所以核对算子时,engine="streaming"这个参数不能省。
限制与结果核对
- 行序不再保证:streaming 引擎对不需要行序的算子(
unpivot、group_by、join 等)不保证行顺序。如果你的下游代码依赖某种偶然顺序,要么显式sort,要么在算子支持时传maintain_order(例如join(..., maintain_order="left"))。 - 回退不是错误:某个算子落回内存引擎意味着该算子执行时仍需要内存,如果数据集对该算子确实装不下,仅靠 streaming 不足以解决问题——此时应结合 Query execution 文档中的分片思路,用
scan子集 + 首尾.head/.collect在小数据上调试查询,再回到全量数据执行。 - 验证查询是否走通的方式就是
collect(engine="streaming")正常返回 DataFrame;算子级别的流式/回退情况以show_graph(plan_stage="physical", engine="streaming")出的物理计划图及其图例为准。
相关文档入口:Streaming、Query execution、Query plan、Version 2.0-rc 升级说明。
【免费下载链接】polarsExtremely fast Query Engine for DataFrames, written in Rust项目地址: https://gitcode.com/GitHub_Trending/po/polars
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考