Apache DataFusion DataFrame API 深度解析与实战指南

支然苹

于 2025-06-06 09:02:16 发布

阅读量341

点赞数 5

CC 4.0 BY-SA版权

本文链接：https://blog.csdn.net/gitblog_00384/article/details/148464418

Apache DataFusion DataFrame API 深度解析与实战指南

前言

Apache DataFusion 是一个高性能的查询执行框架，它提供了 DataFrame API 作为核心接口之一。本文将深入解析 DataFusion 的 DataFrame API，帮助开发者掌握其核心概念和使用方法。

DataFrame 基础概念

什么是 DataFrame？

DataFrame 是 DataFusion 中用于数据处理的核心抽象，它：

借鉴了 Pandas DataFrame 的设计理念
本质上是对 LogicalPlan（逻辑计划）的轻量级封装
提供了构建和执行查询计划的便捷接口

DataFrame 的核心特性

惰性执行：操作不会立即执行，而是在调用执行方法时才进行计算
组合性：支持链式调用，可以方便地组合多个操作
可优化：执行前会进行优化，提高查询效率

DataFrame 创建方式

从内存数据创建

use std::sync::Arc;
use datafusion::prelude::*;
use datafusion::arrow::array::{ArrayRef, Int32Array};
use datafusion::arrow::record_batch::RecordBatch;

let ctx = SessionContext::new();
let data = RecordBatch::try_from_iter(vec![
    ("id", Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef),
    ("bank_account", Arc::new(Int32Array::from(vec![9000, 8000, 7000]))),
])?;

let dataframe = ctx.read_batch(data)?;

从 SQL 查询创建

let ctx = SessionContext::new();
ctx.register_batch("users", data)?;

let dataframe = ctx.sql("SELECT * FROM users;")
    .await?
    .filter(col("bank_account").gt_eq(lit(8000)))?;

DataFrame 操作示例

基本操作链

let dataframe = ctx
    .read_batch(data)?
    .filter(col("bank_account").gt_eq(lit(8000)))?
    .sort(vec![col("bank_account").sort(false, true)])?;

复杂查询构建

let df = ctx.read_csv("data.csv", CsvReadOptions::new()).await?
    .select(vec![col("a"), col("b")])?
    .filter(col("a").gt(lit(5)))?
    .sort(vec![col("b").sort(true, false)])?
    .limit(0, Some(10))?;

执行策略与结果获取

DataFusion 提供了多种执行策略，适应不同场景需求：

1. 全量收集模式

let batches = df.collect().await?;

适用场景：结果集较小，需要全部数据在内存中处理

2. 流式执行模式

let mut stream = df.execute_stream().await?;
while let Some(batch) = stream.next().await {
    // 处理每个批次
}

适用场景：大数据集处理，内存受限环境

3. 缓存模式

let cached_df = df.cache().await?;

适用场景：需要多次使用相同查询结果

数据输出功能

DataFusion 支持将结果写入多种格式：

写入 Parquet 文件

df.write_parquet(
    "output.parquet",
    DataFrameWriteOptions::new(),
    None
).await?;

写入 CSV 文件

df.write_csv(
    "output.csv",
    DataFrameWriteOptions::new()
).await?;

自定义输出格式

通过实现 FileFormat trait 可以支持自定义文件格式

底层原理剖析

DataFrame 与 LogicalPlan 的关系

pub struct DataFrame {
    session_state: Box<SessionState>,
    plan: LogicalPlan,
}

DataFrame 本质上是对 LogicalPlan 的封装，两者可以相互转换：

// DataFrame 转 LogicalPlan
let (state, plan) = df.into_parts();

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

与 LogicalPlanBuilder 的等价性

DataFrame API 和 LogicalPlanBuilder 可以构建相同的逻辑计划：

// 使用 DataFrame API
let df_api = df.select(vec![col("a")])?.filter(col("a").gt(lit(5)))?;

// 使用 LogicalPlanBuilder
let plan = LogicalPlanBuilder::from(plan)
    .project(vec![col("a")])?
    .filter(col("a").gt(lit(5)))?
    .build()?;