sql - 在 postgres 中有效地标记行的连续子集

Question

我有一个包含数十亿行的 Postgres 表，对于机器学习应用程序，需要将其分为训练集和测试集。

我希望测试行通过它们的 id 列大部分是连续的，所以想随机选择每个 1,000 个连续行的几个 block ，并将它们标记为测试行。我在 id 列上有一个索引，因此可以快速选择任意 1,000 个连续的行:

UPDATE table SET test=true WHERE id BETWEEN 100000 AND 101000;

非常高效，并且如您所料使用索引扫描。不幸的是，一旦我随机生成初始 id，即

WITH off AS (SELECT ROUND(random()*maxId))
  UPDATE table SET test=true
    WHERE id BETWEEN (SELECT * FROM off LIMIT 1)
                 AND (SELECT * FROM off LIMIT 1)+1000;

查询规划器现在决定进行全表扫描(慢得多)。

当然，如果我只需要这样做一次，我会手动生成一个随机行，没问题。但是最后我想要一个自动分为测试和训练的功能，如下所示:

CREATE OR REPLACE FUNCTION test_train_divide(chunkSize integer, proportion real)
RETURNS BOOLEAN
AS $$
DECLARE
maxId INTEGER := (SELECT MAX(id) FROM table);
BEGIN
FOR i IN 1 .. round(maxId*proportion/chunkSize) LOOP
  RAISE NOTICE 'Update call %', i;
  WITH off AS (SELECT ROUND(random()*maxId))
  UPDATE table SET test=true
    WHERE id BETWEEN (SELECT * FROM off LIMIT 1)
                 AND (SELECT * FROM off LIMIT 1)+chunkSize;
END LOOP;
return true;
END;
$$ LANGUAGE plpgsql;

SELECT test_train_divide(1000,0.01);

这有效，但速度非常慢!有什么指点吗？

更新

这是架构

    tbl "public.tbl”
  Column   |  Type   | Modifiers 
-----------+---------+-----------
 subid     | integer | 
 id        | bigint  | 
 wordid    | integer | 
 capid     | integer | 
 test      | boolean | 
Indexes:
    “tbl_id_idx" btree (id)

这里有两种不同的查询计划，一种是好的(使用索引)，一种是坏的:

will=# EXPLAIN UPDATE tbl SET test=true WHERE id BETWEEN 1000000 AND 1001000;

                                            QUERY PLAN                                             
---------------------------------------------------------------------------------------------------
 Update on tbl  (cost=0.57..790.45 rows=1079 width=38)
   ->  Index Scan using tbl_id_idx on tbl  (cost=0.57..790.45 rows=1079 width=38)
         Index Cond: ((id >= 1000000) AND (id <= 1001000))
(3 rows)


will=# EXPLAIN WITH start AS (SELECT round(random()*max(id)) FROM tbl) UPDATE tbl c SET test=true WHERE c.id BETWEEN (SELECT * FROM start LIMIT 1) AND (SELECT * FROM start LIMIT 1)+1000;



                                                             QUERY PLAN                                                               
---------------------------------------------------------------------------------------------------------------------------------------
 Update on tbl c  (cost=0.65..14932243.97 rows=1459961 width=38)
   CTE start
     ->  Result  (cost=0.59..0.61 rows=1 width=0)
           InitPlan 1 (returns $0)
             ->  Limit  (cost=0.57..0.59 rows=1 width=8)
                   ->  Index Only Scan Backward using tbl_id_idx on tbl  (cost=0.57..5846291.90 rows=288468819 width=8)
                         Index Cond: (id IS NOT NULL)
   InitPlan 3 (returns $2)
     ->  Limit  (cost=0.00..0.02 rows=1 width=8)
           ->  CTE Scan on start  (cost=0.00..0.02 rows=1 width=8)
   InitPlan 4 (returns $3)
     ->  Limit  (cost=0.00..0.02 rows=1 width=8)
           ->  CTE Scan on start start_1  (cost=0.00..0.02 rows=1 width=8)
   ->  Seq Scan on tbl c  (cost=0.00..14932243.32 rows=1459961 width=38)
         Filter: (((id)::double precision >= $2) AND ((id)::double precision <= ($3 + 1000::double precision)))
(15 rows)

Time: 2.649 ms

Answer 1

初始化后max_id作为max(id) - 1000为 1000 行留出空间，这应该使用索引:

UPDATE table
SET    test = true
FROM  (SELECT (random() * max_id)::bigint AS lower_bound) t
WHERE  id BETWEEN t.lower_bound AND t.lower_bound + 999;

• 不需要 CTE 和子查询的复杂结构。使用单个子查询。

• 您的原始计算结果为 numeric (或 dp )，这可能不适用于 bigint 上的索引柱子。转换到bigint . (在 pg 9.3 中应该不是问题。)

• BETWEEN包括下限和上限。你的上限应该是 lower + 999 , 严格来说。

• random()返回 ( per documentation ) random value in the range 0.0 <= x < 1.0 .为了完全公平，你的lower_bound真的应该这样计算(假设没有间隙):

    trunc(random() * max_id)::bigint + 1
  

如果您需要真正的随机数(或者如果您的 id 有间隙)，请考虑以下相关答案:

• Best way to select random rows PostgreSQL

也许咨询锁或其他方法可能有用。比较这个相关的，稍后的答案:

• How to mark certain nr of rows in table on concurrent access