Apache DataFusion 快速入门与集成示例

Apache DataFusion 快速入门，从环境准备、首个查询与结果解读入手，给出 Python 集成 DataFusion 分析 CSV/Parquet、与 pandas/Arrow 互操作及注册自定义函数的样例代码与实测输出。

Apache DataFusionApache DataFusion2026/10/11

Apache DataFusion 快速入门与集成示例

Apache DataFusion 是用 Rust 实现、基于 Apache Arrow 的可扩展查询引擎，可在应用中嵌入并以 SQL 分析 CSV/Parquet/JSON 等数据。本文带你从零完成第一个查询，并给出在 Python 应用中集成的完整示例，所有输出均来自 Apache DataFusion 55.2.0 / Python 绑定 48.0.0 实测环境。

一、快速入门

1.1 环境准备

安装 Python 绑定（以 Ubuntu 22.04 为例，CentOS 8 需用 python3.9 -m pip）：

$ pip3 install -i https://pypi.tuna.tsinghua.edu.cn/simple "datafusion==48.0.0"
Successfully installed datafusion-48.0.0 pyarrow-25.0.1 typing-extensions-4.16.0

验证：

$ python3 -c "import datafusion; print(datafusion.__version__)"
48.0.0

如需命令行体验，可另行 cargo install datafusion-cli（见安装指南），版本为独立的 55.2.0。

1.2 第一个查询

准备数据文件 /root/employees.csv：

id,name,dept_id,salary
1,Alice,10,5000
2,Bob,20,6000
3,Carol,10,7000
4,David,30,8000

运行：

from datafusion import SessionContext
ctx = SessionContext()
ctx.register_csv("employees", "/root/employees.csv", has_header=True)
ctx.sql("SELECT dept_id, COUNT(*) AS cnt, AVG(salary) AS avg_salary "
        "FROM employees GROUP BY dept_id ORDER BY dept_id").show()

真实输出：

+---------+-----+------------+
| dept_id | cnt | avg_salary |
+---------+-----+------------+
| 10      | 2   | 6000.0     |
| 20      | 1   | 6000.0     |
| 30      | 1   | 8000.0     |
+---------+-----+------------+

1.3 输出解读

  • ctx.sql(...) 返回一个惰性 DataFrame，只有调用 show()/collect()/to_pandas() 时才真正执行。
  • 结果以 Arrow RecordBatch 承载，按列式排列；COUNT/AVG 等聚合由引擎自动并行执行。
  • 表头中的 dept_id、cnt、avg_salary 对应 SQL 的投影列与别名。

二、核心概念

2.1 SessionContext、DataFrame 与 RecordBatch

  • SessionContext：会话入口，持有 catalog、注册表与配置，负责解析并规划 SQL。
  • DataFrame：逻辑查询计划的高层封装，支持 select/filter/join/aggregate 等链式算子。
  • RecordBatch：Arrow 的列式数据批次，是算子间数据传递的基本单元（默认每批 8192 行）。

2.2 逻辑计划与物理计划

EXPLAIN 展示计划树。Python 中实测逻辑计划：

$ python3 -c "..."   # EXPLAIN SELECT dept_id, COUNT(*) FROM employees GROUP BY dept_id
plan_type: logical_plan
Projection: employees.dept_id, count(Int64(1)) AS count(*)
  Aggregate: groupBy=[[employees.dept_id]], aggr=[[count(Int64(1))]]
    TableScan: employees projection=[dept_id]

在 datafusion-cli 中 EXPLAIN 默认输出物理计划（含 AggregateExec、RepartitionExec、DataSourceExec 等算子）。

三、集成示例：CSV 数据分析

对一个 CSV 做过滤、聚合、排序与 JOIN：

from datafusion import SessionContext
ctx = SessionContext()
ctx.register_csv("employees", "/root/employees.csv", has_header=True)
ctx.register_csv("depts", "/root/depts.csv", has_header=True)

# 过滤 + 排序
print(ctx.sql("SELECT name, salary FROM employees WHERE salary > 5500 ORDER BY salary")
      .to_pandas().to_string(index=False))

# JOIN
ctx.sql("SELECT e.name, d.dept_name FROM employees e JOIN depts d "
        "ON e.dept_id = d.dept_id ORDER BY e.name").show()

实测输出：

    name  salary
0    Bob    6000
1  Carol    7000
2  David    8000

+-------+-------------+
| name  | dept_name   |
+-------+-------------+
| Alice | Engineering |
| Bob   | Sales       |
| Carol | Engineering |
| David | Marketing   |
+-------+-------------+

3.1 结果落地为表（CTAS）

ctx.sql("CREATE TABLE high_paid AS SELECT name, salary FROM employees WHERE salary >= 7000")
ctx.sql("SELECT * FROM high_paid ORDER BY salary").show()

+-------+--------+
| name  | salary |
+-------+--------+
| Carol | 7000   |
| David | 8000   |
+-------+--------+

四、集成示例：Parquet 与 Arrow/pandas 互操作

4.1 读取 Parquet 并统计

import pyarrow as pa, pyarrow.parquet as pq
from datafusion import SessionContext

tbl = pa.table({"city": ["bj", "sh", "gz"], "sales": [100, 200, 150]})
pq.write_table(tbl, "/root/sales.parquet")

ctx = SessionContext()
ctx.register_parquet("sales", "/root/sales.parquet")
print(ctx.sql("SELECT city, sales FROM sales ORDER BY sales DESC")
      .to_pandas().to_string(index=False))
print("total:", ctx.sql("SELECT SUM(sales) AS total FROM sales").collect()[0].to_pydict())

实测输出：

city  sales
  sh    200
  gz    150
  bj    100
total: {'total': [450]}

4.2 从 pandas / Arrow 直接建表

import pandas as pd, pyarrow as pa
from datafusion import SessionContext

pdf = pd.DataFrame({"k": ["a", "b", "c"], "v": [1, 2, 3]})
ctx = SessionContext()
ctx.register_record_batches("t", [pa.Table.from_pandas(pdf).to_batches()])
print(ctx.sql("SELECT k, v*10 AS v10 FROM t ORDER BY k").to_pandas().to_string(index=False))

k  v10
a   10
b   20
c   30

也可用 ctx.from_pandas(pdf)、ctx.from_arrow(tbl) 直接构造 DataFrame，用 df.to_pandas()、df.to_arrow_table() 回到 Python 生态，实现零拷贝式的列式互操作。

五、集成示例：在应用中注册自定义函数（UDF）

将 Python 函数注册为 SQL 可调用的标量函数，扩展引擎能力：

import pyarrow as pa
from datafusion import SessionContext, udf

def double_it(x):
    return pa.compute.multiply(x, 2)

ctx = SessionContext()
# 签名：udf(func, input_types, return_type, volatility, name)
ctx.register_udf(udf(double_it, [pa.int64()], pa.int64(), "immutable", "double_it"))
ctx.register_record_batches("nums", [pa.table({"n": [1, 2, 3]}).to_batches()])
print(ctx.sql("SELECT n, double_it(n) AS doubled FROM nums ORDER BY n")
      .to_pandas().to_string(index=False))

实测输出：

 n  doubled
 1        2
 2        4
 3        6

除标量 UDF（register_udf）外，还可用 register_udaf（聚合）、register_udtf（表函数）、register_udwf（窗口函数）扩展引擎。

六、常见问题

Q1：为什么 SHOW 或 information_schema 查询报错？ 未开启 information_schema 时，SHOW [VARIABLE] 会报 SHOW [VARIABLE] is not supported unless information_schema is enabled。用 SessionConfig().with_information_schema(True) 创建会话即可。

Q2：EXPLAIN 通过 Python show() 为什么报 "Explain must be root of the plan"？ show() 会在计划外围再包一层算子，而 EXPLAIN 必须位于计划根部。改用 ctx.sql("EXPLAIN ...").collect() 读取结果，或使用 DataFrame.explain()。

Q3：cargo install datafusion-cli 很慢怎么办？ 首次需从源码编译整棵依赖树（12 核约 15 分钟）。可配置 crates.io 镜像加速下载，编译只需一次。

Q4：CentOS 8 安装时 pyarrow 报 cmake 缺失？ 说明命中了源码构建。固定带 cp39 预编译包的 pyarrow==17.0.0 即可，无需安装 cmake。