为什么从 Order By 切入优化器

查询优化器是关系型数据库的核心,也是代码中最难懂的部分。以 Order By 作为切入口比较合适:它涉及的链路足够完整,从语法解析、逻辑计划、物理计划一直延伸到执行器,但复杂度又相对可控。

分析语句的正确方式是从 parser 入手,先看它构造了哪些相关的 AST 结构。

parser 阶段:构造 AST

Order By 的语法规则最终会构造出 ast.OrderByClause,其核心数据结构是 ast.ByItem

// ByItem represents an item in order by or group by.
type ByItem struct {
    node

    Expr ExprNode
    Desc bool
}

可以看到,每个排序项只包含两件事:排序表达式,以及是否为降序。语法解析阶段做的事情很纯粹,就是把 SQL 文本变成这样的结构列表,不做任何优化判断。

plan 阶段:生成逻辑计划

接下来看 planBuilder.buildSelect。当语句中存在 Order By 子句时,会调用 buildSort 创建一个 Sort 的逻辑计划:

if sel.OrderBy != nil {
    p = b.buildSort(p, sel.OrderBy.Items, orderMap)
    if b.err != nil {
        return nil
    }
}

对应的逻辑计划结构大致是:

// Sort stands for the order by plan.
type Sort struct {
    *basePlan
    baseLogicalPlan
    basePhysicalPlan

    ByItems   []*ByItems
    ExecLimit *Limit
}

聚合函数与 selectFields 的对齐

这里有个容易被忽略的细节。如果 order by 的表达式里带有聚合函数,需要先做一轮处理。以 resolveHavingAndOrderBy 为例,它会用 havingAndOrderbyExprResolver 遍历相关表达式,主要目标是:

  • 找到 order by 中带有聚合函数的相关联的列;
  • 把聚合函数中的参数替换成对应 selectFields 里的项;
  • 得到聚合函数到 selectFields 下标的映射关系,供后续 Sort 算子使用。

由于 TiDB 大量使用 Visitor 模式,处理逻辑体现在 Enter 与 Leave 方法中。走到 AggregateFuncExpr 时,会对每一个列引用参数调用 Accept 来做改写;最终在对应的 Leave 处理中,把自己转换为 selectFields[index].Expr

而对于聚合函数本身,则会在 selectFields 中追加一个辅助的 SelectField,并在映射中记录自己在 selectFields 数组中的下标,方便 Sort 算子使用:

case *ast.AggregateFuncExpr:
    a.inAggFunc = false
    a.aggMapper[v] = len(a.selectFields)
    a.selectFields = append(a.selectFields, &ast.SelectField{
        Auxiliary: true,
        Expr:      v,
        AsName:    model.NewCIStr(fmt.Sprintf("sel_agg_%d", len(a.selectFields))),
    })

buildSort 与表达式改写

再来看 buildSort 本身。它会为每个 byItem 的表达式调用 rewrite:

for _, item := range byItems {
    it, np, err := b.rewrite(item.Expr, p, aggMapper, true)
    if err != nil {
        b.err = err
        return nil
    }
    p = np
    exprs = append(exprs, &ByItems{Expr: it, Desc: item.Desc})
}

rewrite 的含义是:在计划构建阶段就能确定结果的表达式,直接折叠成常量。例如 YEAR('2009-05-19') 可以直接用 2009 表示,1 + 1 也可以是 2。当然还会有其他需要转换的情况。Order By 子句里大部分是列引用,rewrite 会把它转换成对应 schema 中的列对象。

至此 sort.ByItems 构造完成。

Physical Plan:排序要不要真的做

逻辑计划构建完成后,会通过 convert2NewPhysicalPlan 转换为物理计划。对于语句 select * from t order by a;,此时的逻辑计划层级关系是 Sort -> Projection -> DataSource

比较值得注意的是:在 Sort 的物理计划转换过程中,会分别两次调用子节点的逻辑计划转物理计划,区别只在于第二次在 requiredProp 中加入了 order by 相应的列。这个做法的目的是比较两次生成的计划代价,从而决定是否下推这个 order by。

如果发现 order by 刚好可以用上索引的有序性(例如排序的正是主键或某个索引的前缀),那么此时就不需要构造 Sort 的物理计划,直接使用本身就保序的读算子即可。取得的结果里会把 copTask 的 keepOrder 置为 true。具体用到 requiredProp 中列信息的地方,是在 DataSource 的 convertToIndexScan,它会尝试把每个索引都转成 IndexScan,再按代价取最小者。

在 finishCopTask 中会最终计算 task 的代价,并根据 task 中的 tablePlan 与 indexPlan 构造最后的物理计划:

  • PhysicalIndexLookUpReader:tablePlan 与 indexPlan 都提供;
  • PhysicalIndexReader:只有 indexPlan;
  • PhysicalTableReader:其余情况。

getBestTask 的代价比较

func (p *baseLogicalPlan) getBestTask(bestTask task, prop *requiredProp, pp PhysicalPlan) (task, error) {
    newProps := pp.getChildrenPossibleProps(prop)
    for _, newProp := range newProps {
        tasks := make([]task, 0, len(p.basePlan.children))
        for i, child := range p.basePlan.children {
            childTask, err := child.(LogicalPlan).convert2NewPhysicalPlan(newProp[i])
            if err != nil {
                return nil, errors.Trace(err)
            }
            tasks = append(tasks, childTask)
        }
        // 这里是关键, 会计算 task 的代价, 如果当前 resultTask 代价小于 bestTask,
        // 那么会替换当前的 bestTask. 最终选出最优的 task
        resultTask := pp.attach2Task(tasks...)
        if resultTask.cost() < bestTask.cost() {
            bestTask = resultTask
        }
    }
    return bestTask, nil
}

其中 attach2Task 是把子节点的 tasks 挂到当前的物理计划上,再计算总体代价。至此 sort 的计划部分就结束了。

执行器:SortExec 与 TopNExec

之后代码会进入 ExecStmt.buildExecutor 构造最终的执行器。对于 Sort,会先 build 孩子算子的执行器,再构造 SortExec。这里有一个值得注意的优化:

if v.ExecLimit != nil {
    return &TopNExec{
        SortExec: sortExec,
        limit:    v.ExecLimit,
    }
}

也就是说,当 SQL 带有 LIMIT 时,会构造 TopNExec 而不是完整的 SortExec,避免对全量结果排序。这是一个很实际的优化。

有了执行器之后,TiDB 会调用它的 Next 接口。执行时把子算子的所有行取出,为每一行按 ByItems 计算排序键,然后统一排序输出。

小结

自此,Order By 的代码路径基本走完。整条链路可以概括为:parser 构造 AST -> buildSort 生成逻辑计划并做表达式改写 -> 物理计划阶段通过两次代价比较决定排序是否下推 -> 执行器根据是否存在 LIMIT 选择 SortExec 或 TopNExec。

其中还有很多细节没有展开,例如 Visitor 模式的具体实现、requiredProp 的传播规则等,需要读者进一步阅读代码深入了解。

Order By 处理链路速查
AST 结构OrderByClause 持有 ByItem 列表,ByItem 包含排序表达式与 Desc 标记
逻辑计划buildSelect 检测到 OrderBy 后调用 buildSort 生成 Sort 逻辑计划
表达式改写常量折叠、列引用转 schema 列、聚合函数与 selectFields 对齐
下推判定两次转换并比较代价,若可用索引有序性则省去 Sort 算子
执行器无 LIMIT 用 SortExec,有 LIMIT 转为 TopNExec 避免全量排序