SQL HowTo: пишемо while-цикл прямо у запиті, або «Елементарна триходівка»

Періодично виникає завдання пошуку пов'язаних даних по набору ключів, поки не наберемо потрібну сумарну кількість записів.

Найбільш «життєвий» приклад – вивести 20 найстаріших завдань, що числяться на списку співробітників (наприклад, у межах одного підрозділу). Для різних управлінських «дашбордів» з короткими вичавками на ділянках роботи схожа тема потрібна досить часто.

У статті розглянемо реалізацію на PostgreSQL «наївного» варіанта вирішення такого завдання, «розумніший» і дуже складний алгоритм циклу на SQL з умовою виходу від знайдених данихщо може бути корисним як для загального розвитку, так і для застосування в інших схожих випадках.

Візьмемо тестовий набір даних із попередньої статті. Щоб записи, що виводяться, не «скакали» від разу до разу при збігу сортованих значень, розширимо предметний індекс додаванням первинного ключа. Заодно це відразу надасть йому унікальності, і гарантує нам однозначність порядку сортування:

CREATE INDEX ON task(owner_id, task_date, id);
-- а старый - удалим
DROP INDEX task_owner_id_task_date_idx;

Як чується, так і пишеться

Спочатку накидаємо найпростіший варіант запиту, передаючи ID виконавців масивом як вхідний параметр:

SELECT
  *
FROM
  task
WHERE
  owner_id = ANY('{1,2,4,8,16,32,64,128,256,512}'::integer[])
ORDER BY
  task_date, id
LIMIT 20;

[Подивитися на explain.tensor.ru]

Трохи сумно – ми замовляли лише 20 записів, а Index Scan повернув нам 960 рядків, Які потім ще й сортувати довелося ... А давайте спробуємо читати поменше.

unnest + ARRAY

Перше міркування, яке нам допоможе - якщо нам треба всього 20 відсортованих записів, достатньо читати не більше 20 відсортованих у тому ж порядку по кожному ключ. Благо, відповідний індекс (owner_id, task_date, id) ми маємо.

Скористаємося тим самим механізмом вилучення та «розвороту в стовпці» цілісного запису таблиці, що і в минулій статті. А також застосуємо пакунок у масив за допомогою функції ARRAY():

WITH T AS (
  SELECT
    unnest(ARRAY(
      SELECT
        t
      FROM
        task t
      WHERE
        owner_id = unnest
      ORDER BY
        task_date, id
      LIMIT 20 -- ограничиваем тут...
    )) r
  FROM
    unnest('{1,2,4,8,16,32,64,128,256,512}'::integer[])
)
SELECT
  (r).*
FROM
  T
ORDER BY
  (r).task_date, (r).id
LIMIT 20; -- ... и тут - тоже

[Подивитися на explain.tensor.ru]

О, вже набагато краще! На 40% швидше, і в 4.5 рази менше даних довелося читати.

Матеріалізація записів таблиць через CTEЗверну увагу, що в деяких випадках спроба одразу працювати з полями запису після її пошуку в підзапиті, без «обертання» в CTE може призводити до «множення» InitPlan пропорційно кількості цих полів:

SELECT
  ((
    SELECT
      t
    FROM
      task t
    WHERE
      owner_id = 1
    ORDER BY
      task_date, id
    LIMIT 1
  ).*);

Result  (cost=4.77..4.78 rows=1 width=16) (actual time=0.063..0.063 rows=1 loops=1)
  Buffers: shared hit=16
  InitPlan 1 (returns $0)
    ->  Limit  (cost=0.42..1.19 rows=1 width=48) (actual time=0.031..0.032 rows=1 loops=1)
          Buffers: shared hit=4
          ->  Index Scan using task_owner_id_task_date_id_idx on task t  (cost=0.42..387.57 rows=500 width=48) (actual time=0.030..0.030 rows=1 loops=1)
                Index Cond: (owner_id = 1)
                Buffers: shared hit=4
  InitPlan 2 (returns $1)
    ->  Limit  (cost=0.42..1.19 rows=1 width=48) (actual time=0.008..0.009 rows=1 loops=1)
          Buffers: shared hit=4
          ->  Index Scan using task_owner_id_task_date_id_idx on task t_1  (cost=0.42..387.57 rows=500 width=48) (actual time=0.008..0.008 rows=1 loops=1)
                Index Cond: (owner_id = 1)
                Buffers: shared hit=4
  InitPlan 3 (returns $2)
    ->  Limit  (cost=0.42..1.19 rows=1 width=48) (actual time=0.008..0.008 rows=1 loops=1)
          Buffers: shared hit=4
          ->  Index Scan using task_owner_id_task_date_id_idx on task t_2  (cost=0.42..387.57 rows=500 width=48) (actual time=0.008..0.008 rows=1 loops=1)
                Index Cond: (owner_id = 1)
                Buffers: shared hit=4"
  InitPlan 4 (returns $3)
    ->  Limit  (cost=0.42..1.19 rows=1 width=48) (actual time=0.009..0.009 rows=1 loops=1)
          Buffers: shared hit=4
          ->  Index Scan using task_owner_id_task_date_id_idx on task t_3  (cost=0.42..387.57 rows=500 width=48) (actual time=0.009..0.009 rows=1 loops=1)
                Index Cond: (owner_id = 1)
                Buffers: shared hit=4

Один і той самий запис «пошукався» 4 рази… Аж до PostgreSQL 11 така поведінка зустрічається регулярно, і рішенням є «обертання» в CTE, що є безумовною межею для оптимізатора в цих версіях.

Рекурсивний акумулятор

У попередньому варіанті сумарно ми прочитали 200 рядків заради потрібних 20. Вже не 960, але ще менше – можна?

Спробуємо скористатися знанням, що нам треба всього 20 записів. Тобто ітеруватимемо вичитування даних лише до досягнення потрібної нам кількості.

Крок 1: стартовий список

Очевидно, що наш «цільовий» список із 20 записів має починатися з «перших» записів по одному з наших owner_id-ключів. Тому спочатку знайдемо такі «найперші» по кожному з ключів і занесемо до списку, відсортувавши його в порядку, що хочемо - (task_date, id).

Крок 2: знаходимо «наступні» записи

Тепер, якщо ми візьмемо з нашого списку перший запис та почнемо крокувати далі за індексом зі збереженням owner_id-ключа, то всі знайдені записи - якраз наступні в результуючій вибірці. Звичайно, тільки поки ми не перетнемо прикладний ключ другий запис у списку.

Якщо вийшло, що ми другий запис «перетнули», то останній прочитаний запис повинен бути доданий до списку замість першого (з тим самим owner_id), після чого список знову пересортуємо.

Тобто у нас весь час виходить, що в списку є не більше одного запису по кожному з ключів (якщо записи скінчилися, а ми не перетнули, то зі списку перший запис просто пропаде і нічого не додасться), причому вони завжди відсортовані у порядку зростання прикладного ключа (task_date, id).

Крок 3: фільтруємо та «розвертаємо» записи

У частині рядків нашої рекурсивної вибірки деякі записи rv дублюються - спочатку ми знаходимо такі як «перетинає кордон 2-го запису списку», а потім підставляємо як 1-й зі списку. Так ось перша поява треба відфільтрувати.

Страшний підсумковий запит

WITH RECURSIVE T AS (
  -- #1 : заносим в список "первые" записи по каждому из ключей набора
  WITH wrap AS ( -- "материализуем" record'ы, чтобы обращение к полям не вызывало умножения InitPlan/SubPlan
    WITH T AS (
      SELECT
        (
          SELECT
            r
          FROM
            task r
          WHERE
            owner_id = unnest
          ORDER BY
            task_date, id
          LIMIT 1
        ) r
      FROM
        unnest('{1,2,4,8,16,32,64,128,256,512}'::integer[])
    )
    SELECT
      array_agg(r ORDER BY (r).task_date, (r).id) list -- сортируем список в нужном порядке
    FROM
      T
  )
  SELECT
    list
  , list[1] rv
  , FALSE not_cross
  , 0 size
  FROM
    wrap
UNION ALL
  -- #2 : вычитываем записи 1-го по порядку ключа, пока не перешагнем через запись 2-го
  SELECT
    CASE
      -- если ничего не найдено для ключа 1-й записи
      WHEN X._r IS NOT DISTINCT FROM NULL THEN
        T.list[2:] -- убираем ее из списка
      -- если мы НЕ пересекли прикладной ключ 2-й записи
      WHEN X.not_cross THEN
        T.list -- просто протягиваем тот же список без модификаций
      -- если в списке уже нет 2-й записи
      WHEN T.list[2] IS NULL THEN
        -- просто возвращаем пустой список
        '{}'
      -- пересортировываем словарь, убирая 1-ю запись и добавляя последнюю из найденных
      ELSE (
        SELECT
          coalesce(T.list[2] || array_agg(r ORDER BY (r).task_date, (r).id), '{}')
        FROM
          unnest(T.list[3:] || X._r) r
      )
    END
  , X._r
  , X.not_cross
  , T.size + X.not_cross::integer
  FROM
    T
  , LATERAL(
      WITH wrap AS ( -- "материализуем" record
        SELECT
          CASE
            -- если все-таки "перешагнули" через 2-ю запись
            WHEN NOT T.not_cross
              -- то нужная запись - первая из спписка
              THEN T.list[1]
            ELSE ( -- если не пересекли, то ключ остался как в предыдущей записи - отталкиваемся от нее
              SELECT
                _r
              FROM
                task _r
              WHERE
                owner_id = (rv).owner_id AND
                (task_date, id) > ((rv).task_date, (rv).id)
              ORDER BY
                task_date, id
              LIMIT 1
            )
          END _r
      )
      SELECT
        _r
      , CASE
          -- если 2-й записи уже нет в списке, но мы хоть что-то нашли
          WHEN list[2] IS NULL AND _r IS DISTINCT FROM NULL THEN
            TRUE
          ELSE -- ничего не нашли или "перешагнули"
            coalesce(((_r).task_date, (_r).id) < ((list[2]).task_date, (list[2]).id), FALSE)
        END not_cross
      FROM
        wrap
    ) X
  WHERE
    T.size < 20 AND -- ограничиваем тут количество
    T.list IS DISTINCT FROM '{}' -- или пока список не кончился
)
-- #3 : "разворачиваем" записи - порядок гарантирован по построению
SELECT
  (rv).*
FROM
  T
WHERE
  not_cross; -- берем только "непересекающие" записи

[Подивитися на explain.tensor.ru]

Таким чином, ми обміняли 50% читань даних на 20% часу виконання. Тобто якщо у вас є причини вважати, що читання може бути довгим (наприклад, дані часто не в кеші і доводиться за ними ходити на диск), то в такий спосіб можна залежати від читання менше.

У будь-якому випадку час виконання вийшов краще, ніж у «наївному» першому варіанті. Але яким із цих 3 варіантів користуватися - вибирати вам.

Джерело: habr.com