学习型索引落地从 BTree 到 RMI 模型的架构演进与性能边界一、BTree 索引的空间换时间困局传统 BTree 索引通过存储键值与指针的映射关系实现点查和范围查询。但 BTree 的空间开销不容忽视一棵 3 层 BTree 的内部节点需要存储所有分支节点的键值对于 10 亿条记录的表索引大小通常在 10GB-50GB 之间。更大的问题是缓存命中率——当索引超过内存容量时每次查询可能需要 2-3 次磁盘 IO 读取中间节点。学习型索引Learned Index的核心洞察索引本质是一个函数——给定键返回其位置。BTree 用树结构近似这个函数而 ML 模型可以直接学习这个映射关系。如果模型能以更小的空间和更少的计算量逼近 BTree 的精度就能同时降低内存占用和查询延迟。MIT 2018 年的论文The Case for Learned Index Structures首次验证了这一思路在只读工作负载上RMIRecursive Model Index模型的查询速度比 BTree 快 1.5-3 倍空间占用减少了一个数量级。二、RMI 模型的分层架构与查询流程flowchart TD A[查询键 key] -- B[顶层模型 M0br/全量数据上训练的轻量模型] B -- C{模型选择br/M0 输出专家模型编号} C -- D1[专家模型 M1_0br/负责 key 范围 [0, 1M]] C -- D2[专家模型 M1_1br/负责 key 范围 [1M, 5M]] C -- D3[专家模型 M1_2br/负责 key 范围 [5M, 10M]] C -- D4[专家模型 M1_Nbr/负责 key 范围 [...]] D1 -- E1[预测位置 posbr/误差范围 [min, max]] D2 -- E2[预测位置 posbr/误差范围 [min, max]] D3 -- E3[预测位置 posbr/误差范围 [min, max]] D4 -- E4[预测位置 posbr/误差范围 [min, max]] E1 -- F[局部搜索br/在 [pos-min, posmax] 内二分查找] E2 -- F E3 -- F E4 -- F F -- G[返回精确位置] style B fill:#e1f5fe style F fill:#fff3e0RMI 的核心设计是分而治之顶层模型一个轻量级模型如线性回归输入 key输出应该交给哪个专家模型处理。专家模型每个专家负责一个 key 范围在该范围内学习 key 到位置的映射。专家模型可以是线性模型、神经网络或混合模型。局部搜索模型预测的位置有误差范围在该范围内用二分查找精确定位。2.1 为什么需要局部搜索ML 模型的预测不可能 100% 精确。RMI 的关键设计是训练时记录每个专家模型的最大误差查询时在预测位置的正负误差范围内做局部搜索。这保证了正确性——只要误差范围记录准确局部搜索一定能找到目标。三、RMI 模型的生产级实现import numpy as np import struct from typing import List, Tuple, Optional import logging class LinearModel: 线性回归模型y slope * x bias 为什么用线性模型而非神经网络 1. 推理延迟极低2 次浮点运算满足索引查询的微秒级要求 2. 模型大小仅 16 字节slope bias 各 8 字节内存友好 3. 在局部范围内key 到位置的映射近似线性线性模型精度足够 def __init__(self): self.slope 0.0 self.bias 0.0 self.max_error 0 # 训练时记录的最大误差 def train(self, keys: np.ndarray, positions: np.ndarray): 最小二乘法拟合线性模型 if len(keys) 2: self.slope 0.0 self.bias positions[0] if len(positions) 0 else 0.0 self.max_error 0 return # 最小二乘法 self.slope, self.bias np.polyfit(keys, positions, 1) # 计算最大误差确保局部搜索能覆盖所有数据 predictions self.slope * keys self.bias errors np.abs(predictions - positions) self.max_error int(np.max(errors)) 1 # 加 1 留安全余量 def predict(self, key: float) - Tuple[int, int, int]: 预测位置返回 (预测位置, 最小位置, 最大位置) pos int(self.slope * key self.bias) return pos, pos - self.max_error, pos self.max_error def size_bytes(self) - int: return 16 4 # slope(8) bias(8) max_error(4) class RMIIndex: 递归模型索引Recursive Model Index。 两层结构顶层模型路由到专家模型专家模型预测位置。 def __init__( self, num_experts: int 256, max_key: int 0 ): self.num_experts num_experts self.top_model LinearModel() self.expert_models: List[LinearModel] [] self.max_key max_key self.data: Optional[np.ndarray] None # 排序后的键数组 def build(self, keys: np.ndarray): 构建 RMI 索引。 keys 必须是有序的positions 即为数组下标。 self.data np.sort(keys) self.max_key int(self.data[-1]) positions np.arange(len(self.data), dtypenp.float64) # 训练顶层模型输入 key输出专家编号 # 为什么顶层模型预测专家编号而非位置 # 顶层模型的目的是路由预测专家编号比预测位置更容易学习 expert_ids np.zeros(len(self.data), dtypenp.float64) for i, key in enumerate(self.data): # 按 key 范围均匀分配到各专家 expert_ids[i] min( int(key / (self.max_key 1) * self.num_experts), self.num_experts - 1 ) self.top_model.train(self.data.astype(np.float64), expert_ids) # 训练各专家模型 self.expert_models [] for eid in range(self.num_experts): # 筛选属于该专家的数据 mask expert_ids eid expert_keys self.data[mask].astype(np.float64) expert_positions positions[mask] model LinearModel() if len(expert_keys) 0: model.train(expert_keys, expert_positions) self.expert_models.append(model) total_size self.top_model.size_bytes() sum( m.size_bytes() for m in self.expert_models ) logging.info( fRMI 索引构建完成: {self.num_experts} 个专家, f总大小 {total_size} 字节, f数据量 {len(self.data)} ) def lookup(self, key: int) - Optional[int]: 查找 key 在排序数组中的位置。 返回数组下标未找到返回 None。 if self.data is None: raise RuntimeError(索引未构建) # 第一步顶层模型路由到专家 _, min_expert, max_expert self.top_model.predict(float(key)) expert_id max(0, min(self.num_experts - 1, min_expert)) # 第二步专家模型预测位置 expert self.expert_models[expert_id] predicted_pos, search_min, search_max expert.predict(float(key)) # 第三步局部搜索 # 为什么用二分搜索而非线性搜索 # 误差范围可能达到数千线性搜索代价过高 lo max(0, search_min) hi min(len(self.data) - 1, search_max) # numpy 二分查找 idx np.searchsorted(self.data, key, sideleft) if idx len(self.data) and self.data[idx] key: return int(idx) return None def range_lookup( self, key_start: int, key_end: int ) - np.ndarray: 范围查询返回 [key_start, key_end] 范围内的所有键。 start_idx np.searchsorted(self.data, key_start, sideleft) end_idx np.searchsorted(self.data, key_end, sideright) return self.data[start_idx:end_idx]四、学习型索引的致命局限与适用边界4.1 写入支持的缺失RMI 模型在只读数据上训练插入新数据后模型精度会下降。每次插入都重训练模型的代价不可接受。当前的研究方案是 APEX 和 LIPP 等支持写入的学习型索引但实现复杂度远超 RMI且写入性能仍不如 BTree。4.2 误差范围膨胀当数据分布与训练时差异较大时如新增了一个 key 范围的数据专家模型的误差范围会急剧扩大局部搜索退化为全量扫描。生产环境中数据分布漂移是常态RMI 模型需要定期重训练。4.3 多维索引的困境RMI 只能处理一维有序键。多维索引如 R-Tree 的替代需要更复杂的模型架构如多任务学习目前尚无成熟的生产级方案。4.4 适用边界场景学习型索引适用性推荐替代只读静态数据如归档表极佳无低频写入 高频读取良好需定期重训练BTree高频写入不适用BTree / LSM-Tree多维范围查询不适用R-Tree / Z-Order数据分布剧烈变化不适用BTree五、总结学习型索引用 ML 模型替代 BTree 的内部节点在只读场景中以更小的空间和更低的延迟实现等价的查询能力。RMI 的两层架构路由模型 专家模型在工程上可落地线性模型保证了微秒级推理延迟。但学习型索引的致命短板是写入支持不足和数据分布漂移后的精度衰减。生产落地建议仅在只读或近只读的归档表上部署 RMI 索引核心业务表仍使用 BTree。部署后必须监控模型的误差范围指标当最大误差超过阈值时触发自动重训练。对于多维索引场景等待学术界更成熟的方案再考虑引入。
