如何基于传入参数动态过滤Polars延迟DataFrame并选择列？

我需要修改下面的query_parquet函数，让它的filter表达式能根据传入的参数动态生成：

• 当传入eq_cols = {"a": "aaa", "b": "bbb"}时，自动生成两个相等判断的过滤条件；如果不传该参数，则不添加这类判断
• isin_cols参数逻辑同理：传入{"C": ["ccc1", "ccc2"]}时生成对应的is_in过滤条件，不传则不添加
• select_cols参数逻辑相同：传入指定列名列表时选择这些列，不传则使用默认的列选择（或保留所有列）

原始代码如下：

import polars as pl
from typing import Dict, List

def query_parquet(
        file_path: str,
        eq_cols: Dict[str, str] = None,
        isin_cols: Dict[str, List[str]] = None,
        select_cols: List[str] = None,
    ) -> pl.DataFrame:

    query = (
        pl.scan_parquet(file_path)
        .filter(
            pl.col("A") == "aaa",
            # pl.col("B") == "bbb",
            # pl.col("C").is_in(["ccc1", "ccc2"])
            # pl.col("D").is_in(["ddd1", "ddd2"])
        )
        .select(["F", "G", "H"])
        .collect()
    )

    return query

实现方法 #

核心思路是动态构建过滤表达式列表，再根据参数决定是否添加过滤条件和选择列：

1. 初始化查询对象为Parquet文件的扫描器
2. 遍历eq_cols生成相等判断表达式，遍历isin_cols生成is_in判断表达式，将所有有效表达式收集到列表中
3. 如果过滤表达式列表不为空，就将其传入filter方法
4. 根据select_cols参数决定选择的列：传入则用指定列，否则保留默认的列选择（或改为选择所有列，根据需求调整）

修改后的完整代码：

import polars as pl
from typing import Dict, List, Optional

def query_parquet(
        file_path: str,
        eq_cols: Optional[Dict[str, str]] = None,
        isin_cols: Optional[Dict[str, List[str]]] = None,
        select_cols: Optional[List[str]] = None,
    ) -> pl.DataFrame:
    # 初始化扫描器
    query = pl.scan_parquet(file_path)
    
    # 构建过滤表达式列表
    filters = []
    # 处理eq_cols：生成相等判断
    if eq_cols:
        for col, val in eq_cols.items():
            filters.append(pl.col(col) == val)
    # 处理isin_cols：生成is_in判断
    if isin_cols:
        for col, vals in isin_cols.items():
            filters.append(pl.col(col).isin(vals))
    
    # 应用过滤条件（如果有）
    if filters:
        query = query.filter(*filters)
    
    # 处理列选择：传入指定列则用它，否则用默认列
    if select_cols:
        query = query.select(select_cols)
    else:
        query = query.select(["F", "G", "H"])
    
    # 执行并返回结果
    return query.collect()

关键细节说明 #

• 使用Optional类型标注参数，让类型提示更准确（Python 3.10+直接支持，低版本需从typing导入）
• 过滤表达式通过列表收集后，用*filters解包传入filter方法，Polars会自动将多个表达式用逻辑与（&）组合
• 若不需要默认列选择，不传select_cols时返回所有列，只需删除else分支：

  if select_cols:
      query = query.select(select_cols)
  # 无else分支则默认返回所有列
  

• 可轻松扩展其他过滤逻辑（如大于、小于判断），只需在filters列表中添加对应表达式即可

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