risinglightDB 核心代码学习笔记

1 整体文件列表

.
├── CODE_OF_CONDUCT.md
├── CONTRIBUTING.md
├── Cargo.toml
├── GOVERNANCE.md
├── LICENSE
├── Makefile
├── README.md
├── benches # 对于一些宏的定义
│   ├── array.rs
│   └── e2e.rs
├── build.rs
├── clippy.toml
├── docs
│   └── #...
├── proto # 对于risinglight数据库物理存储格式的定义
│   ├── Cargo.toml
│   ├── README.md
│   ├── build.rs
│   └── src
│       ├── lib.rs
│       └── proto
│           └── rowset.proto
├── pyproject.toml
├── rust-toolchain
├── rustfmt.toml
├── src # 核心代码
│   ├── array # 内存数据库数据对象
│   │   ├── data_chunk.rs
│   │   ├── data_chunk_builder.rs
│   │   ├── internal_ext.rs
│   │   ├── mod.rs
│   │   ├── ops.rs
│   │   ├── primitive_array.rs
│   │   ├── shuffle_ext.rs
│   │   └── utf8_array.rs
│   ├── binder # SqlStatement绑定器, 将sqlparser的Statement类绑定到对应执行函数
│   │   ├── copy.rs
│   │   ├── create_table.rs
│   │   ├── delete.rs
│   │   ├── drop.rs
│   │   ├── expr.rs
│   │   ├── insert.rs
│   │   ├── mod.rs
│   │   ├── select.rs
│   │   └── table.rs
│   ├── catalog # DB执行的元信息, 定义Database->Schema->Table->Column->Types
│   │   ├── column.rs
│   │   ├── mod.rs
│   │   ├── root.rs
│   │   ├── schema.rs
│   │   └── table.rs
│   ├── db.rs # DB sql执行入口, 传入SqlString返回结果
│   ├── executor # 所有执行模块的入口, 包括Sql执行的入口, 物理存储的入口等
│   │   ├── copy_from_file.rs
│   │   ├── copy_to_file.rs
│   │   ├── create.rs
│   │   ├── delete.rs
│   │   ├── drop.rs
│   │   ├── evaluator.rs
│   │   ├── explain.rs
│   │   ├── filter.rs
│   │   ├── hash_agg.rs
│   │   ├── hash_join.rs
│   │   ├── insert.rs
│   │   ├── internal.rs
│   │   ├── limit.rs
│   │   ├── merge_join.rs
│   │   ├── mod.rs
│   │   ├── nested_loop_join.rs
│   │   ├── order.rs
│   │   ├── projection.rs
│   │   ├── simple_agg.rs
│   │   ├── sort_agg.rs
│   │   ├── table_scan.rs
│   │   ├── top_n.rs
│   │   ├── values.rs
│   │   └── window.rs
│   ├── lib.rs
│   ├── main.rs
│   ├── parser
│   │   └── mod.rs
│   ├── planner # 接入Egg的核心流程
│   │   ├── cost.rs
│   │   ├── explain.rs
│   │   ├── mod.rs
│   │   └── rules # 核心优化规则
│   │       ├── agg.rs
│   │       ├── expr.rs
│   │       ├── mod.rs
│   │       ├── order.rs
│   │       ├── plan.rs
│   │       ├── range.rs
│   │       ├── rows.rs
│   │       ├── schema.rs
│   │       └── type_.rs
│   ├── python # Python绑定接口
│   │   └── mod.rs
│   ├── server # 本地DB服务器接口
│   │   ├── mod.rs
│   │   └── processor.rs
│   ├── storage # 数据库存储操作类
│   │   ├── chunk.rs
│   │   ├── error.rs
│   │   ├── memory # 内存DB操作类
│   │   │   ├── iterator.rs
│   │   │   ├── mod.rs
│   │   │   ├── row_handler.rs
│   │   │   ├── table.rs
│   │   │   └── transaction.rs
│   │   ├── mod.rs
│   │   └── secondary # 硬盘存储DB的数据操作类
│   │       ├── block
│   │       │   ├── blob_block_builder.rs
│   │       │   ├── blob_block_iterator.rs
│   │       │   ├── block_index_builder.rs
│   │       │   ├── char_block_builder.rs
│   │       │   ├── char_block_iterator.rs
│   │       │   ├── dict_block_builder.rs
│   │       │   ├── dict_block_iterator.rs
│   │       │   ├── fake_block_iterator.rs
│   │       │   ├── nullable_block_builder.rs
│   │       │   ├── nullable_block_iterator.rs
│   │       │   ├── primitive_block_builder.rs
│   │       │   ├── primitive_block_iterator.rs
│   │       │   ├── rle_block_builder.rs
│   │       │   └── rle_block_iterator.rs
│   │       ├── block.rs
│   │       ├── checksum.rs
│   │       ├── column
│   │       │   ├── blob_column_builder.rs
│   │       │   ├── blob_column_factory.rs
│   │       │   ├── char_column_builder.rs
│   │       │   ├── char_column_factory.rs
│   │       │   ├── column_builder.rs
│   │       │   ├── column_iterator.rs
│   │       │   ├── concrete_column_iterator.rs
│   │       │   ├── primitive_column_builder.rs
│   │       │   ├── primitive_column_factory.rs
│   │       │   └── row_handler_column_iterator.rs
│   │       ├── column.rs
│   │       ├── compactor.rs
│   │       ├── concat_iterator.rs
│   │       ├── delete_vector.rs
│   │       ├── encode.rs
│   │       ├── index.rs
│   │       ├── index_builder.rs
│   │       ├── manifest.rs
│   │       ├── merge_iterator.rs
│   │       ├── mod.rs
│   │       ├── options.rs
│   │       ├── row_handler.rs
│   │       ├── rowset
│   │       │   ├── disk_rowset.rs
│   │       │   ├── encoded.rs
│   │       │   ├── mem_rowset.rs
│   │       │   ├── mod.rs
│   │       │   ├── rowset_builder.rs
│   │       │   ├── rowset_iterator.rs
│   │       │   └── rowset_writer.rs
│   │       ├── statistics
│   │       │   ├── distinct_value.rs
│   │       │   ├── row_count.rs
│   │       │   └── statistics_builder.rs
│   │       ├── statistics.rs
│   │       ├── storage.rs
│   │       ├── table.rs
│   │       ├── tests.rs
│   │       ├── transaction.rs
│   │       ├── transaction_manager.rs
│   │       ├── txn_iterator.rs
│   │       └── version_manager.rs
│   ├── types
│   │   ├── blob.rs
│   │   ├── date.rs
│   │   ├── interval.rs
│   │   ├── mod.rs
│   │   ├── native.rs
│   │   ├── timestamp.rs
│   │   └── value.rs
│   └── utils
│       ├── mod.rs
│       └── time
│           ├── duration.rs
│           └── mod.rs
└── tests
    ├── planner_test # 优化器测试
    │   └── # ... 
    ├── sql # sql测试
    │   └── # ...
    ├── sqllogictest # 逻辑优化器测试
    │   └── # ...
    ├── sqlplannertest # 物理规划器测试
    │   └── # ...
    └── tpch-full.sh # 全测试脚本

大致流程:

Sql String进入Database类
经过parser类: Sql -> Statement
经过Binder类: Statement -> egg::RecExpr<crate::planner::Expr>
由Optimizer进行优化: egg::RecExpr<crate::planner::Expr> -> egg::RecExpr<crate::planner::Expr>
由Executor进行计划: egg::RecExpr<crate::planner::Expr> --crate::executor::build--> std::pin::Pin<Box<dyn Stream<Item = Result<DataChunk, crate::executor::ExecutorError>> + Send>>
执行执行计划: 执行计划 --executor.try_collect--> Vec<DataChunk>
进行结果的打印

2 Egg优化类型转化逻辑

risinglight/src/binder
├── copy.rs
├── create_table.rs
├── delete.rs
├── drop.rs
├── expr.rs
├── insert.rs
├── mod.rs
├── select.rs
└── table.rs

对于mod.rs中bind函数的学习:

其中使用了egg::Extractor对于转化好的e-graph进行了以AstSize作为损失函数的提取
- 相当于这一步是做了Statement到e-graph的转化, 并且做了简单的简化
其主函数的核心目的是将statement绑定到egraph中并且根据cost生成一个best的RecExpr

其中的核心优化逻辑在对于Statement::Query类型的Statement进行绑定之后

 fn bind_stmt(&mut self, stmt: Statement) -> Result {
        match stmt {
            //....
            Statement::Query(query) => 
            	self.bind_query(*query) // 返回值为 Result<(Id, Context), BindError>
            		// Context类型中存储的是档期可访问的tableName, columnName, 
            		// 和命令外可以访问的字段
            		.map(|(id, _)| id), // 这是对于Reslt进行映射的函数, (id, _) -> {return id;}
            // 将传入的Statement中的query字段进行了绑定,然后构成了map
            //....
        }
    }
    pub(super) fn bind_query(&mut self, query: Query) -> Result<(Id, Context)> {
        // 处理context
        self.contexts.push(Context::default());
        // 处理查询逻辑, 并将其输入Binder的e-graph
        // 对于Statement中的每个节点, 通过绑定的形式将其添加到e-graph中
        // 相当于做的是一个手动的深拷贝, 将statement节点替换为自定义的符合e-graph要求的类型的节点 -- 其节点类型定义写在planner包的mod文件中
        let ret = self.bind_query_internal(query);
        // 处理context
        let ctx = self.contexts.pop().unwrap();
        ret.map(|id| (id, ctx))
    }

对于主流程的分析

    /// Run SQL queries and return the outputs.
    pub async fn run(&self, sql: &str) -> Result<Vec<Chunk>, Error> {
        // ... 前置字符串处理

        let optimizer = crate::planner::Optimizer::new(
            // ... 创建了一个新的优化器
        );

        // sql -> Statement
        let stmts: Vec<sqlparser::ast::Statement> = parse(sql)?;
        let mut outputs: Vec<Chunk> = vec![];
        for stmt in stmts {
            // Statement -> RecExpr<Expr>
            let mut binder: crate::binder::Binder = crate::binder::Binder::new(self.catalog.clone());
            let bound: egg::RecExpr<crate::planner::Expr> = binder.bind(stmt)?;
            // RecExpr<Expr> -> RecExpr<Expr> 本身的优化
            let optimized: egg::RecExpr<crate::planner::Expr> = optimizer.optimize(&bound);
            // 物理执行器
            let executor: std::pin::Pin<Box<dyn Stream<Item = Result<DataChunk, crate::executor::ExecutorError>> + Send>> = match self.storage.clone() {
                //... 执行物理查询或写入
            };
            // 回收结果
            let output: Vec<DataChunk> = executor.try_collect().await?;
            let chunk: Chunk = Chunk::new(output);
            // 结果输出
            outputs.push(chunk);
        }
        Ok(outputs)
    }
}

3 核心Egg优化逻辑子包

/risinglight/src/planner
├── cost.rs
├── explain.rs
├── mod.rs
└── rules
    ├── agg.rs
    ├── expr.rs
    ├── mod.rs
    ├── order.rs
    ├── plan.rs
    ├── range.rs
    ├── rows.rs
    ├── schema.rs
    └── type_.rs

分为四个大的模块:

mod.rs: 定义了基础类型和优化入口

Expr该类型是实际在e-graph中存储的数据类型. 该数据类型被定义在了egg::define_language!宏中, 用于高效地定义节点间的依赖关系

Optimizer该类型就是全局进行egg优化的客户端, 提供了三个函数

// 构造函数
pub fn new(catalog: RootCatalogRef, config: Config) -> Self; 
// 执行优化的主函数
pub fn optimize(&self, expr: &RecExpr) -> RecExpr; 
// 对一条RecExpr路径上的所有节点计算cost
// 被用于在进行Explain的时候展示优化的策略
pub fn costs(&self, expr: &RecExpr) -> Vec<f32>;

explain.rs: 定义了Explain类型用于在RecExpr上进行包装并在前端展示, 相当于是一个可视化模块

cost.rs: 定义了CostFn类型用于进行最终e-graph的择优和Explain中显示每个RecExpr的cost

官方提供的一个cost函数的示例

struct SillyCostFn; // cost函数类型
impl CostFunction<SymbolLang> for SillyCostFn {
    // 定义了当前的cost的数据类型是一个double
    type Cost = f64;
    // 定义了cost函数为一个泛型函数
    fn cost<C>(
        &mut self, // 表明其是一个实例函数
        enode: &SymbolLang, // 由自定义Symbo形成的语法树
        mut costs: C) // 一个costs函数指针 -- 应该是用于进行递归
		-> Self::Cost // 返回值为一个Cost类型的结果
    where
        C: FnMut(Id) -> Self::Cost // 这个函数指针的输入是egg::Id, 输出是Cost
    {
        // 这里没有使用传入的函数指针, 而是直接mock了三个cost
        let op_cost = match enode.op.as_str() {
            "foo" => 100.0,
            "bar" => 0.7,
            _ => 1.0
        };
        enode.fold(op_cost, |sum, id| sum + costs(id))
    }
}

这里的cost函数实现:

impl egg::CostFunction<Expr> for CostFn<'_> {
    type Cost = f32;
    fn cost<C>(&mut self, enode: &Expr, mut costs: C) -> Self::Cost
    where
        C: FnMut(Id) -> Self::Cost,
    {
        use Expr::*;
        // 获取当前需要计算cost的图的Id
        let id: &Id = &self.egraph.lookup(enode.clone()).unwrap();
        // 构造一个新的costs函数, 其等于直接调用传入的costs函数
        let mut costs = |i: &Id| costs(*i);
        // 根据一个Id获取其输出的行数
        let rows = |i: &Id| self.egraph[*i].data.rows;
        // 根据一个Id获取其输出的列数
        let cols = |i: &Id| self.egraph[*i].data.schema.len() as f32;
        // 定义了 n * log(n) 的函数
        let nlogn = |x: f32| x * (x + 1.0).log2();
        // The cost of output chunks of a plan.
        // 定义了输出数据的复杂度, 等于输出的行数 * 输出的列数
        let out = || rows(id) * cols(id);

        let c: f32 = match enode {
            // 如果匹配到 Scan 和 Values节点, 返回数据输出的复杂度
            Scan(_) | Values(_) => out(),
            // 如果匹配到 Order 节点, 
            // 返回使用快排的 nlogn 子语句时间复杂度 + 数据输出复杂度 + 子语句的cost
            Order([_, c]) => nlogn(rows(c)) + out() + costs(c),
            // 匹配Filter节点,
            // 返回 (内部条件节点的cost) * 子语句结果行数 + 数据输出复杂度 + 子语句cost
            Filter([exprs, c]) => costs(exprs) * rows(c) + out() + costs(c),
            // 匹配Proj节点或window节点
            // 返回 (内部条件节点复杂度) * 子语句行数 + 子语句cost
            Proj([exprs, c]) | Window([exprs, c]) => costs(exprs) * rows(c) + costs(c),
            // 匹配到Agg节点
            // 返回 (内部条件节点复杂度) * 子语句行数 + 子语句cost
            Agg([exprs, c]) => costs(exprs) * rows(c) + out() + costs(c),
            // 匹配到HashAgg节点
            // 返回 (log2(当前行数) + 内部条件节点复杂度 + groupBy节点复杂度) * 子语句行数 + 输出复杂度 + 子语句cost
            HashAgg([exprs, groupby, c]) => {
                ((rows(id) + 1.0).log2() + costs(exprs) + costs(groupby)) * rows(c)
                    + out()
                    + costs(c)
            }
            // SortAgg: (内部条件节点复杂度 + groupBy节点复杂度) * 子语句行数 + 输出复杂度 + 子语句cost
            SortAgg([exprs, groupby, c]) => {
                (costs(exprs) + costs(groupby)) * rows(c) + out() + costs(c)
            }
            // Limit: 输出复杂度 + 之前的cost + 输出复杂度 + 子语句cost
            Limit([_, _, c]) => out() + costs(c),
            // TopN: log2(行数) * 子语句行数 + 输出复杂度 + 子语句复杂度
            TopN([_, _, _, c]) => (rows(id) + 1.0).log2() * rows(c) + out() + costs(c),
            // Join: on子句复杂度 * 左子句输出长度 * 右子句输出长度 + 输出复杂度 + 左子句复杂度 + 右子句复杂度
            Join([_, on, l, r]) => costs(on) * rows(l) * rows(r) + out() + costs(l) + costs(r),
            // HashJoin: log2(左子句长度) * (左右子句长度和) + 输出复杂度 + 左右子句的复杂度
            HashJoin([_, _, _, l, r]) => {
                (rows(l) + 1.0).log2() * (rows(l) + rows(r)) + out() + costs(l) + costs(r)
            }
            // MergeJoin: 输出复杂度 + 左子句复杂度 + 右子句复杂度
            MergeJoin([_, _, _, l, r]) => out() + costs(l) + costs(r),
            // Insert, CopyTo: 子句行数 * 子句列数 + 子句cost
            Insert([_, _, c]) | CopyTo([_, c]) => rows(c) * cols(c) + costs(c),
            // 对于空的语句, 没有复杂度
            Empty(_) => 0.0,
            // 对于column和Ref的访问是0.01常量
            Column(_) | Ref(_) => 0.01,
            // for expressions, the cost is 0.1x AST size
            // 对于表达式, 其cost大小是 0.1 * AST的大小
            _ => enode.fold(0.1, |sum: f32, id: Id| sum + costs(&id)),
        };
        debug!(
            "{id}\t{enode:?}\tcost={c}, rows={}, cols={}",
            rows(id),
            cols(id)
        );
        c
    }
}

rules包: 对于egg的规则核心
- mod.rs
  - 定义了六类代码分析规则
    - expr: 针对查询条件: 常量分析
    - range: 针对过滤扫描规则: 范围条件分析
    - plan: 规划优化: 定义和使用列
    - agg: 聚合函数: 查询条件中的聚合函数
    - schema: 列Id到Index: 返回规划的返回Schema
    - type_: 数据类型优化
    - rows: 评估行优化
    - order: 排序优化: 针对mergeJoin进行的排序键的优化
  - 定义了优化器优化的两个Stage
    - STAGE1_RULES: 使用expr, plan, order进行优化
    - STAGE2_RULES: 使用expr, plan, order进行优化
- 详细规则:
  - agg.rs
  - expr.rs
  - order.rs
  - plan.rs
  - range.rs
  - rows.rs
  - schema.rs
  - type_.rs

posted @ 2023-08-02 00:24 NoobSir 阅读(45) 评论(0) 编辑收藏举报

会员力量，点亮园子希望

刷新页面返回顶部

NoobSir

risinglightDB 核心代码学习笔记

1 整体文件列表

2 Egg优化类型转化逻辑

3 核心Egg优化逻辑子包

公告