Apache Arrow DataFusion 自定义表提供器开发指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00821/article/details/148464401

Apache Arrow DataFusion 自定义表提供器开发指南

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

概述

Apache Arrow DataFusion 是一个高性能的查询引擎，它允许开发者通过实现特定接口来扩展其功能。本文将详细介绍如何实现自定义表提供器（Custom Table Provider），这是 DataFusion 中连接外部数据源的核心机制。

核心概念

表提供器（TableProvider）

表提供器是 DataFusion 中表示数据源的抽象接口，主要职责包括：

提供数据源的元信息（如表结构）
创建执行计划来读取数据
支持数据写入操作（可选）

执行计划（ExecutionPlan）

执行计划是 DataFusion 查询执行的核心抽象，它定义了如何实际获取数据。表提供器的关键方法 scan 需要返回一个执行计划实例。

实现自定义表提供器

1. 定义数据结构

首先需要定义表示数据源的数据结构。以下是一个用户数据源的示例：

#[derive(Clone, Debug)]
pub struct CustomDataSource {
    inner: Arc<Mutex<CustomDataSourceInner>>,
}

#[derive(Debug)]
struct CustomDataSourceInner {
    data: HashMap<u8, User>,  // 用户数据存储
    bank_account_index: BTreeMap<u64, u8>,  // 银行账号索引
}

2. 实现 ExecutionPlan

执行计划需要实现 ExecutionPlan trait，核心是 execute 方法：

impl ExecutionPlan for CustomExec {
    fn execute(
        &self,
        _partition: usize,
        _context: Arc<TaskContext>,
    ) -> Result<SendableRecordBatchStream> {
        // 1. 从数据源获取数据
        let users: Vec<User> = {
            let db = self.db.inner.lock().unwrap();
            db.data.values().cloned().collect()
        };

        // 2. 构建列数据
        let mut id_array = UInt8Builder::with_capacity(users.len());
        let mut account_array = UInt64Builder::with_capacity(users.len());

        for user in users {
            id_array.append_value(user.id);
            account_array.append_value(user.bank_account);
        }

        // 3. 创建记录批次流
        Ok(Box::pin(MemoryStream::try_new(
            vec![RecordBatch::try_new(
                self.projected_schema.clone(),
                vec![
                    Arc::new(id_array.finish()),
                    Arc::new(account_array.finish()),
                ],
            )?],
            self.schema(),
            None,
        )?))
    }
}

3. 实现 TableProvider

实现 TableProvider trait 的核心方法：

#[async_trait]
impl TableProvider for CustomDataSource {
    fn schema(&self) -> SchemaRef {
        SchemaRef::new(Schema::new(vec![
            Field::new("id", DataType::UInt8, false),
            Field::new("bank_account", DataType::UInt64, true),
        ]))
    }

    async fn scan(
        &self,
        _state: &dyn Session,
        projection: Option<&Vec<usize>>,
        _filters: &[Expr],
        _limit: Option<usize>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        self.create_physical_plan(projection, self.schema()).await
    }
}

高级功能

谓词下推（Predicate Pushdown）

通过实现 supports_filters_pushdown 方法，可以支持谓词下推优化：

fn supports_filters_pushdown(
    &self,
    filters: &[&Expr],
) -> Result<Vec<TableProviderFilterPushDown>> {
    Ok(filters.iter().map(|_| TableProviderFilterPushDown::Exact).collect())
}

谓词下推有三种级别：

Unsupported：不支持下推
Exact：完全下推，数据源保证完全过滤
Inexact：部分下推，DataFusion会进行二次过滤

使用自定义表提供器

注册并使用自定义表提供器：

#[tokio::main]
async fn main() -> Result<()> {
    let ctx = SessionContext::new();
    
    let custom_table_provider = CustomDataSource {
        inner: Arc::new(Mutex::new(CustomDataSourceInner {
            data: Default::default(),
            bank_account_index: Default::default(),
        })),
    };

    // 注册表
    ctx.register_table("customers", Arc::new(custom_table_provider));
    
    // 执行查询
    let df = ctx.sql("SELECT id, bank_account FROM customers").await?;

    Ok(())
}