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

Question

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

我希望测试行的id 列大部分是连续的，所以想随机选择 几块 1,000 连续行，并将它们标记为测试行。我在 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