Apache Arrow DataFusion表达式(Expr)操作完全指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00884/article/details/148527009

Apache Arrow DataFusion表达式(Expr)操作完全指南

arrow-datafusion Apache Arrow DataFusion SQL Query Engine 项目地址: https://gitcode.com/gh_mirrors/arr/arrow-datafusion

表达式(Expr)基础概念

在Apache Arrow DataFusion项目中，Expr(表达式)是表示计算逻辑的核心抽象。它采用编译器与数据库中常见的"表达式树"结构，能够灵活地表示各种计算过程。

表达式树的基本特点：

每个节点代表一个操作或值
子节点代表操作数
从下往上递归求值

表达式树结构解析

以SQL表达式a + b * c为例，其表达式树结构如下：

            ┌────────────────────┐
            │     BinaryExpr     │
            │       op: +        │
            └────────────────────┘
                   ▲     ▲
           ┌───────┘     └────────────────┐
           │                              │
┌────────────────────┐         ┌────────────────────┐
│     Expr::Col      │         │     BinaryExpr     │
│       col: a       │         │       op: *        │
└────────────────────┘         └────────────────────┘
                                        ▲    ▲
                               ┌────────┘    └─────────┐
                               │                       │
                    ┌────────────────────┐  ┌────────────────────┐
                    │     Expr::Col      │  │     Expr::Col      │
                    │       col: b       │  │       col: c       │
                    └────────────────────┘  └────────────────────┘

创建与使用表达式

创建标量UDF表达式

标量UDF(用户定义函数)是常见的表达式类型。以下示例展示如何创建和使用一个简单的"加一"函数：

// 创建UDF定义
let add_one_udf = create_udf(
    "add_one",
    vec![DataType::Int64],  // 输入参数类型
    Arc::new(DataType::Int64),  // 返回类型
    Volatility::Immutable,  // 函数特性
    make_scalar_function(add_one),  // 实际函数实现
);

// 创建调用字面量的表达式：add_one(5)
let expr = add_one_udf.call(vec![lit(5)]);

// 创建调用列的表达式：add_one(my_column)
let expr = add_one_udf.call(vec![col("my_column")]);

表达式重写技术

表达式重写是将一个表达式转换为另一个表达式的过程，常用于：

简化表达式结构
优化计算性能
类型转换
函数内联

实现表达式转换

以下示例展示如何将add_one UDF内联为加法表达式：

fn rewrite_add_one(expr: Expr) -> Result<Expr> {
    expr.transform(&|expr| {
        Ok(match expr {
            Expr::ScalarUDF(scalar_fun) if scalar_fun.fun.name == "add_one" => {
                let input_arg = scalar_fun.args[0].clone();
                let new_expression = input_arg + lit(1i64);
                Transformed::yes(new_expression)
            }
            _ => Transformed::no(expr),
        })
    })
}

创建优化规则

将表达式重写逻辑封装为优化规则，可应用于整个查询计划：

struct AddOneInliner {}

impl OptimizerRule for AddOneInliner {
    fn name(&self) -> &str {
        "add_one_inliner"
    }

    fn try_optimize(
        &self,
        plan: &LogicalPlan,
        config: &dyn OptimizerConfig,
    ) -> Result<Option<LogicalPlan>> {
        let new_expressions = plan
            .expressions()
            .into_iter()
            .map(|expr| rewrite_add_one(expr))
            .collect::<Result<Vec<_>>>()?;

        let inputs = plan.inputs().into_iter().cloned().collect::<Vec<_>>();
        let plan = plan.with_new_exprs(&new_expressions, &inputs);
        plan.map(Some)
    }
}

表达式类型推导

获取表达式的数据类型是常见需求，可通过get_type方法实现：

let expr = col("c1") + col("c2");
let schema = DFSchema::new_with_metadata(
    vec![
        DFField::new_unqualified("c1", DataType::Int32, true),
        DFField::new_unqualified("c2", DataType::Float32, true),
    ],
    HashMap::new(),
).unwrap();

println!("type = {}", expr.get_type(&schema).unwrap());
// 输出: type = Float32