Apache Arrow DataFusion原理與架構

金筆書生呂落第發表於2023-05-15
Apache架構

本篇主要介紹了一種使用Rust語言編寫的查詢引擎——DataFusion，其使用了基於Arrow格式的記憶體模型，結合Rust語言本身的優勢，達成了非常優秀的效能指標

DataFusion是一個查詢引擎而非資料庫，因此其本身不具備儲存資料的能力。但正因為不依賴底層儲存的格式，使其成為了一個靈活可擴充套件的查詢引擎。它原生支援了查詢CSV，Parquet，Avro，Json等儲存格式，也支援了本地，AWS S3，Azure Blob Storage，Google Cloud Storage等多種資料來源。同時還提供了豐富的擴充套件介面，可以方便的讓我們接入自定義的資料格式和資料來源。

DataFusion具有以下特性：

  • 高效能：基於Rust，不用進行垃圾回收；基於Arrow記憶體模型，列式儲存，方便向量化計算
  • 連線簡單：能夠與Arrow的其他生態互通
  • 整合和定製簡單：可以擴充套件資料來源，方法和運算元等
  • 完全基於Rust編寫：高質量

基於DataFusion我們可以輕鬆構建高效能、高質量、可擴充套件的資料處理系統。

DBMS 與 Query Engine 的區別

DBMS: DataBase Management System

DBMS是一個包含完整資料庫管理特性的系統，主要包含以下幾個模組：

  • 儲存系統
  • 後設資料(Catalog)
  • 查詢引擎(Query Engine)
  • 訪問控制和許可權
  • 資源管理
  • 管理工具
  • 客戶端
  • 多節點管理

Query Engine

DataFusion是一種查詢引擎，查詢引擎屬於資料庫管理系統的一部分。查詢引擎是使用者與資料庫互動的主要介面，主要作用是將面向使用者的高階查詢語句翻譯成可被具體執行的資料處理單元操作，然後執行操作獲取資料。

DataFusion架構

架構詳情

image

DataFusion查詢引擎主要由以下幾部分構成：

  1. 前端
    • 語法解析
    • 語義分析
    • Planner：語法樹轉換成邏輯計劃

主要涉及DFParser和SqlToRel這兩個struct

  1. 查詢中間表示
    • Expression（表示式）/ Type system（型別系統）
    • Query Plan / Relational Operators（關係運算元）
    • Rewrites / Optimizations（邏輯計劃最佳化）

主要涉及LogicalPlan和Expr這兩個列舉類

  1. 查詢底層表示
    • Statistics（物理計劃運算元的統計資訊，輔助物理計劃最佳化）
    • Partitions（分塊，多執行緒執行物理計劃運算元）
    • Sort orders（物理計劃運算元對資料是否排序）
    • Algorithms（物理計劃運算元的執行演算法，如Hash join和Merge join）
    • Rewrites / Optimizations（物理計劃最佳化）

主要涉及PyhsicalPlanner這個trait實現的邏輯計劃到物理計劃的轉換，其中主要的關鍵點是ExecutionPlan和PhysicalExpr

  1. 執行執行時（運算元）
    • 分配資源
    • 向量化計算

主要涉及所有執行運算元，如GroupedHashAggregateStream

擴充套件點

DataFusion查詢引擎的架構還是比較簡單的，其中的擴充套件點也非常清晰，我們可以從以下幾個方面對DataFusion進行擴充套件：

使用者自定義函式UDF

無狀態方法

/// 邏輯表示式列舉類
pub enum Expr {
    ...
    ScalarUDF {
        /// The function
        fun: Arc<ScalarUDF>,
        /// List of expressions to feed to the functions as arguments
        args: Vec<Expr>,
    },
    ...
}
/// UDF的邏輯表示式
pub struct ScalarUDF {
    /// 方法名
    pub name: String,
    /// 方法簽名
    pub signature: Signature,
    /// 返回值型別
    pub return_type: ReturnTypeFunction,
    /// 方法實現
    pub fun: ScalarFunctionImplementation,
}
/// UDF的物理表示式
pub struct ScalarFunctionExpr {
    fun: ScalarFunctionImplementation,
    name: String,
    /// 參數列達式列表
    args: Vec<Arc<dyn PhysicalExpr>>,
    return_type: DataType,
}

使用者自定義聚合函式UADF

有狀態方法

/// 邏輯表示式列舉類
pub enum Expr {
    ...
    AggregateUDF {
        /// The function
        fun: Arc<AggregateUDF>,
        /// List of expressions to feed to the functions as arguments
        args: Vec<Expr>,
        /// Optional filter applied prior to aggregating
        filter: Option<Box<Expr>>,
    },
    ...
}
/// UADF的邏輯表示式
pub struct AggregateUDF {
    /// 方法名
    pub name: String,
    /// 方法簽名
    pub signature: Signature,
    /// 返回值型別
    pub return_type: ReturnTypeFunction,
    /// 方法實現
    pub accumulator: AccumulatorFunctionImplementation,
    /// 需要儲存的狀態的型別
    pub state_type: StateTypeFunction,
}
/// UADF的物理表示式
pub struct AggregateFunctionExpr {
    fun: AggregateUDF,
    args: Vec<Arc<dyn PhysicalExpr>>,
    data_type: DataType,
    name: String,
}

使用者自定義最佳化規則

Optimizer定義了承載最佳化規則的結構體，其中optimize方法實現了邏輯計劃最佳化的過程。最佳化規則列表中的每個最佳化規則會被以TOP-DOWN或BOTTOM-UP方式作用於邏輯計劃樹，最佳化規則列表會被實施多個輪次。我們可以透過實現OptimizerRule這個trait來實現自己的最佳化邏輯。

pub struct Optimizer {
    /// All rules to apply
    pub rules: Vec<Arc<dyn OptimizerRule + Send + Sync>>,
}

pub trait OptimizerRule {
    /// Try and rewrite `plan` to an optimized form, returning None if the plan cannot be
    /// optimized by this rule.
    fn try_optimize(
        &self,
        plan: &LogicalPlan,
        config: &dyn OptimizerConfig,
    ) -> Result<Option<LogicalPlan>>;

    ...
}

使用者自定義邏輯計劃運算元

/// 邏輯計劃運算元列舉類
pub enum LogicalPlan {
    ...
    Extension(Extension),
    ...
}
/// 自定義邏輯計劃運算元
pub struct Extension {
    /// The runtime extension operator
    pub node: Arc<dyn UserDefinedLogicalNode>,
}
/// 自定義邏輯計劃運算元需要實現的trait
pub trait UserDefinedLogicalNode: fmt::Debug + Send + Sync { ... }

使用者自定義物理計劃運算元

/// 為自定義的邏輯計劃運算元`UserDefinedLogcialNode`生成對應的物理計劃運算元
pub trait ExtensionPlanner {
    async fn plan_extension(
        &self,
        planner: &dyn PhysicalPlanner,
        node: &dyn UserDefinedLogicalNode,
        logical_inputs: &[&LogicalPlan],
        physical_inputs: &[Arc<dyn ExecutionPlan>],
        session_state: &SessionState,
    ) -> Result<Option<Arc<dyn ExecutionPlan>>>;
}
/// DataFusion預設的邏輯計劃到物理計劃的轉換器提供了自定義轉換過程的結構體
pub struct DefaultPhysicalPlanner {
    extension_planners: Vec<Arc<dyn ExtensionPlanner + Send + Sync>>,
}
/// 自定義物理計劃運算元需要實現的trait
pub trait ExecutionPlan: Debug + Send + Sync { ... }

使用者自定義資料來源

可以看出，自定義資料來源其實就是生成一個對應的ExecutionPlan執行計劃，這個執行計劃實施的是掃表的任務。如果資料來源支援下推的能力，我們在這裡可以將projection filters limit等操作下推到掃表時。

/// 自定義資料來源需要實現的trait
pub trait TableProvider: Sync + Send {
    ...
    async fn scan(
        &self,
        state: &SessionState,
        projection: Option<&Vec<usize>>,
        filters: &[Expr],
        limit: Option<usize>,
    ) -> Result<Arc<dyn ExecutionPlan>>;
    ...
}

使用者自定義後設資料

pub trait CatalogProvider: Sync + Send {
    ...
	
    /// 根據名稱獲取Schema
    fn schema(&self, name: &str) -> Option<Arc<dyn SchemaProvider>>;
    /// 註冊Schema
    fn register_schema(
        &self,
        name: &str,
        schema: Arc<dyn SchemaProvider>,
    ) -> Result<Option<Arc<dyn SchemaProvider>>> {
        // use variables to avoid unused variable warnings
        let _ = name;
        let _ = schema;
        Err(DataFusionError::NotImplemented(
            "Registering new schemas is not supported".to_string(),
        ))
    }
}

pub trait SchemaProvider: Sync + Send {
    ...
    /// 根據表名獲取資料來源
    async fn table(&self, name: &str) -> Option<Arc<dyn TableProvider>>;
    /// 註冊資料來源
    fn register_table(
        &self,
        name: String,
        table: Arc<dyn TableProvider>,
    ) -> Result<Option<Arc<dyn TableProvider>>> {
        Err(DataFusionError::Execution(
            "schema provider does not support registering tables".to_owned(),
        ))
    }
    ...
}

邏輯計劃（LogicalPlan）

邏輯計劃其實就是資料流圖，資料從葉子節點流向根節點

let df: DataFrame = ctx.read_table("http_api_requests_total")?
            .filter(col("path").eq(lit("/api/v2/write")))?
            .aggregate([col("status")]), [count(lit(1))])?;

這裡我們就使用DataFusion的API介面構造了一個資料流，首先read_table節點會從資料來源中掃描資料到記憶體中，然後經過filter節點按照條件進行過濾，最後經過aggregate節點進行聚合。資料流過最後的節點時，就生成了我們需要的資料。

上述鏈式呼叫的API介面實際上並沒有真正執行對資料的操作，這裡實際上是使用了建造者模式構造了邏輯計劃樹。最終生成的DataFrame實際上只是包含了一下資訊：

pub struct DataFrame {
    /// 查詢上下文資訊，包含了後設資料，使用者註冊的UDF和UADF，使用的最佳化器，使用的planner等資訊
    session_state: SessionState,
    /// 邏輯計劃樹的根節點
    plan: LogicalPlan,
}

支援的邏輯計劃運算元

點選檢視程式碼 
Projection
Filter
Window
Aggregate
Sort
Join
TableScan

Repartition
Union
Subquery
Limit
Extension
Distinct

Values
Explain
Analyze
SetVariable
Prepare
Dml(...)

CreateExternalTable
CreateView
CreateCatalogSchema
CreateCatalog
DropTable
DropView

邏輯計劃最佳化

目標：確保結果相同的情況下，執行更快

image

初始的邏輯計劃，需要經過多個輪次的最佳化，才能生成執行效率更高的邏輯計劃。DataFusion本身的最佳化器內建了很多最佳化規則，使用者也可以擴充套件自己的最佳化規則。

內建最佳化輪次

  1. 下推（Pushdown）：減少從一個節點到另一個節點的資料的行列數

    • PushDownProjection
    • PushDownFilter
    • PushDownLimit

  2. 簡化（Simplify）：簡化表示式，減少執行時的運算。例如使用布林代數的法則，將b > 2 AND b > 2簡化成b > 2。

    • SimplifyExpressions
    • UnwrapCastInComparison

  3. 簡化（Simplify）：刪除無用的節點

  4. 平鋪子查詢（Flatten Subqueries）：將子查詢用join重寫

    • DecorrelateWhereExists
    • DecorrelatedWhereIn
    • ScalarSubqueryToJoin

  5. 最佳化join：識別join謂詞

    • ExtractEqualJoinPredicate
    • RewriteDisjunctivePredicate
    • FilterNullJoinKeys

  6. 最佳化distinct

    • SingleDistinctToGroupBy
    • ReplaceDistinctWithAggregate

表示式運算（Expression Evaluation）

假設現在有這樣一個謂詞表示式

path = '/api/v2/write' or path is null

經過語法解析和轉換後，可以用如下表示式樹表示：

image

DataFusion在實施表示式運算時，使用了Arrow提供的向量化計算方法來加速運算

image

物理計劃（ExecutionPlan）

image

呼叫DataFusion提供的DefaultPhysicalPlanner中的create_physical_plan方法，可以將邏輯計劃樹轉換成物理計劃樹。其中物理計劃樹中的每個節點都是一個ExecutionPlan。執行物理計劃樹時，會從根節點開始呼叫execute方法，呼叫該方法還沒有執行對資料的操作，僅僅是將每個物理計劃運算元轉換成一個RecordBatchStream運算元，形成資料流運算元樹。這些RecordBatchStream運算元都實現了future包提供的Stream特性，當我們最終呼叫RecordBatchStream的collect方法時，才會從根節點開始poll一次來獲取一下輪要處理的資料，根節點的poll方法內會呼叫子節點的poll方法，最終每poll一次，整棵樹都會進行一次資料從葉子節點到根節點的流動，生成一個RecordBatch。

image

DataFusion實現的物理計劃運算元具有以下特性：

  • 非同步：避免了阻塞I/O
  • 流式：資料是流式處理的
  • 向量化：每次可以向量化地處理一個RecordBatch
  • 分片：每個運算元都可以並行，可以產生多個分片
  • 多核

相關文章