Apache Arrow DataFusion SQL API 使用指南

Apache Arrow DataFusion SQL API 使用指南

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

概述

Apache Arrow DataFusion 是一个高性能的查询引擎，它提供了完整的 SQL 接口，允许开发者使用标准的 SQL 语句与数据进行交互。本文将详细介绍如何使用 DataFusion 的 SQL API 来执行各种数据操作。

核心概念

SessionContext

SessionContext 是 DataFusion 中最高级别的 API 入口点，它提供了执行 SQL 查询的主要接口。使用 SQL API 的基本流程是：

1. 创建 SessionContext 实例
2. 注册数据源（表）
3. 执行 SQL 查询

数据源注册

在 DataFusion 中执行 SQL 查询前，需要先将数据源注册为表。DataFusion 支持多种数据格式，包括 CSV、Parquet 和 Avro。

注册 CSV 文件

use datafusion::error::Result;
use datafusion::prelude::*;

#[tokio::main]
async fn main() -> Result<()> {
    let ctx = SessionContext::new();
    // 注册CSV文件为"example"表
    ctx.register_csv("example", "tests/data/example.csv", CsvReadOptions::new()).await?;
    
    // 执行SQL查询
    let df = ctx.sql("SELECT a, min(b) FROM example WHERE a <= b GROUP BY a LIMIT 100").await?;
    
    // 收集结果
    let results = df.collect().await?;
    
    // 验证结果
    datafusion::assert_batches_eq!(vec![
        "+---+----------------+",
        "| a | min(example.b) |",
        "+---+----------------+",
        "| 1 | 2              |",
        "+---+----------------+",
        ],
        &results
    );
    Ok(())
}

注册 Parquet 文件

Parquet 是一种列式存储格式，DataFusion 提供了专门的注册方法：

use datafusion::error::Result;
use datafusion::prelude::*;

#[tokio::main]
async fn main() -> Result<()> {
    let ctx = SessionContext::new();
    let testdata = datafusion::test_util::parquet_test_data();

    // 注册Parquet文件
    ctx.register_parquet(
        "alltypes_plain",
        &format!("{testdata}/alltypes_plain.parquet"),
        ParquetReadOptions::default(),
    ).await?;

    // 执行复杂查询
    let df = ctx.sql(
        "SELECT int_col, double_col, CAST(date_string_col as VARCHAR) \
         FROM alltypes_plain \
         WHERE id > 1 AND tinyint_col < double_col"
    ).await?;

    // 验证结果
    let results = df.collect().await?;
    datafusion::assert_batches_eq!(vec![
        "+---------+------------+--------------------------------+",
        "| int_col | double_col | alltypes_plain.date_string_col |",
        "+---------+------------+--------------------------------+",
        "| 1       | 10.1       | 03/01/09                       |",
        "| 1       | 10.1       | 04/01/09                       |",
        "| 1       | 10.1       | 02/01/09                       |",
        "+---------+------------+--------------------------------+",
        ],
        &results
    );
    Ok(())
}

注册 Avro 文件

Avro 是另一种流行的数据序列化格式，DataFusion 也提供了支持：

use datafusion::error::Result;
use datafusion::prelude::*;

#[tokio::main]
async fn main() -> Result<()> {
    let ctx = SessionContext::new();
    let testdata = datafusion::test_util::arrow_test_data();
    
    // 注册Avro文件
    let avro_file = &format!("{testdata}/avro/alltypes_plain.avro");
    ctx.register_avro("alltypes_plain", avro_file, AvroReadOptions::default()).await?;

    // 执行查询
    let df = ctx.sql(
        "SELECT int_col, double_col, CAST(date_string_col as VARCHAR) \
         FROM alltypes_plain \
         WHERE id > 1 AND tinyint_col < double_col"
    ).await?;

    // 验证结果
    let results = df.collect().await?;
    datafusion::assert_batches_eq!(vec![
        "+---------+------------+--------------------------------+",
        "| int_col | double_col | alltypes_plain.date_string_col |",
        "+---------+------------+--------------------------------+",
        "| 1       | 10.1       | 03/01/09                       |",
        "| 1       | 10.1       | 04/01/09                       |",
        "| 1       | 10.1       | 02/01/09                       |",
        "+---------+------------+--------------------------------+",
        ],
        &results
    );
    Ok(())
}

使用 SQL 语句注册表

除了编程式注册外，还可以直接使用 SQL 的 CREATE EXTERNAL TABLE 语句来注册数据源：

use datafusion::error::Result;
use datafusion::prelude::*;

#[tokio::main]
async fn main() -> Result<()> {
    let ctx = SessionContext::new();
    let testdata = datafusion::test_util::parquet_test_data();

    // 使用SQL语句注册Parquet文件
    let ddl = format!(
        "CREATE EXTERNAL TABLE alltypes_plain \
        STORED AS PARQUET LOCATION '{testdata}/alltypes_plain.parquet'"
    );
    ctx.sql(&ddl).await?;

    // 执行查询
    let df = ctx.sql(
        "SELECT int_col, double_col, CAST(date_string_col as VARCHAR) \
         FROM alltypes_plain \
         WHERE id > 1 AND tinyint_col < double_col"
    ).await?;

    // 验证结果
    let results = df.collect().await?;
    datafusion::assert_batches_eq!(vec![
        "+---------+------------+--------------------------------+",
        "| int_col | double_col | alltypes_plain.date_string_col |",
        "+---------+------------+--------------------------------+",
        "| 1       | 10.1       | 03/01/09                       |",
        "| 1       | 10.1       | 04/01/09                       |",
        "| 1       | 10.1       | 02/01/09                       |",
        "+---------+------------+--------------------------------+",
        ],
        &results
    );
    Ok(())
}

高级用法

多文件读取

DataFusion 支持将多个文件作为一个表来读取，这是通过 ListingTableProvider 实现的。它会自动匹配各个文件的模式（schema），并将它们合并为一个逻辑表。

// 多文件读取示例（即将推出）

查询选项控制

对于需要更精细控制的情况，可以使用 SessionContext::sql_with_options 方法或直接操作 SessionState API。这些方法允许你：

• 禁用特定的 SQL 语句类型（如 DDL）
• 自定义查询优化器规则
• 调整执行参数

最佳实践

1. 资源管理：确保及时释放不再需要的表和上下文
2. 错误处理：妥善处理可能出现的错误，特别是 I/O 相关的错误
3. 性能考虑：对于大型数据集，考虑使用分区和并行处理
4. 类型转换：注意 SQL 中的类型转换，必要时使用 CAST 函数

总结

DataFusion 的 SQL API 提供了强大而灵活的方式来处理各种数据源。通过本文介绍的方法，你可以轻松地将不同格式的数据注册为表，并使用标准的 SQL 语句进行查询和分析。无论是简单的 CSV 文件还是复杂的 Parquet 数据集，DataFusion 都能提供高效的查询能力。

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