学习索引：当机器学习遇上数据库

def learned_lookup(model, data, key):
    # 模型预测位置
    predicted_pos = model.predict(key)
    predicted_pos = clamp(predicted_pos, 0, len(data) - 1)

    # 在误差范围内做局部搜索
    lo = max(0, predicted_pos - model.max_error)
    hi = min(len(data) - 1, predicted_pos + model.max_error)

    # 二分查找精确定位
    while lo <= hi:
        mid = (lo + hi) // 2
        if data[mid] == key:
            return mid
        elif data[mid] < key:
            lo = mid + 1
        else:
            hi = mid - 1

    return -1  # not found

def train_rmi(data, keys, stages):
    """
    data:   排序后的 key 数组
    stages: 每层的模型数量，如 [1, 100, 10000]
    """
    models = []
    # Stage 0：在所有数据上训练根模型
    root = LinearRegression()
    root.fit(keys, positions)
    models.append([root])

    for stage_idx in range(1, len(stages)):
        stage_models = [None] * stages[stage_idx]
        # 用上一层的预测来分配 key 到当前层的模型
        for key, pos in zip(keys, positions):
            predicted = models[stage_idx - 1][...].predict(key)
            model_idx = int(predicted * stages[stage_idx] / len(data))
            model_idx = clamp(model_idx, 0, stages[stage_idx] - 1)
            # 将 (key, pos) 分配给 stage_models[model_idx]
            assign(stage_models[model_idx], key, pos)

        # 训练当前层的每个模型
        for m in stage_models:
            m.fit()

        models.append(stage_models)

    return models

// 最优分段线性近似的核心算法（简化版）
struct Segment {
    double slope;
    double intercept;
    size_t start_key;
    size_t start_pos;
};

std::vector<Segment> optimal_pla(
    const std::vector<int64_t>& keys,
    size_t epsilon
) {
    std::vector<Segment> segments;
    size_t seg_start = 0;

    while (seg_start < keys.size()) {
        // 初始化可行域（一个锥形区域）
        double upper_slope = INFINITY;
        double lower_slope = -INFINITY;
        double intercept_at_start = seg_start;

        size_t seg_end = seg_start;

        for (size_t i = seg_start + 1; i < keys.size(); ++i) {
            double dx = keys[i] - keys[seg_start];
            if (dx == 0) { seg_end = i; continue; }

            double upper = (double(i - seg_start) + epsilon) / dx;
            double lower = (double(i - seg_start) - epsilon) / dx;

            // 收缩可行域
            upper_slope = std::min(upper_slope, upper);
            lower_slope = std::max(lower_slope, lower);

            if (upper_slope < lower_slope) {
                // 可行域为空，当前线段到此结束
                break;
            }
            seg_end = i;
        }

        double slope = (upper_slope + lower_slope) / 2.0;
        segments.push_back({slope, intercept_at_start,
                           (size_t)keys[seg_start], seg_start});
        seg_start = seg_end + 1;
    }

    return segments;
}

// 指数搜索：从预测位置开始，步长倍增
size_t exponential_search(
    const int64_t* data,
    size_t predicted_pos,
    int64_t key,
    size_t data_size
) {
    // 先判断方向
    if (data[predicted_pos] == key) return predicted_pos;

    if (data[predicted_pos] < key) {
        // 向右搜索
        size_t step = 1;
        size_t pos = predicted_pos;
        while (pos + step < data_size && data[pos + step] < key) {
            step *= 2;
        }
        // 在 [pos + step/2, min(pos + step, data_size-1)] 内二分
        return binary_search(data, pos + step / 2,
                           std::min(pos + step, data_size - 1), key);
    } else {
        // 向左搜索
        size_t step = 1;
        size_t pos = predicted_pos;
        while (pos >= step && data[pos - step] > key) {
            step *= 2;
        }
        size_t lo = (pos >= step) ? pos - step : 0;
        return binary_search(data, lo, pos - step / 2, key);
    }
}

// LIPP 节点重建（概念代码）
void rebuild_node(LIPPNode* node) {
    // 1. 收集所有 key-value 对
    auto all_kvs = collect_all_kvs(node);

    // 2. 计算新的容量（通常是 key 数量的 1.2-1.5 倍）
    size_t new_capacity = all_kvs.size() * expansion_factor;

    // 3. 训练新的线性模型
    LinearModel new_model;
    new_model.train(all_kvs, new_capacity);

    // 4. 重新分配 key 到新数组
    node->resize(new_capacity);
    node->model = new_model;
    for (auto& kv : all_kvs) {
        size_t pos = new_model.predict(kv.key);
        node->insert_at(pos, kv);
    }
}

def learned_bloom_filter(model, backup_filter, key):
    # 模型先做预测
    score = model.predict(key)  # 输出 [0, 1] 之间的概率

    if score > threshold:
        return True   # 模型认为 key 存在
    else:
        # 模型不确定的情况，用备份 Bloom filter 兜底
        return backup_filter.query(key)

// 学习化 Bloom filter 的简化实现
class LearnedBloomFilter {
    NeuralNetwork model_;
    BloomFilter backup_;
    double threshold_;

public:
    void build(const std::vector<int64_t>& keys,
               const std::vector<int64_t>& non_keys,
               double target_fpr) {
        // 训练二分类模型
        model_.train(keys, non_keys);

        // 找到使 false negative rate = 0 的阈值
        threshold_ = find_threshold(model_, keys);

        // 把模型漏掉的 key 放入备份 filter
        std::vector<int64_t> missed;
        for (auto k : keys) {
            if (model_.predict(k) < threshold_) {
                missed.push_back(k);
            }
        }
        backup_.build(missed, target_fpr);
    }

    bool query(int64_t key) {
        if (model_.predict(key) >= threshold_) {
            return true;
        }
        return backup_.query(key);
    }
};

// learned_index.hpp
// 一个完整的简单学习索引实现
// 支持：构建、点查询、范围查询
// 模型：两层 RMI，所有层使用线性回归

#pragma once
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <iostream>
#include <numeric>
#include <vector>

namespace learned_index {

// 线性模型：pos = slope * key + intercept
struct LinearModel {
    double slope = 0.0;
    double intercept = 0.0;

    // 用最小二乘法训练
    void train(const std::vector<int64_t>& keys,
               const std::vector<double>& positions) {
        assert(keys.size() == positions.size());
        size_t n = keys.size();
        if (n == 0) { slope = 0; intercept = 0; return; }
        if (n == 1) { slope = 0; intercept = positions[0]; return; }

        double sum_x = 0, sum_y = 0, sum_xy = 0, sum_xx = 0;
        for (size_t i = 0; i < n; ++i) {
            double x = static_cast<double>(keys[i]);
            double y = positions[i];
            sum_x += x;
            sum_y += y;
            sum_xy += x * y;
            sum_xx += x * x;
        }

        double denom = n * sum_xx - sum_x * sum_x;
        if (std::abs(denom) < 1e-10) {
            slope = 0;
            intercept = sum_y / n;
        } else {
            slope = (n * sum_xy - sum_x * sum_y) / denom;
            intercept = (sum_y - slope * sum_x) / n;
        }
    }

    double predict(int64_t key) const {
        return slope * static_cast<double>(key) + intercept;
    }
};

// 叶子模型：线性模型 + 误差界
struct LeafModel {
    LinearModel model;
    size_t start_pos = 0;     // 该模型负责的数据起始位置
    size_t end_pos = 0;       // 该模型负责的数据结束位置
    int64_t min_key = 0;
    int64_t max_key = 0;
    size_t max_error = 0;     // 最大预测误差
};

class SimpleLearnedIndex {
public:
    // 构建索引
    void build(const std::vector<int64_t>& sorted_data,
               size_t num_leaf_models = 0) {
        data_ = sorted_data;
        size_t n = data_.size();
        if (n == 0) return;

        // 默认叶子模型数量：sqrt(N)
        if (num_leaf_models == 0) {
            num_leaf_models = std::max<size_t>(
                1, static_cast<size_t>(std::sqrt(n)));
        }
        num_leaves_ = num_leaf_models;

        // --- 训练根模型 ---
        std::vector<int64_t> all_keys(n);
        std::vector<double> all_positions(n);
        for (size_t i = 0; i < n; ++i) {
            all_keys[i] = data_[i];
            all_positions[i] = static_cast<double>(i);
        }
        root_model_.train(all_keys, all_positions);

        // --- 将 key 分配到叶子模型 ---
        leaves_.resize(num_leaves_);
        std::vector<std::vector<int64_t>> leaf_keys(num_leaves_);
        std::vector<std::vector<double>> leaf_positions(num_leaves_);

        for (size_t i = 0; i < n; ++i) {
            double pred = root_model_.predict(data_[i]);
            // 将预测值映射到叶子模型索引
            size_t leaf_idx = static_cast<size_t>(
                pred / static_cast<double>(n) * num_leaves_);
            leaf_idx = std::min(leaf_idx, num_leaves_ - 1);

            leaf_keys[leaf_idx].push_back(data_[i]);
            leaf_positions[leaf_idx].push_back(
                static_cast<double>(i));
        }

        // --- 训练每个叶子模型 ---
        for (size_t i = 0; i < num_leaves_; ++i) {
            auto& leaf = leaves_[i];
            if (leaf_keys[i].empty()) {
                // 空叶子节点，设置为默认
                leaf.start_pos = 0;
                leaf.end_pos = 0;
                leaf.max_error = 0;
                continue;
            }

            leaf.model.train(leaf_keys[i], leaf_positions[i]);
            leaf.min_key = leaf_keys[i].front();
            leaf.max_key = leaf_keys[i].back();
            leaf.start_pos = static_cast<size_t>(
                leaf_positions[i].front());
            leaf.end_pos = static_cast<size_t>(
                leaf_positions[i].back());

            // 计算最大误差
            leaf.max_error = 0;
            for (size_t j = 0; j < leaf_keys[i].size(); ++j) {
                double pred = leaf.model.predict(leaf_keys[i][j]);
                double actual = leaf_positions[i][j];
                size_t err = static_cast<size_t>(
                    std::abs(pred - actual)) + 1;
                leaf.max_error = std::max(leaf.max_error, err);
            }
        }
    }

    // 点查询：返回 key 的位置，未找到返回 -1
    int64_t lookup(int64_t key) const {
        if (data_.empty()) return -1;
        size_t n = data_.size();

        // Stage 0：根模型选择叶子
        double root_pred = root_model_.predict(key);
        size_t leaf_idx = static_cast<size_t>(
            root_pred / static_cast<double>(n) * num_leaves_);
        leaf_idx = std::min(leaf_idx, num_leaves_ - 1);

        const auto& leaf = leaves_[leaf_idx];
        if (leaf.start_pos == leaf.end_pos &&
            leaf.max_error == 0) {
            return -1;
        }

        // Stage 1：叶子模型预测位置
        double pred = leaf.model.predict(key);
        int64_t predicted_pos = static_cast<int64_t>(
            std::round(pred));

        // 限制搜索范围
        int64_t lo = std::max<int64_t>(
            0, predicted_pos - static_cast<int64_t>(
                leaf.max_error));
        int64_t hi = std::min<int64_t>(
            static_cast<int64_t>(n) - 1,
            predicted_pos + static_cast<int64_t>(
                leaf.max_error));
        lo = std::max<int64_t>(lo, 0);
        hi = std::min<int64_t>(hi,
            static_cast<int64_t>(n) - 1);

        // 二分搜索精确定位
        while (lo <= hi) {
            int64_t mid = lo + (hi - lo) / 2;
            if (data_[mid] == key) return mid;
            if (data_[mid] < key) lo = mid + 1;
            else hi = mid - 1;
        }

        return -1;
    }

    // 范围查询：返回 [lo_key, hi_key] 范围内的所有 key
    std::vector<int64_t> range_query(
        int64_t lo_key, int64_t hi_key) const {
        std::vector<int64_t> result;
        if (data_.empty()) return result;

        // 找到 lo_key 的起始位置
        int64_t start = find_lower_bound(lo_key);
        if (start < 0) start = 0;

        // 扫描直到超出 hi_key
        for (size_t i = static_cast<size_t>(start);
             i < data_.size() && data_[i] <= hi_key; ++i) {
            if (data_[i] >= lo_key) {
                result.push_back(data_[i]);
            }
        }

        return result;
    }

    // 获取统计信息
    void print_stats() const {
        size_t total_max_err = 0;
        size_t non_empty = 0;
        for (const auto& leaf : leaves_) {
            if (leaf.max_error > 0) {
                total_max_err += leaf.max_error;
                ++non_empty;
            }
        }
        double avg_err = non_empty > 0
            ? static_cast<double>(total_max_err) / non_empty
            : 0;

        std::cout << "=== Learned Index Stats ===" << '\n';
        std::cout << "Data size:       " << data_.size()
                  << '\n';
        std::cout << "Leaf models:     " << num_leaves_
                  << '\n';
        std::cout << "Non-empty leaves:" << non_empty
                  << '\n';
        std::cout << "Avg max error:   " << avg_err
                  << '\n';
        std::cout << "Model memory:    "
                  << model_size_bytes() << " bytes"
                  << '\n';
    }

    size_t model_size_bytes() const {
        // 根模型：2 * sizeof(double) = 16 bytes
        // 每个叶子模型：2 * sizeof(double) + metadata
        return 16 + num_leaves_ * sizeof(LeafModel);
    }

private:
    int64_t find_lower_bound(int64_t key) const {
        size_t n = data_.size();
        double root_pred = root_model_.predict(key);
        size_t leaf_idx = static_cast<size_t>(
            root_pred / static_cast<double>(n) * num_leaves_);
        leaf_idx = std::min(leaf_idx, num_leaves_ - 1);

        const auto& leaf = leaves_[leaf_idx];
        double pred = leaf.model.predict(key);
        int64_t pos = static_cast<int64_t>(std::round(pred));

        int64_t lo = std::max<int64_t>(
            0, pos - static_cast<int64_t>(leaf.max_error));
        int64_t hi = std::min<int64_t>(
            static_cast<int64_t>(n) - 1,
            pos + static_cast<int64_t>(leaf.max_error));

        // 标准 lower_bound
        while (lo < hi) {
            int64_t mid = lo + (hi - lo) / 2;
            if (data_[mid] < key) lo = mid + 1;
            else hi = mid;
        }
        return lo;
    }

    std::vector<int64_t> data_;
    LinearModel root_model_;
    std::vector<LeafModel> leaves_;
    size_t num_leaves_ = 0;
};

}  // namespace learned_index

// benchmark.cpp
// 编译：g++ -O2 -std=c++17 -o benchmark benchmark.cpp
// 运行：./benchmark

#include "learned_index.hpp"
#include <chrono>
#include <random>

using namespace learned_index;
using Clock = std::chrono::high_resolution_clock;

// 生成不同分布的数据
std::vector<int64_t> generate_uniform(size_t n) {
    std::vector<int64_t> data(n);
    for (size_t i = 0; i < n; ++i) data[i] = i * 10;
    return data;
}

std::vector<int64_t> generate_normal(size_t n) {
    std::mt19937_64 rng(42);
    std::normal_distribution<double> dist(5e8, 1e7);
    std::vector<int64_t> data(n);
    for (size_t i = 0; i < n; ++i) {
        data[i] = static_cast<int64_t>(dist(rng));
    }
    std::sort(data.begin(), data.end());
    data.erase(std::unique(data.begin(), data.end()),
               data.end());
    return data;
}

std::vector<int64_t> generate_lognormal(size_t n) {
    std::mt19937_64 rng(42);
    std::lognormal_distribution<double> dist(20.0, 2.0);
    std::vector<int64_t> data(n);
    for (size_t i = 0; i < n; ++i) {
        data[i] = static_cast<int64_t>(dist(rng));
    }
    std::sort(data.begin(), data.end());
    data.erase(std::unique(data.begin(), data.end()),
               data.end());
    return data;
}

double benchmark_learned_index(
    const std::vector<int64_t>& data,
    const std::vector<int64_t>& queries
) {
    SimpleLearnedIndex idx;
    idx.build(data);

    auto start = Clock::now();
    volatile int64_t sink = 0;
    for (auto q : queries) {
        sink = idx.lookup(q);
    }
    auto end = Clock::now();
    double ns = std::chrono::duration_cast<
        std::chrono::nanoseconds>(end - start).count();
    return ns / queries.size();
}

double benchmark_binary_search(
    const std::vector<int64_t>& data,
    const std::vector<int64_t>& queries
) {
    auto start = Clock::now();
    volatile bool sink = false;
    for (auto q : queries) {
        sink = std::binary_search(
            data.begin(), data.end(), q);
    }
    auto end = Clock::now();
    double ns = std::chrono::duration_cast<
        std::chrono::nanoseconds>(end - start).count();
    return ns / queries.size();
}

int main() {
    const size_t N = 10'000'000;
    const size_t Q = 1'000'000;

    struct TestCase {
        std::string name;
        std::vector<int64_t> data;
    };

    std::vector<TestCase> tests = {
        {"Uniform",   generate_uniform(N)},
        {"Normal",    generate_normal(N)},
        {"Lognormal", generate_lognormal(N)},
    };

    for (auto& tc : tests) {
        // 生成查询（随机选取已有 key）
        std::mt19937_64 rng(123);
        std::uniform_int_distribution<size_t> dist(
            0, tc.data.size() - 1);
        std::vector<int64_t> queries(Q);
        for (size_t i = 0; i < Q; ++i) {
            queries[i] = tc.data[dist(rng)];
        }

        std::cout << "\n--- " << tc.name
                  << " (N=" << tc.data.size() << ") ---\n";

        SimpleLearnedIndex idx;
        idx.build(tc.data);
        idx.print_stats();

        double learned_ns = benchmark_learned_index(
            tc.data, queries);
        double binary_ns = benchmark_binary_search(
            tc.data, queries);

        std::cout << "Learned index: "
                  << learned_ns << " ns/query\n";
        std::cout << "Binary search: "
                  << binary_ns << " ns/query\n";
        std::cout << "Speedup:       "
                  << binary_ns / learned_ns << "x\n";
    }

    return 0;
}

# 范围查询伪代码
def range_query(index, lo, hi):
    start_pos = index.lookup_lower_bound(lo)  # 这一步有差异
    result = []
    for i in range(start_pos, len(data)):     # 这一步没差异
        if data[i] > hi:
            break
        result.append(data[i])
    return result