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

本文链接：https://blog.youkuaiyun.com/gitblog_00783/article/details/148464416

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

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

表达式(Expr)基础概念

在 Apache Arrow DataFusion 项目中，Expr（表达式）是表示计算逻辑的核心抽象。它遵循大多数编译器和数据库中常见的"表达式树"抽象模式。理解并掌握Expr的使用对于开发基于 DataFusion 的数据处理应用至关重要。

表达式树结构

表达式在 DataFusion 中以树形结构组织。例如，SQL 表达式 a + b * c 会被表示为：

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

这种树形结构使得复杂表达式的表示和计算变得直观且高效。

Schema 与 DFSchema 详解

在 DataFusion 中处理表达式时，理解数据结构定义至关重要。

Arrow Schema

基础数据结构定义
包含列名和数据类型信息
轻量级结构，适用于单表场景

DataFusion DFSchema

扩展自 Arrow Schema
增加了表名等限定信息
支持多表查询场景
包含函数依赖关系信息

转换方法

从 Schema 到 DFSchema：

DFSchema::try_from_qualified_schema(table_name, original_schema)

从 DFSchema 到 Schema：

let arrow_schema: Schema = df_schema.into();

表达式创建与评估

DataFusion 提供了丰富的 API 来创建和操作表达式：

字面量表达式：lit(5)
列引用：col("my_column")
函数调用：udf.call(vec![expr])
算术运算：col("a") + col("b")

表达式评估通常在查询执行阶段进行，DataFusion 会自动优化和并行化评估过程。

自定义标量函数(UDF)实现

以下是一个完整的标量 UDF 实现示例：

use std::sync::Arc;
use datafusion::arrow::array::{ArrayRef, Int64Array};
use datafusion::common::{Result, cast::as_int64_array};
use datafusion::logical_expr::{ColumnarValue, Volatility, create_udf};

// UDF 实现：对输入值加1
pub fn add_one(args: &[ColumnarValue]) -> Result<ColumnarValue> {
    let args = ColumnarValue::values_to_arrays(args)?;
    let i64s = as_int64_array(&args[0])?;
    
    let new_array = i64s
        .iter()
        .map(|elem| elem.map(|v| v + 1))
        .collect::<Int64Array>();
    
    Ok(ColumnarValue::from(Arc::new(new_array) as ArrayRef))
}

// 注册UDF
let add_one_udf = create_udf(
    "add_one",
    vec![DataType::Int64],  // 输入类型
    DataType::Int64,        // 输出类型
    Volatility::Immutable,  // 函数特性
    Arc::new(add_one),      // 实现
);

// 使用UDF创建表达式
let expr = add_one_udf.call(vec![lit(5)]);      // add_one(5)
let expr = add_one_udf.call(vec![col("col")]);  // add_one(col)

表达式重写技术

表达式重写是 DataFusion 中强大的功能，可用于：

表达式简化
性能优化
类型转换
UDF 内联

基本重写模式

use datafusion::common::tree_node::{Transformed, TreeNode};

fn rewrite_expr(expr: Expr) -> Result<Transformed<Expr>> {
    expr.transform(&|expr| {
        Ok(match expr {
            // 匹配需要重写的表达式模式
            Expr::ScalarFunction(func) if func.name() == "target" => {
                // 构造新表达式
                Transformed::yes(new_expr)
            }
            _ => Transformed::no(expr),  // 不匹配则保留原样
        })
    })
}

优化器规则实现

将重写逻辑封装为优化器规则：

#[derive(Default)]
struct MyOptimizerRule;

impl OptimizerRule for MyOptimizerRule {
    fn name(&self) -> &str { "my_rule" }

    fn rewrite(
        &self,
        plan: LogicalPlan,
        _config: &dyn OptimizerConfig,
    ) -> Result<Transformed<LogicalPlan>> {
        let new_exprs = plan.expressions()
            .into_iter()
            .map(rewrite_expr)
            .collect::<Result<Vec<_>>>()?
            .into_iter()
            .map(|t| t.data)
            .collect();
        
        plan.with_new_exprs(new_exprs, plan.inputs())
            .map(Transformed::yes)
    }
}

表达式类型推断

DataFusion 提供了表达式类型推断功能：

let expr = col("a") + col("b");
let schema = DFSchema::new_with_metadata(
    vec![
        Field::new("a", DataType::Int32, true),
        Field::new("b", DataType::Float32, true),
    ],
    HashMap::new(),
).unwrap();

assert_eq!(
    expr.get_type(&schema).unwrap(),
    DataType::Float32  // 类型提升规则
);