PostgreSQL如何用FOR循环实现按所选属性各抽取1000条样本

常见问题说明 #

你写的匿名DO块无法直接返回结果，是因为PostgreSQL的DO语句属于匿名代码块，没有返回值设计，仅用于执行无结果返回的过程操作，另外你代码中的LIKE没有加通配符，和精确匹配=效果一致，建议直接用=避免不必要的匹配逻辑。

方案1：窗口函数实现（最简洁，推荐大多数场景使用） #

不需要循环、不需要动态SQL，仅用普通查询即可实现，后续新增属性只需要修改IN列表中的属性值即可：

SELECT * FROM (
  SELECT 
    *,
    ROW_NUMBER() OVER (PARTITION BY atribut ORDER BY random()) AS rn
  FROM data_table
  WHERE atribut IN ('atrib1', 'atrib2', 'atrib3') -- 新增属性直接在此处添加即可
) t
WHERE rn <= 1000;

• 实现逻辑：按atribut字段分组，给每个分组内的行随机排序生成序号，仅取每个分组序号前1000的行
• 优点：写法简单易维护，不需要额外创建存储对象

方案2：动态生成SQL（适配超大数据量场景） #

如果每个属性对应的行数达到百万级以上，窗口函数全量排序的性能会低于你原来的分段LIMIT方案，这时候可以用以下方式自动生成拼接好的UNION语句，不需要手写每一段子查询：

-- 执行以下查询会直接输出你需要的完整查询语句，复制即可执行
SELECT string_agg(
  format('(SELECT * FROM data_table WHERE atribut = %L ORDER BY random() LIMIT 1000)', atr),
  ' UNION ALL ' -- 确认无跨属性重复行时用UNION ALL比UNION性能高数倍
) || ';' AS final_query
FROM unnest(array['atrib1', 'atrib2', 'atrib3']) AS t(atr); -- 新增属性直接修改数组即可

如果需要封装成可反复调用的函数，可以使用以下写法：

CREATE OR REPLACE FUNCTION get_sample_by_attr(atribut_array text[])
RETURNS SETOF data_table AS $$
DECLARE
  atr text;
BEGIN
  FOREACH atr IN ARRAY atribut_array
  LOOP
    RETURN QUERY
    SELECT * FROM data_table
    WHERE atribut = atr
    ORDER BY random()
    LIMIT 1000;
  END LOOP;
  RETURN;
END;
$$ LANGUAGE plpgsql;

-- 调用方式
SELECT * FROM get_sample_by_attr(array['atrib1', 'atrib2', 'atrib3']);

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