Postgres. Выборка N случайных записей

При работе над одним проектом возникла необходимость написать некое подобие тестовой системы. Задача формулировалась примерно так:

• из N записей в базе необходимо выбрать m (3-5) случайных строк в серии из k выборок (преимущественно k=2).

А теперь то же самое человеческим языком: из таблицы нужно два раза выбрать по 3-5 случайных записей. При этом не должно быть дубликатов и выборка должна происходить случайным образом.

Первое, что приходит в голову:

 SELECT *
  FROM data_set
  WHERE id NOT IN (1,2,3,4, 5)
  ORDER BY random()
  LIMIT 5;

И это даже будет работать. Вот только цена такого решения…

Поэтому пришлось воспользоваться «высшим разумом», который выдал намёк на решение.

WITH RECURSIVE r AS ( 
    WITH b AS (SELECT min(id), max(id) FROM table1) 
    ( 
        SELECT id, min, max, array[]::integer[] AS a, 0 AS n 
        FROM table1, b 
        WHERE id > min + (max - min) * random() 
        LIMIT 1 
    ) UNION ALL ( 
        SELECT t.id, min, max, a || t.id, r.n + 1 AS n 
        FROM table1 AS t, r 
        WHERE 
            t.id > min + (max - min) * random() AND 
            t.id <> all(a) AND 
            r.n + 1 < 10 
        LIMIT 1 
    ) 
) 
SELECT t.id FROM table1 AS t, r WHERE r.id = t.id;

Вот только решение это… «слегка» непонятное. А тащить непонятные решения в проект как-то неохота. Поэтому было решено пойти уже проверенным путём, где удалось найти пояснение сущности рекурсивных запросов.

Что же. Логика в общих чертах стала более-менее понятной: n раз выбираем по одной строке со случайным смещением от минимального значения ID (primary key). Таким образом запрос затрагивает максимум N строк вместо полного перебора содержимого таблицы.

Но «гладко было на бумаге». При попытке использования «как есть» (с поправкой на имена таблицы/полей) всплыли «сюрпризы»:

1. Массив в строке 4 создаётся пустым. Из-за этого в финальной выборке могут оказываться дубликаты. Скажем, запрос в строчках 4-7 возвращает id==5. После этого отрабатывает запрос в UNION блоке (строчки 9-15) и на каком-то этапе возвращает то же id==5, поскольку предыдущее значение id не попало в массив "а" и проверка в строке 13 «t.id <> all(a)» прошла успешно;
2. Количество значений на «выходе» было не более чем указано (стр. 14). Но меньше этого — запросто. Даже если оное количество гарантировано было в таблице. Иногда возвращалась пустая выборка (количество значений «0»).

С первой «особенностью» справиться оказалось сравнительно легко — достаточно было чуть поменять строчку с инициализацией массива. Примерно так:

SELECT id, min, max, array[]::integer[] || id AS a, 0 AS n 

А вот второй пункт заставил пораскинуть мозгами. Подвох обнаружился в самом «сердце» алгоритма — выборке записи из диапазона. Действительно, рекурсивный запрос пытается выбрать строчку, для которой выполнялось бы условие:

id > min + (max - min) * random()

Но в случае, когда random() возвращает «1», это условие трансформируется в:

id > max

Естественно, что в этом случае запрос ничего не найдёт и прекратит работу. А ежели такое случится при первом же проходе запроса, то «на выходе» будет пусто. Даже если в базе заведомо присутствуют нужные записи.

Первым побуждением было слегка изменить условие и привести его примерно к такому виду:

id >= min + (max - min) * random()

Это, так сказать, решение позволило получать минимум одну строчку на выходе. Но вовсе не гарантировало получения заданного количества строк в результате. Пришлось думать дальше.

После длительных размышлений, множества попыток, и литров стимулятора появился на свет
такой вот вариант:

код запроса

WITH RECURSIVE r AS (
 WITH b AS (
   SELECT
  min(t.id),
  (
   SELECT t.id
   FROM table1 AS t
   ORDER BY t.id DESC
   LIMIT 1
   OFFSET 9
  ) max
  FROM table1 AS t
 )
 (
  SELECT
    id, min, max, array[]::integer[] || id AS a, 0 AS n
  FROM table1 AS t, b
  WHERE
    id >= min + ((max - min) * random())::int
  LIMIT 1
 ) UNION ALL (
  SELECT t.id, min, max, a || t.id, r.n + 1 AS n
  FROM {table} AS t, r
  WHERE
   t.id >= min + ((max - min) * random())::int AND
   t.id <> all(a) AND
   r.n + 1 < 10
  LIMIT 1
 )
)
SELECT t.* FROM table1 AS t, r WHERE r.id = t.id

Вся соль в строчках 5-11. Т.е. дабы гарантировать, что на «выходе» будет N элементов нужно исходить из худшего из возможных сценариев. В данном случае — что random N раз подряд вернёт 1. Следовательно нужно знать/иметь N-1 значений перед max. Как этого достичь в запросе? Как вариант — отсортировать все записи по ID в порядке убывания и произвести смещение «вниз» на N строк. А поскольку в строках 19 и 25 используется ">=", то смещение можно указать на единицу меньше (N-1).

Вот и хорошо — основные затруднения решены и остались «мелочи»: запрос в текущем виде мало полезен. Нужно ведь делать выборку с учётом некоторых условий. В простейшем случае — исключать из выборки ID записей, выбранных на предыдущем этапе. Кроме того, нельзя исключать возможность использования каких-то дополнительных условий налагаемых на строки в таблице (is_active = true, is_deleted=false, ...). После некоторых размышлений становится понятно, что условия придётся ставить во все существенные части запроса (практически все подзапросы). Примерно как в следующем шаблоне:

код шаблона

WITH RECURSIVE r AS (
 WITH b AS (
   SELECT
  min(t.{pk}),
  (
   SELECT t.{pk}
   FROM {table} AS t
   WHERE {exclude} t.is_active
   ORDER BY t.{pk} DESC
   LIMIT 1
   OFFSET {limit} - 1
  ) max
  FROM {table} AS t WHERE {exclude} t.is_active
 )
 (
  SELECT
    t.{pk}, min, max, array[]::integer[] || t.{pk} AS a, 0 AS n
  FROM {table} AS t, b
  WHERE
    t.{pk} >= min + ((max - min) * random())::int AND
    {exclude}
    t.is_active
  LIMIT 1
 ) UNION ALL (
  SELECT t.{pk}, min, max, a || t.{pk}, r.n + 1 AS n
  FROM {table} AS t, r
  WHERE
   t.{pk} >= min + ((max - min) * random())::int AND
   t.{pk} <> all(a) AND
   r.n + 1 < {limit} AND
   {exclude}
   t.is_active
  LIMIT 1
 )
)
SELECT {fields} FROM {table} AS t, r WHERE r.{pk} = t.{pk}

Тут в фигурных скобках указаны именованные параметры, подлежащие замене:

• {table} — название таблицы;
• {pk} — имя PrimaryKey-поля;
• {fields} — список полей для выборки (можно указать и "*");
• {exclude} — условие (набор условий) для исключения записей из выборки. Например «t.id NOT IN (1,2,3,4)»;
• {limit} — количество записей в финальной выборке

И, наконец, последний в списке, но не по важности вопрос: «стоит ли овчинка выделки»? Каков «профит» от этой мозголомной конструкции? Будем проверять.

Для начала создадим таблицу, в которой будут проводиться эксперименты.

DROP TABLE IF EXISTS ttbl;
CREATE TABLE IF NOT EXISTS ttbl
(
  id serial NOT NULL,
  is_active BOOL NOT NULL DEFAULT true,
  CONSTRAINT ttbl_pkey PRIMARY KEY (id)
)
WITH (OIDS=FALSE);

Теперь наполним её данными. При этом нужно, чтобы значения id не шли последовательно, а имели «дырки». Т.е. не «1, 2, 3, 4, 5...» а хотя-бы «1, 4, 5, 8....». Для этого набросаем простенький скриптик.

код скрипта

import random

import psycopg2


DB_TABLE = 'ttbl'
PK_NAME = 'id'
DATABASES = {
    'NAME': 'test_db',
    'USER': 'user_test',
    'PASSWORD': 'R#Q9Jw*ZEKWOhBX+EP|3/xGkQn3',
    'HOST': 'localhost',
    'PORT': '',
}

TOTAL = 10000000
current = 0
step = 10000
dev = 8
while current <= TOTAL:
    data = set()
    for x in range(current, current+step, 10):
        data.add(x + int(random.random() * dev))

    x = cur.execute(
        "INSERT INTO {t_table} VALUES {t_items};".format(
                t_table=DB_TABLE,
                t_items='(' + '), ('.join(map(str, data)) + ')'
        )
    )
    current += step
    print(x, current)
    
cur.execute('COMMIT;')

Как видно из кода — каждый запрос вставляет по сотне записей. Значения меняются с шагом примерно «10». Примерно, т.к. каждое из значений может отклоняться на случайную величину в диапазоне 0-dev. Т.е. при двух последовательных значениях x «340» и «350» в таблицу могут быть вставлены любые значения из диапазона 340-348 и 350-358 (342, 347, 340...; 351, 355, 358...).

Итого в таблице оказалось

select count(id) from ttbl;

1001000 записей

Довольно приличное количество. Попробуем сделать выборку. Понятно, что одиночная выборка не показатель. Поэтому произведём серию последовательных запусков. Для определённости — серию из 5 запусков запросов каждого типа. Результаты сведём в таблицу и посчитаем среднее.

Сначала простой запрос

select t.*
from ttbl t
where 
  t.id not in (1, 3, 10, 89, 99, 22, 24, 25, 28, 30) 
  AND t.is_active
order by random()
limit 5;

Результаты:

№ п/п	время, ms
1	697
2	605
3	624
4	593
5	611

Как видно, среднее время выполнения запроса около* 600ms.

А сейчас — «монстр».

смотреть монстра

WITH RECURSIVE r AS (
  WITH b AS (
    SELECT
    min(t.id),
    (
      SELECT t.id
      FROM ttbl AS t
      WHERE
        t.id NOT IN (1, 3, 10, 89, 99, 22, 24, 25, 28, 30)
        AND t.is_active
      ORDER BY t.id DESC
      LIMIT 1
      OFFSET 5 - 1
    ) max
    FROM ttbl AS t
    WHERE 
      t.id NOT IN (1, 3, 10, 89, 99, 22, 24, 25, 28, 30)
      AND t.is_active
  )
  (
    SELECT
      id, min, max, array[]::integer[] || id AS a, 0 AS n
    FROM ttbl AS t, b
    WHERE
      id >= min + ((max - min) * random())::int AND
      t.id NOT IN (1, 3, 10, 89, 99, 22, 24, 25, 28, 30) AND
      t.is_active
    LIMIT 1
  ) UNION ALL (
    SELECT t.id, min, max, a || t.id, r.n + 1 AS n
    FROM ttbl AS t, r
    WHERE
      t.id > min + ((max - min) * random())::int AND
      t.id <> all( a ) AND
      r.n + 1 < 5 AND
      t.id NOT IN (1, 3, 10, 89, 99, 22, 24, 25, 28, 30) AND
      t.is_active
    LIMIT 1
  )
)
SELECT t.* FROM ttbl AS t, r WHERE r.id = t.id

Результаты:

№ п/п	время, ms
1	15
2	17
3	8
4	12
5	12

Среднее время около* 15ms.

Итого разница примерно на полтора порядка (в 40-50 раз). Оно таки того стоило.

Запросы проверялись в т.ч. и при «холодном» старте (после перезапуска машины/демона). И хотя время выполнения в абсолютном выражении менялось, но соотношение оставалось постоянным (насколько это возможно). Т.е. рекурсивный запрос всегда в несколько десятков раз был быстрее решения «в лоб».

Примечания

*Около, т.к. точное значение роли не играет из-за отклонений, вызванных кешем Postgres-а, влиянием операционной системы, железа, остального софта и т.д.