[WindowFunction] 系列二：TiDB 中的 Window Function 实现

接下来几篇文章将分析一些热门数据库中的窗口函数实现方式，主要包括 节点间并发、节点内并发、具体实现、内存控制，以及其它值得注意的细节。

TiDB

代码：代码
文档：文档
design doc: design doc

窗口函数分组

在 parser 将 sql 解析为 ast 之后，tidb 的 planbuilder 会将 windowFunction 进行分组，合并具有相同 spec name 的窗口，核心逻辑在以下函数中：

func (b *PlanBuilder) groupWindowFuncs(windowFuncs []*ast.WindowFuncExpr) (map[*ast.WindowSpec][]*ast.WindowFuncExpr, []*ast.WindowSpec, error) {
	// updatedSpecMap is used to handle the specifications that have frame clause changed.
	updatedSpecMap := make(map[string][]*ast.WindowSpec)
	groupedWindow := make(map[*ast.WindowSpec][]*ast.WindowFuncExpr)
	orderedSpec := make([]*ast.WindowSpec, 0, len(windowFuncs))
	for _, windowFunc := range windowFuncs {
		if windowFunc.Spec.Name.L == "" {
			spec := &windowFunc.Spec
			if spec.Ref.L != "" {
				ref, ok := b.windowSpecs[spec.Ref.L]
				if !ok {
					return nil, nil, ErrWindowNoSuchWindow.GenWithStackByArgs(getWindowName(spec.Ref.O))
				}
				err := mergeWindowSpec(spec, ref)
				if err != nil {
					return nil, nil, err
				}
			}
			spec, _ = b.handleDefaultFrame(spec, windowFunc.F)
			groupedWindow[spec] = append(groupedWindow[spec], windowFunc)
			orderedSpec = appendIfAbsentWindowSpec(orderedSpec, spec)
			continue
		}

		name := windowFunc.Spec.Name.L
		spec, ok := b.windowSpecs[name]
		if !ok {
			return nil, nil, ErrWindowNoSuchWindow.GenWithStackByArgs(windowFunc.Spec.Name.O)
		}
		newSpec, updated := b.handleDefaultFrame(spec, windowFunc.F)
		if !updated {
			groupedWindow[spec] = append(groupedWindow[spec], windowFunc)
			orderedSpec = appendIfAbsentWindowSpec(orderedSpec, spec)
		} else {
			var updatedSpec *ast.WindowSpec
			if _, ok := updatedSpecMap[name]; !ok {
				updatedSpecMap[name] = []*ast.WindowSpec{newSpec}
				updatedSpec = newSpec
			} else {
				for _, spec := range updatedSpecMap[name] {
					eq, err := specEqual(spec, newSpec)
					if err != nil {
						return nil, nil, err
					}
					if eq {
						updatedSpec = spec
						break
					}
				}
				if updatedSpec == nil {
					updatedSpec = newSpec
					updatedSpecMap[name] = append(updatedSpecMap[name], newSpec)
				}
			}
			groupedWindow[updatedSpec] = append(groupedWindow[updatedSpec], windowFunc)
			orderedSpec = appendIfAbsentWindowSpec(orderedSpec, updatedSpec)
		}
	}
	// Unused window specs should also be checked in b.buildWindowFunctions,
	// so we add them to `groupedWindow` with empty window functions.
	for _, spec := range b.windowSpecs {
		if _, ok := groupedWindow[spec]; !ok {
			if _, ok = updatedSpecMap[spec.Name.L]; !ok {
				groupedWindow[spec] = nil
				orderedSpec = appendIfAbsentWindowSpec(orderedSpec, spec)
			}
		}
	}
	return groupedWindow, orderedSpec, nil

这个函数有两个返回值，第一个是 spec 到 func 的 map，第二个是 spec 的切片，函数流程如下：

依次遍历所有 windowFunction：

1. 如果 spec 匿名：
   • 首先调用 mergeWindowSpec 处理 ref spec（形如 w2 as w1 order by deptid），填充 partitionBy 和 orderBy
   • 调用 handleDefaultFrame 函数处理 frame，会根据 func 的默认 frame 与 spec frame 进行组合。
   • 将这对 spec/func 添加到返回值中
   • continue
2. 对于具名 spec：
   • 首先处理 frame，注意对于相同的 spec，不同的 func 可能会产生不同的 frame，从而改变 spec。
   • 如果 spec 不会因 frame 改变，添加到返回值中，返回。
3. 否则，利用 specEqual 函数判断是否之前已经生成出来过相同的 spec。如果有就合并，如果没有就创建。
   然后添加到返回值中，返回。

想到一个优化点：为什么不用 specEqual 判断匿名 spec 和具名 spec 是否相等从而进行更多的合并。。联系了开发者，称最后没来得及做，所以留了个 todo。

一个问题：
在 tidb 中，在同一个窗口 w 上执行 row_number 和 rank，并不会将它们分到一个组中。

mysql> explain select *, row_number() over w, rank() over w from employee window w as (partition by deptid);
+--------------------------------+---------+-----------+----------------+---------------------------------------------------------------------------------------------------------+
| id                             | estRows | task      | access object  | operator info                                                                                           |
+--------------------------------+---------+-----------+----------------+---------------------------------------------------------------------------------------------------------+
| Projection_8                   | 17.00   | root      |                | test.employee.empid, test.employee.deptid, test.employee.salary, Column#8, Column#7                     |
| └─Window_10                    | 17.00   | root      |                | row_number()->Column#8 over(partition by test.employee.deptid rows between current row and current row) |
|   └─Window_11                  | 17.00   | root      |                | rank()->Column#7 over(partition by test.employee.deptid)                                                |
|     └─Sort_15                  | 17.00   | root      |                | test.employee.deptid                                                                                    |
|       └─TableReader_14         | 17.00   | root      |                | data:TableFullScan_13                                                                                   |
|         └─TableFullScan_13     | 17.00   | cop[tikv] | table:employee | keep order:false, stats:pseudo                                                                          |
+--------------------------------+---------+-----------+----------------+---------------------------------------------------------------------------------------------------------+

答案：因为 handleDefaultFrame 会分别对这两个函数进行处理：

1. 给 row_number 加了一个 deafult frame（between current row and current row），rank 没加，所以不认为他们是一组了。

至于为什么要给 row_number 加 default_frame，开发者的回复是，为了避免将整个分区物化的内存消耗，所以给 row_number 加了一个宽度为 1 的窗口。