PL/pgSQL按权重占比随机选取姓名生成测试样本数据方案咨询

核心优化思路 #

要满足100万级样本的生成效率，必须避免逐行循环查表的低效模式，改用批量集合操作：

• 仅1次预计算名字/姓氏的权重区间，无需每次生成都重复查表
• 一次性生成所有需要的随机值，批量匹配出对应姓名
• 尽量减少PL/pgSQL循环次数，多用原生SQL执行批量操作

方案1：无需修改原有表结构（适合临时使用场景） #

这个方案不需要改动现有表，直接通过CTE预计算累计权重，执行效率足够支撑100万条数据生成：

CREATE OR REPLACE FUNCTION create_sample_data(p_number_of_records INT)
RETURNS VOID AS $$
DECLARE
    v_records INT := coalesce(p_number_of_records, 0);
    v_full_name TEXT;
BEGIN
    IF v_records <= 0 THEN
        RETURN;
    END IF;

    -- 批量生成所有姓名后统一处理，避免循环内多次查表
    FOR v_full_name IN
        WITH 
        -- 预计算名字的权重区间
        first_name_intervals AS (
            SELECT 
                first_name,
                SUM(ratio) OVER (ORDER BY first_name) AS cum_upper,
                LAG(SUM(ratio) OVER (ORDER BY first_name), 1, 0) OVER (ORDER BY first_name) AS cum_lower
            FROM common_first_name
        ),
        -- 预计算姓氏的权重区间
        last_name_intervals AS (
            SELECT 
                last_name,
                SUM(ratio) OVER (ORDER BY last_name) AS cum_upper,
                LAG(SUM(ratio) OVER (ORDER BY last_name), 1, 0) OVER (ORDER BY last_name) AS cum_lower
            FROM common_last_name
        ),
        -- 生成指定数量的随机数对，每个对应一个姓名的随机权重
        random_seeds AS (
            SELECT 
                random() * 100 AS rnd_first,
                random() * 100 AS rnd_last
            FROM generate_series(1, v_records)
        )
        -- 匹配出符合权重分布的姓名
        SELECT f.first_name || ' ' || l.last_name AS full_name
        FROM random_seeds s
        JOIN first_name_intervals f ON s.rnd_first >= f.cum_lower AND s.rnd_first < f.cum_upper
        JOIN last_name_intervals l ON s.rnd_last >= l.cum_lower AND s.rnd_last < l.cum_upper
    LOOP
        -- 调用你已有的人员生成函数
        PERFORM create_person(v_full_name);
    END LOOP;
END
$$ LANGUAGE plpgsql VOLATILE;

方案2：修改表结构预存权重区间（适合频繁生成数据场景） #

如果需要经常调用该函数，建议预存累计权重并加索引，进一步提升性能：

步骤1：新增累计权重字段并初始化 #

-- 名字表加字段
ALTER TABLE common_first_name ADD COLUMN cum_lower NUMERIC, ADD COLUMN cum_upper NUMERIC;
-- 初始化权重区间
WITH intervals AS (
    SELECT 
        first_name,
        SUM(ratio) OVER (ORDER BY first_name) AS cum_upper,
        LAG(SUM(ratio) OVER (ORDER BY first_name), 1, 0) OVER (ORDER BY first_name) AS cum_lower
    FROM common_first_name
)
UPDATE common_first_name f
SET cum_lower = i.cum_lower, cum_upper = i.cum_upper
FROM intervals i WHERE f.first_name = i.first_name;
-- 加索引加速匹配
CREATE INDEX idx_first_name_interval ON common_first_name(cum_lower, cum_upper);

-- 姓氏表同理执行相同操作
ALTER TABLE common_last_name ADD COLUMN cum_lower NUMERIC, ADD COLUMN cum_upper NUMERIC;
WITH intervals AS (
    SELECT 
        last_name,
        SUM(ratio) OVER (ORDER BY last_name) AS cum_upper,
        LAG(SUM(ratio) OVER (ORDER BY last_name), 1, 0) OVER (ORDER BY last_name) AS cum_lower
    FROM common_last_name
)
UPDATE common_last_name l
SET cum_lower = i.cum_lower, cum_upper = i.cum_upper
FROM intervals i WHERE l.last_name = i.last_name;
CREATE INDEX idx_last_name_interval ON common_last_name(cum_lower, cum_upper);

步骤2：优化后的函数（无需每次预计算区间） #

CREATE OR REPLACE FUNCTION create_sample_data(p_number_of_records INT)
RETURNS VOID AS $$
DECLARE
    v_records INT := coalesce(p_number_of_records, 0);
    v_full_name TEXT;
BEGIN
    IF v_records <= 0 THEN
        RETURN;
    END IF;

    FOR v_full_name IN
        SELECT f.first_name || ' ' || l.last_name AS full_name
        FROM (
            SELECT random() * 100 AS rnd_first, random() * 100 AS rnd_last
            FROM generate_series(1, v_records)
        ) s
        JOIN common_first_name f ON s.rnd_first >= f.cum_lower AND s.rnd_first < f.cum_upper
        JOIN common_last_name l ON s.rnd_last >= l.cum_lower AND s.rnd_last < l.cum_upper
    LOOP
        PERFORM create_person(v_full_name);
    END LOOP;
END
$$ LANGUAGE plpgsql VOLATILE;

额外性能提升建议 #

如果你的create_person函数仅做插入操作，建议修改为支持批量输入的版本，直接将批量匹配出来的姓名集合一次性插入，完全去掉PL/pgSQL循环，100万条数据生成速度可提升至10秒以内。

内容的提问来源于stack exchange，提问作者jfbaro