Apache Arrow DataFusion DataFrame API 使用指南

Apache Arrow DataFusion DataFrame API 使用指南

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

什么是DataFrame

在Apache Arrow DataFusion项目中，DataFrame是一个基于Rust实现的高级数据操作抽象，其设计灵感来源于Pandas DataFrame。DataFrame本质上是对LogicalPlan（逻辑计划）的一层薄封装，提供了更友好的API来构建和执行查询计划。

从技术实现上看，DataFrame结构体包含两个核心组件：

pub struct DataFrame {
    session_state: SessionState,  // 会话状态信息
    plan: LogicalPlan,           // 底层的逻辑执行计划
}

DataFrame与LogicalPlan的关系

DataFusion提供了两种构建查询计划的方式：

1. DataFrame API：面向用户的友好接口，类似Pandas的操作风格
2. LogicalPlanBuilder：更底层的构建方式，直接操作逻辑计划

这两种方式可以相互转换，实现相同的功能。例如，以下两种方式都会创建一个按id排序的查询：

// 使用DataFrame API
let new_df = df.select(vec![col("id"), col("bank_account")])?
    .sort(vec![col("id")])?;

// 使用LogicalPlanBuilder
let plan = LogicalPlanBuilder::from(&df.to_logical_plan())
    .project(vec![col("id"), col("bank_account")])?
    .sort(vec![col("id")])?
    .build()?;

创建DataFrame的两种方式

1. 从SQL查询创建

let ctx = SessionContext::new();
ctx.register_table("users", Arc::new(create_memtable()?))?;
let dataframe = ctx.sql("SELECT * FROM users;").await?;

2. 使用DataFrame API创建

let ctx = SessionContext::new();
ctx.register_table("users", Arc::new(create_memtable()?))?;
let dataframe = ctx
  .table("users")
  .filter(col("a").lt_eq(col("b")))?
  .sort(vec![col("a").sort(true, true), col("b").sort(false, false)])?;

DataFrame执行模式

DataFusion采用惰性计算模式，只有在显式调用执行方法时才会真正执行查询。这为查询优化提供了空间。

1. 全量收集模式(collect)

let batches = df.collect().await?;  // 将所有结果收集到内存中

适用场景：结果集较小，需要一次性处理所有数据

2. 流式执行模式(execute_stream)

let mut stream = df.execute_stream().await?;
while let Some(rb) = stream.next().await {
    println!("{rb:?}");  // 逐批次处理数据
}

适用场景：大数据集处理，减少内存占用

3. 缓存模式(cache)

let cached_df = df.cache().await?;  // 执行查询并将结果缓存为新DataFrame

适用场景：需要重复使用查询结果

数据导出功能

DataFusion支持将DataFrame结果导出为多种文件格式：

CSV导出示例

dataframe
    .write_csv("output.csv", DataFrameWriteOptions::default(), None)
    .await;

支持的文件格式

1. CSV：适合表格数据的通用交换格式
2. JSON：适合结构化数据的Web应用场景
3. Parquet：列式存储格式，适合大数据分析场景

DataFrame与LogicalPlan的转换

由于DataFrame只是LogicalPlan的封装，两者可以轻松转换：

// DataFrame转LogicalPlan
let plan = dataframe.logical_plan().clone();

// LogicalPlan转DataFrame
let new_df = DataFrame::new(ctx.state(), plan);

这种灵活性使得开发者可以在高级API和底层计划之间自由切换，满足不同场景的需求。

最佳实践建议

1. 对于简单查询，优先使用DataFrame API，代码更简洁
2. 对于复杂查询，可以考虑直接操作LogicalPlan获得更精细的控制
3. 大数据集处理时使用流式执行模式，避免内存溢出
4. 频繁使用的中间结果考虑使用cache()缓存
5. 结果导出时根据下游系统选择合适的文件格式

通过掌握DataFusion的DataFrame API，开发者可以高效地构建复杂的数据处理管道，同时享受Rust语言带来的性能优势。

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