如何在Polars IO插件中解析并转换predicate（谓词）？

问题背景 #

Polars的自定义IO插件能实现自定义数据读取逻辑，查询引擎还会自动传递谓词（predicate）来做读取阶段的过滤优化，但官方未明确说明如何将这些Polars表达式转成纯Python可用的对象。

比如你写了这样的自定义scan_data函数：

def scan_data(path: str) -> pl.LazyFrame:

    def source_generator(
        with_columns: list[str] | None,
        predicate: pl.Expr | None,
        n_rows: int | None,
        batch_size: int | None,
    ) -> Iterator[pl.DataFrame]:
        print(predicate.to_string())
        # 这里需要用谓词优化数据读取，但不知道怎么解析它

    return register_io_source(io_source=source_generator, schema=schema)

当你执行带filter的查询时：

from polars_data import scan_data # 你的自定义库

df = scan_data("file-path")
df = df.filter(pl.col("col-name").is_in(["a", "b"]))
df.count()

插件里打印出的谓词是col("channel").is_in([Series])——这是个Polars表达式，里面的列表是Polars Series对象，没法直接给asammdf用，你需要把它转成Python原生列表来过滤通道：

from asammdf import MDF

with MDF(mdf_path, channels=<要填筛选后的通道列表>) as mdf:
    signals = []
    for channel in mdf.iter_channels():
        signal = pa.Table.from_pydict({"timestamp": channel.timestamps, "samples": channel.samples.astype(float)})
        signals.append(signal)
    
    signals = pa.concat_tables(signals)
    signals = signals.sort_by("timestamp")

    return pl.from_arrow(signals)

解决方法 #

你可以通过Polars表达式的元数据（meta）API来拆解谓词，提取出需要的列名和筛选值：

1. 编写谓词解析函数 #

针对is_in类型的谓词，写一个解析函数把它转成Python原生的列名和值列表：

def parse_is_in_predicate(predicate: pl.Expr | None) -> tuple[str, list] | None:
    if not predicate:
        return None
    
    # 判断是否是is_in表达式
    if predicate.meta.has_named_function("is_in"):
        # 获取目标列名
        col_name = predicate.meta.root_names()[0].meta.output_name()
        # 获取is_in里的Series参数，转成Python列表
        filter_series = predicate.meta.args()[0]
        filter_values = filter_series.to_list()
        
        return (col_name, filter_values)
    return None

2. 在读取逻辑中应用解析结果 #

修改你的source_generator，用解析出来的值来过滤asammdf的通道：

def source_generator(
    with_columns: list[str] | None,
    predicate: pl.Expr | None,
    n_rows: int | None,
    batch_size: int | None,
) -> Iterator[pl.DataFrame]:
    parsed_pred = parse_is_in_predicate(predicate)
    target_channels = None
    
    # 如果解析出了针对channel列的is_in筛选，就用这些值过滤通道
    if parsed_pred:
        col_name, filter_vals = parsed_pred
        if col_name == "channel":
            target_channels = filter_vals
    
    with MDF(mdf_path, channels=target_channels) as mdf:
        signals = []
        for channel in mdf.iter_channels():
            signal = pa.Table.from_pydict({
                "timestamp": channel.timestamps, 
                "samples": channel.samples.astype(float)
            })
            signals.append(signal)
        
        signals = pa.concat_tables(signals)
        signals = signals.sort_by("timestamp")
        
        yield pl.from_arrow(signals)

扩展说明 #

如果需要支持更多类型的谓词（比如等于、大于这类比较操作），可以继续扩展解析函数：

• 用predicate.meta.has_named_function("eq")判断是否是等于操作
• 用predicate.meta.args()获取左右操作数，左边是列，右边是常量值
• 类似地拆解其他表达式结构，把Polars Expr转成Python能直接用的条件判断

内容的提问来源于stack exchange，提问作者Herman Jonsson