DataFusion:用于 Apache Arrow 的 Rust 原生查询引擎


已发布 2019 年 2 月 4 日
作者 Andy Grove (agrove)

我们很高兴地宣布,DataFusion 已捐赠给 Apache Arrow 项目。 DataFusion 是 Apache Arrow 的 Rust 实现的内存查询引擎。

虽然 DataFusion 在两年前就开始了,但最近重新实现为 Arrow 原生,目前功能有限,但支持针对 RecordBatch 迭代器的 SQL 查询,并支持 CSV 文件。 计划添加对 Parquet 文件的支持

SQL 支持仅限于投影(SELECT)、选择(WHERE)和简单的聚合(MINMAXSUM),以及可选的GROUP BY子句。

支持的表达式包括标识符、文字、简单数学运算(+-*/)、二元表达式(ANDOR)、相等和比较运算符(=!=<<=>=>)和 CAST(expr AS type)

例子

以下示例演示了针对 CSV 文件运行一个简单的聚合 SQL 查询。

// create execution context
let mut ctx = ExecutionContext::new();

// define schema for data source (csv file)
let schema = Arc::new(Schema::new(vec![
    Field::new("c1", DataType::Utf8, false),
    Field::new("c2", DataType::UInt32, false),
    Field::new("c3", DataType::Int8, false),
    Field::new("c4", DataType::Int16, false),
    Field::new("c5", DataType::Int32, false),
    Field::new("c6", DataType::Int64, false),
    Field::new("c7", DataType::UInt8, false),
    Field::new("c8", DataType::UInt16, false),
    Field::new("c9", DataType::UInt32, false),
    Field::new("c10", DataType::UInt64, false),
    Field::new("c11", DataType::Float32, false),
    Field::new("c12", DataType::Float64, false),
    Field::new("c13", DataType::Utf8, false),
]));

// register csv file with the execution context
let csv_datasource =
    CsvDataSource::new("test/data/aggregate_test_100.csv", schema.clone(), 1024);
ctx.register_datasource("aggregate_test_100", Rc::new(RefCell::new(csv_datasource)));

let sql = "SELECT c1, MIN(c12), MAX(c12) FROM aggregate_test_100 WHERE c11 > 0.1 AND c11 < 0.9 GROUP BY c1";

// execute the query
let relation = ctx.sql(&sql).unwrap();
let mut results = relation.borrow_mut();

// iterate over the results
while let Some(batch) = results.next().unwrap() {
    println!(
        "RecordBatch has {} rows and {} columns",
        batch.num_rows(),
        batch.num_columns()
    );

    let c1 = batch
        .column(0)
        .as_any()
        .downcast_ref::<BinaryArray>()
        .unwrap();

    let min = batch
        .column(1)
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    let max = batch
        .column(2)
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    for i in 0..batch.num_rows() {
        let c1_value: String = String::from_utf8(c1.value(i).to_vec()).unwrap();
        println!("{}, Min: {}, Max: {}", c1_value, min.value(i), max.value(i),);
    }
}

路线图

DataFusion 的路线图将取决于 Rust 社区的兴趣,但以下是一些计划中的短期项目

  • 扩展现有功能的测试覆盖率
  • 添加对 Parquet 数据源的支持
  • 实现更多 SQL 功能,例如JOINORDER BYLIMIT
  • 实现 DataFrame API 作为 SQL 的替代方案
  • 使用 Rust 的 async 和 await 功能添加对分区和并行查询执行的支持
  • 创建一个 Docker 镜像,以便轻松地将 DataFusion 用作交互式和批量查询的独立查询工具

欢迎贡献者!

如果您很高兴能够使用 Rust 进行数据科学,并希望为这项工作做出贡献,那么有很多方法可以参与其中。 最简单的入门方法是针对您自己的数据源试用 DataFusion 并提交您发现的任何问题的错误报告。 您还可以查看当前的问题列表并尝试修复一个。 您还可以加入用户邮件列表以提出问题。