使用PilotScope将learned CE与Postgres真实交互

研究生期间在做ai4db的课题，总感觉华而不实，也跟system一样做的人很少了，没什么前景，而且现在市面上也很少把ai的东西真实放到内核中使用，还是可维护和动态性太难做到了

PilotScope

前段时间关注到alibaba的工作pilotscope，把ai4db的通过pilotscope这个中间层调用到Postgres中的优化器中，还是很有意思的一个工作，想作为大论文的第三个点，实现起来也显得有点意义，但是在跑JOB时候踩好几个坑，我寻思阿里也是水个论文就结束了？也不管后续仓库维护了，我之前提的issue问有没有data-driven的代码也是没有人回的。难顶，最后都是自己实现然后自己修bug，下面说一下心路历程和如何修的bug吧

bug1: HASH_TABLE_SIZE

这个bug是在SUBQUERY_CARD_PULL_ANCHOR和CARD_PUSH_ANCHOR时会出现的，也就是获取subquery时候的报错，因为JOB中存在非常多的subquery，定义在src/backend/pilotscope/parse_json.c中的HASH_TABLE_SIZE定义的只有1023，肯定不够，我设置了16384是没有问题的，反正肯定要大于10k。

这个问题肯定会遇到，所以放第一个说，后面需要分条件说明了，就是是否设置了geqo = off

bug2: geqo

geqo是是否启用遗传算法来优化join order，默认为true

在subquery手动注入到pg的工作中，可以看到他们是关闭了grqo的：https://github.com/heriec/End-to-End-CardEst-Benchmark/blob/master/dockerfile/init_pgsql.sh

但是pilotscope是没有这样说的要求的，这在JOB中生成非常多的subquery时，会触发geqo（stats不会存在这个问题）。

geqo = on

具体调用路径在src/backend/optimizer/path/allpaths.c中的make_rel_from_joinlist

这个函数里面的

if (levels_needed == 1)
{
	/*
	 * Single joinlist node, so we're done.
	 */
	return (RelOptInfo *) linitial(initial_rels);
}
else
{
	/*
	 * Consider the different orders in which we could join the rels,
	 * using a plugin, GEQO, or the regular join search code.
	 *
	 * We put the initial_rels list into a PlannerInfo field because
	 * has_legal_joinclause() needs to look at it (ugly :-().
	 */
	root->initial_rels = initial_rels;

	if (join_search_hook)
		return (*join_search_hook) (root, levels_needed, initial_rels);
	else if (enable_geqo && levels_needed >= geqo_threshold)
		return geqo(root, levels_needed, initial_rels);
	else
		return standard_join_search(root, levels_needed, initial_rels);
}

当开启enable_geqo 同时levels_needed >= geqo_threshold就会调用geqo_eval()

也就是说：

如果 表的数量 < geqo_threshold（默认 12） → 用动态规划算法。
如果 表的数量 ≥ geqo_threshold → 启动 GEQO。

调用geqo() 会在 planner 决定用 GEQO 代替标准的动态规划 join search 时触发。

geqo 的调用在 src/backend/optimizer/geqo/geqo_main.c然后调用一些其他函数，直到调用到位于src/backend/optimizer/geqo/geqo_eval.c的

geqo_eval(PlannerInfo *root, Gene *tour, int num_gene)

就不贴详细代码了，自己看源码，具体就是遗传算法期间临时创建一个内存 cxt，用完会被释放

mycontext = AllocSetContextCreate(CurrentMemoryContext,
                                  "GEQO",
                                  ALLOCSET_DEFAULT_SIZES);

所以在 subquery 调用的时候会调用src/backend/pilotscope/utils/hashtable.c的create_entry函数创建 entry：

Entry* entry = (Entry*)palloc(sizeof(Entry));
entry->key   = (char*)palloc(strlen(key) + 1);
entry->value = (char*)palloc(strlen(value) + 1);

分配内存空间会被分配到GEQO这个 cxt 上

mycontext = AllocSetContextCreate(CurrentMemoryContext,
									  "GEQO",
									  ALLOCSET_DEFAULT_SIZES);

如果没有使用就是使用MessageContext，GEQO是他的 children（pg里的内存是树状的）

PostgreSQL 的常见上下文包括：

TopMemoryContext

CacheMemoryContext

ErrorContext

TopTransactionContext

PortalContext

QueryContext（规划器或执行器的上下文）

MessageContext

PlannerContext

GEQO（遗传算法评估临时用）

还有函数调用级别的上下文，比如 ExprContext 等。

所以当 GEQO 的内存在释放后，之前创建的Entry是存在 count_table 的，但是GEQO会释放这块的内存，导致后面从count_table中get时候内存找不到对应数据，只是为了解决这个bug的话，我就把create_entry时单独处理了一下，感觉不是很好的处理方式，仅作思路参考：

// create entry
Entry* create_entry(const char* key, const char* value) 
{
    MemoryContext oldcxt = NULL;
    if (strcmp(CurrentMemoryContext->name, "MessageContext") != 0)
    {
        oldcxt = MemoryContextSwitchTo(CurrentMemoryContext->parent);
    }
    Entry* entry = (Entry*)palloc(sizeof(Entry));
    entry->key   = (char*)palloc(strlen(key) + 1);
    entry->value = (char*)palloc(strlen(value) + 1);
    entry->next  = NULL;
    strcpy(entry->key, key);
    strcpy(entry->value, value);
    if (oldcxt != NULL)
    {
        MemoryContextSwitchTo(oldcxt);
    }
    return entry;
}

geqo = off

好，你说我不这么麻烦了，我直接关了不就行了，哈哈，那你就遇到下一个坑，还是JOB过大的subquery导致的，在上面的HASH_TABLE_SIZE我们给count_table分配容量时候太小出现问题了，现在在CARD_PUSH_ANCHOR时存subquery2card又出现问题了，具体在调用函数store_aimodel_subquery2card()里面的，在src/backend/pilotscope/anchor2struct.c中：

int table_size = card_push_anchor->card_num * card_push_anchor->card_num;
table      = create_hashtable(table_size);

table是个全局对象

table_size定义为$card_num^2$，这就导致了card_num大于10k时，非常大，我们知道他是为了防止碰撞，但是这样需要分配太大的空间，毕竟都是简化实现的，这里也不存在什么动态扩容的代码（你要想实现你就自己写），所以也就选合理的table_size，保守一点就用「略大于 n 的素数」作为哈希表容量：

int is_prime(int n) {
    if (n < 2) return 0;
    if (n == 2) return 1;
    if (n % 2 == 0) return 0;
    for (int i = 3; i * i <= n; i += 2)
        if (n % i == 0)
            return 0;
    return 1;
}

int next_prime(int x) {
    if (x <= 2) return 2;
    if (x % 2 == 0) x++;
    while (!is_prime(x)) x += 2;
    return x;
}

使用：

// avoid hash confict
int n = card_push_anchor->card_num;
int table_size = next_prime(n * 2);

多说一句，qeqo生成的subquery_2_card顺序和普通的dp生成的顺序是不一样的，具体可以自己debug看下

上述流程贴一个：

planner()
 └── subquery_planner()
     └── grouping_planner()
         └── query_planner()
             └── make_one_rel()
                 ├── 普通情况: standard join search
                 └── 启动 GEQO (>= geqo_threshold)
                      └── geqo()
                          ├── random_init_pool()
                          │     └── geqo_eval()  ← 第一次评估
                          ├── evolution loop
                          │     └── geqo_eval()  ← 每一代评估
                          └── 选出最优 joinrel

learned Cardinality Estimator

研究生期间的主要工作还是集中于data-driven的cardinality estimator，说白了就是用模型拟合数据分布，比起query-driven，我个人感觉还是可靠一些的，现有的方法还是有一些的，在与PilotScope结合时候就体现出来每个表训练一个模型然后进行基数估计的优点了，方便实现很多，我花了点时间将我的工作实现起来，总体执行时间肯定好于Postgres，在geqo=off下，JOB的结果是：

TrueCard Total Time: 95.0308s
Postgres Total Time: 228.6600s
ours Total Time: 157.6513s

stats结果是：

TrueCard Total Time: 3165.6807s
Postgres Total Time: 4950.4897s
ours Total Time: 3322.1512s

这个结果和我们论文中手动注入pg的结果没什么出入

需要说明的是，这里我们开启了并行，所以发现TrueCard的时间是变快的，同时stats的整体时间都减少了

后续还是有工作量的，我遇到的几个问题是：

在数据库层面，我是不知道哪些 attr 是 join key 的。如何获取？
模型训练时候到什么时候停止，如何判断？根据 loss 变化
模型超参定义？中间层多少，隐藏 size 多少？这个感觉和上面这条有一点类似
模型更新？定时+数据分布变化？用什么来衡量数据分布变化？

我想我大论文第三个点的一些流程图可以从这里面搞了

总结

这样一来，我研究生期间的工作基本就是这样子了，后续听天由命了，工作上面的东西也可以开始搞一搞，我个人还是比较喜欢未雨绸缪的，不然怕到时候来不及，但是整个节奏我还是比较喜欢的，提前完成一些任务，然后后面可以留足充足的时间，无论是继续卷还是放开玩都没有顾虑~

Heriecの Blog