Polars LazyFrame自定义函数apply查询计划异常，执行耗时过长

问题描述 #

需要将spaCy的nlp(<string>).vector函数应用于DataFrame的字符串列，该函数平均耗时13毫秒，返回包含300个Float64的ndarray，需将这些数值展开为独立列。采用如下代码实现：

import spacy
import polars as pl

nlp = spacy.load('en_core_web_lg')
full = pl.LazyFrame([["apple", "banana", "orange"]], schema=['keyword'])

VECTOR_FIELD_NAMES = ['dim_' + str(x) for x in range(300)]

full = full.with_columns(
    pl.col('keyword').map_elements(
        lambda x: tuple(nlp(x).vector), return_dtype=pl.List(pl.Float64)
    ).list.to_struct(fields=VECTOR_FIELD_NAMES).struct.unnest()
)

full.collect()

代码执行耗时11.5秒，比在Polars外部执行计算慢100多倍。查看查询计划发现：

naive plan: (run LazyFrame.explain(optimized=True) to see the optimized plan)
WITH_COLUMNS:
[col("keyword").map_list().list.to_struct().struct.field_by_name(dim_0)(),
 col("keyword").map_list().list.to_struct().struct.field_by_name(dim_1)(),
 col("keyword").map_list().list.to_struct().struct.field_by_name(dim_2)(),
 ...

计划中会为全部300个维度重复执行上述操作，怀疑每个输出单元格都会重复计算nlp(<keyword>)。

问题原因 #

Polars的LazyFrame在处理链式操作时，若未显式缓存中间结果，会对每个输出列重复执行整个前置操作链。这里struct.unnest()会将结构体的每个字段拆分为独立列，而Polars默认会为每个拆分出的列重新执行map_elements和list.to_struct步骤，导致每个关键词被调用300次nlp()，直接放大了计算量。

优化方案 #

核心思路是先计算出所有向量的列表列并缓存，再基于这个缓存列拆分出300个维度列，避免重复计算。

基础优化代码 #

import spacy
import polars as pl

nlp = spacy.load('en_core_web_lg')
full = pl.LazyFrame([["apple", "banana", "orange"]], schema=['keyword'])

VECTOR_FIELD_NAMES = ['dim_' + str(x) for x in range(300)]

full = (
    full
    # 先生成向量列表列，作为中间结果缓存
    .with_columns(
        pl.col('keyword').map_elements(
            lambda x: tuple(nlp(x).vector), return_dtype=pl.List(pl.Float64)
        ).alias('vector_list')
    )
    # 基于缓存的列表列拆分维度列
    .with_columns(
        pl.col('vector_list').list.to_struct(fields=VECTOR_FIELD_NAMES).struct.unnest()
    )
    # 可选：删除中间列
    .drop('vector_list')
)

full.collect()

进阶批量优化 #

spaCy的nlp.pipe()支持批量处理文本，比单条调用效率更高，可替换map_elements为map_batches结合nlp.pipe()，进一步降低耗时：

import spacy
import polars as pl

nlp = spacy.load('en_core_web_lg')
full = pl.LazyFrame([["apple", "banana", "orange"]], schema=['keyword'])

VECTOR_FIELD_NAMES = ['dim_' + str(x) for x in range(300)]

def batch_process(texts):
    return [tuple(doc.vector) for doc in nlp.pipe(texts)]

full = (
    full
    .with_columns(
        pl.col('keyword').map_batches(batch_process, return_dtype=pl.List(pl.Float64)).alias('vector_list')
    )
    .with_columns(
        pl.col('vector_list').list.to_struct(fields=VECTOR_FIELD_NAMES).struct.unnest()
    )
    .drop('vector_list')
)

full.collect()

内容的提问来源于Stack Exchange，提问作者Falcondance