Polars技术问题：如何获取表达式列数并验证参数列数要求？

获取表达式列数的方法 #

Polars中可以通过表达式的meta.output_shape()方法获取输出维度，返回值的第一个元素就是列数。对于IntoExpr类型的输入（支持列名字符串、单个Expr、多Expr组成的列表等），需要先通过pl.expr.to_expr_list()将其统一转换为表达式列表，再通过列表长度判断列数，或者针对单个表达式用meta.output_shape()确认输出列数。

参数验证实现逻辑 #

针对你的业务需求，验证步骤如下：

• 将输入的x和y转换为表达式列表，统一处理不同类型的输入
• 检查y对应的表达式列表长度必须为1，否则抛出异常
• x允许多列，只需确保转换后的列表长度至少为1（可根据业务需求调整）

完整代码示例 #

import polars as pl
from polars import IntoExpr

def func(x: IntoExpr, y: IntoExpr):
    # 将输入转换为表达式列表，统一处理不同类型的输入
    x_exprs = pl.expr.to_expr_list(x)
    y_exprs = pl.expr.to_expr_list(y)
    
    # 验证y必须为单列
    if len(y_exprs) != 1:
        raise ValueError("参数y必须是单列（单个列名或单个表达式）")
    
    # 验证x至少包含一列（可选，根据业务需求调整）
    if len(x_exprs) == 0:
        raise ValueError("参数x至少需要包含一列")
    
    # 可选：进一步确认单个表达式的输出列数（避免返回多列的表达式，比如pl.struct）
    for expr in y_exprs:
        if expr.meta.output_shape()[0] != 1:
            raise ValueError("参数y对应的表达式输出必须为单列")
    
    # 后续业务逻辑
    ...

# 测试用例
if __name__ == "__main__":
    # 合法调用：x多列，y单列
    func(["col1", "col2"], "col3")
    func(pl.col("col1") + pl.col("col2"), pl.col("col3"))
    
    # 非法调用：y多列，触发异常
    try:
        func("col1", ["col2", "col3"])
    except ValueError as e:
        print(e)  # 输出：参数y必须是单列（单个列名或单个表达式）

关键说明 #

• pl.expr.to_expr_list()是处理IntoExpr类型的核心方法，能将各种输入（列名、单个Expr、列表）统一转为Expr列表，方便统计列数
• 额外的expr.meta.output_shape()检查是为了规避本身返回多列的表达式（比如pl.struct(["a", "b"])），如果业务场景不需要可以省略

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