SQL HowTo: собираем «цепочки» с помощью window functions

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

Это может быть как условие от данных самой записи, так и сложное выражение относительно одной или нескольких предыдущих записей — например, длина интервала между близкими временными отсчетами.



Традиционные решения предусматривают разные варианты «self join», когда выборка соединяется с собой же, либо использование некоторых фактов «за пределами данных» — например, что записи должны иметь строго определенный шаг (N+1, «за каждый день», ...).

Первый вариант зачастую приводит к квадратичной сложности алгоритма от количества записей, что недопустимо на больших выборках, а второй может легко «развалиться», если каких-то отсчетов в исходных данных вдруг не окажется.

Но эту задачу нам помогут эффективно решить оконные функции в PostgreSQL.

Задача: считаем чужие деньги

Рассмотрим самый простой случай цепочки — когда условие непрерывности определяется данными самой записи.

Все дальнейшие операции не обязательно делать отдельно. Но ради наглядности алгоритма я разобью его на последовательные шаги, а уже в конце покажу, что и как можно оптимизировать.

Давайте представим, что у нас есть маленький банк, который ведет в таблице балансы по счетам клиентов. Как только происходит приходно-расходная операция — этой датой и фиксируется итоговая сумма счета на конец дня.

После длинных новогодних каникул банк решил вознаградить своих клиентов — и каждому открывшему счет в этом году дополнительно начислить +1% от среднесуточного остатка за самый длинный непрерывный период, когда счет не «обнулялся».

Вот он наш критерий непрерывности «цепочки». Ну а упорядоченность данных будет определяться датами балансов.

Нам принесли вот такой CSV, и попросили быстро подсчитать, кому и в каком размере такая щедрость от банка должна достаться:

date;client;balance
01.01.2020;Алиса;150
01.01.2020;Боб;100
02.01.2020;Алиса;100
02.01.2020;Боб;150
03.01.2020;Алиса;200
05.01.2020;Алиса;0
06.01.2020;Алиса;50
08.01.2020;Алиса;0
08.01.2020;Боб;200
09.01.2020;Алиса;0
09.01.2020;Боб;0
10.01.2020;Алиса;5

Сразу отметим несколько фактов, заметных на этих данных:

• 07.01 был праздник, и банк не работал. Поэтому ни у кого из клиентов записей об изменении баланса в этот день нет, а деньги на счетах — есть. То есть «переборные» алгоритмы, итерирующие по дням, уже нормально не пройдут.
• 04.01 Алиса не проводила никаких операций, поэтому записи нет. Но до 05.01 сумма на счету у нее была ненулевая — это придется учесть при анализе.
• Мы проводим анализ за 01.01-12.01, но баланс счета Алисы на конец этого периода ненулевой. Учтем и необходимость ограничения периода.

CSV-to-table

Самый правильный путь для импорта из CSV — воспользоваться оператором COPY. Но мы для разминки попробуем сделать это через регулярные выражения:

CREATE TEMPORARY TABLE tbl AS
SELECT
  to_date(prt[1], 'DD.MM.YYYY') dt
, prt[2] client
, prt[3]::numeric(32,2) balance
FROM
  (
    SELECT
      regexp_split_to_array(str, ';') prt
    FROM
      (
        SELECT
          regexp_split_to_table(
$$
date;client;balance
01.01.2020;Алиса;150
01.01.2020;Боб;100
02.01.2020;Алиса;100
02.01.2020;Боб;150
03.01.2020;Алиса;200
05.01.2020;Алиса;0
06.01.2020;Алиса;50
08.01.2020;Алиса;0
08.01.2020;Боб;200
09.01.2020;Алиса;0
09.01.2020;Боб;0
10.01.2020;Алиса;5
$$
        , E'\\n') str
      ) T
    WHERE
      str <> ''
    OFFSET 1
  ) T;

Это «нечестный» способ в том смысле, что не переварит корректно, например, экранирование разделителя в теле поля. Но для большинства простых применений — подходит.

Шаг 1: Фиксируем прикладное условие

В нашем случае условие непрерывности цепочки — ненулевой баланс. Выведем его отдельным полем, для наглядности хронологически упорядочивая по клиенту:

SELECT
  *
, balance > 0 cond
FROM
  tbl
ORDER BY
  client, dt;

dt         | client | balance | cond
------------------------------------
2020-01-01 | Алиса  |  150.00 | t
2020-01-02 | Алиса  |  100.00 | t
2020-01-03 | Алиса  |  200.00 | t
2020-01-05 | Алиса  |    0.00 | f
2020-01-06 | Алиса  |   50.00 | t
2020-01-08 | Алиса  |    0.00 | f
2020-01-09 | Алиса  |    0.00 | f
2020-01-10 | Алиса  |    5.00 | t
2020-01-01 | Боб    |  100.00 | t
2020-01-02 | Боб    |  150.00 | t
2020-01-08 | Боб    |  200.00 | t
2020-01-09 | Боб    |    0.00 | f

Шаг 2: Вычисляем недостающее

Обратим внимание, что сумма у Боба не менялась с 02.01 по 08.01. А по условию задачи мы должны вычислить именно среднесуточный остаток — то есть нам необходима информация об этих «пропущенных» днях. Или хотя бы само количество таких дней, когда значение оставалось одинаковым:

coalesce(lead(dt) OVER(PARTITION BY client ORDER BY dt), '2020-01-12') - dt days

dt         | client | balance | cond | days
-------------------------------------------
2020-01-01 | Алиса  |  150.00 | t    |    1
2020-01-02 | Алиса  |  100.00 | t    |    1
2020-01-03 | Алиса  |  200.00 | t    |    2
2020-01-05 | Алиса  |    0.00 | f    |    1
2020-01-06 | Алиса  |   50.00 | t    |    2
2020-01-08 | Алиса  |    0.00 | f    |    1
2020-01-09 | Алиса  |    0.00 | f    |    1
2020-01-10 | Алиса  |    5.00 | t    |    2
2020-01-01 | Боб    |  100.00 | t    |    1
2020-01-02 | Боб    |  150.00 | t    |    6
2020-01-08 | Боб    |  200.00 | t    |    1
2020-01-09 | Боб    |    0.00 | f    |    3

С помощью оконной функции lead() мы узнали дату из следующей по порядку записи, а через coalesce ограничили интервал для последней. Заодно воспользовались полезным свойством, что разность двух дат в PostgreSQL возвращает целочисленное количество дней между ними.

В качестве почти бесплатного бонуса мы получили всю ту же информацию и для записей с нулевым балансом. Но если строк с невыполняющимся условием, которые нас не интересуют, достаточно много, имеет смысл такие вычисления загнать под CASE, чтобы сэкономить ресурсы сервера.

Шаг 3: Находим точки разрывов

Начало каждой интересующей нас цепочки — это точка, где значение вычисленного ранее условия меняется относительно предыдущей записи. Воспользуемся функцией lag(), чтобы найти такие точки:

lag(cond) OVER(PARTITION BY client ORDER BY dt) IS DISTINCT FROM cond chain_start

dt         | client | balance | cond | days | chain_start
---------------------------------------------------------
2020-01-01 | Алиса  |  150.00 | t    |    1 | t
2020-01-02 | Алиса  |  100.00 | t    |    1 | f
2020-01-03 | Алиса  |  200.00 | t    |    2 | f
2020-01-05 | Алиса  |    0.00 | f    |    1 | t
2020-01-06 | Алиса  |   50.00 | t    |    2 | t
2020-01-08 | Алиса  |    0.00 | f    |    1 | t
2020-01-09 | Алиса  |    0.00 | f    |    1 | f
2020-01-10 | Алиса  |    5.00 | t    |    2 | t
2020-01-01 | Боб    |  100.00 | t    |    1 | t
2020-01-02 | Боб    |  150.00 | t    |    6 | f
2020-01-08 | Боб    |  200.00 | t    |    1 | f
2020-01-09 | Боб    |    0.00 | f    |    3 | t

С помощью оператора IS DISTINCT FROM вместо <> мы избежали проблем сравнения с NULL для первых записей по каждому из клиентов. Соответственно, все строки, где значение TRUE — начало новой цепочки, а FALSE — ее продолжение.

Шаг 4: Нанизываем звенья

Чтобы сгруппировать данные в рамках каждой отдельной цепочки, проще всего присвоить всем ее записям один и тот же идентификатор. В качестве него отлично подходит порядковый номер самой цепочки. А он как раз равен количеству «начал» цепочек, встретившихся выше по выборке.

Их можно посчитать или через «оконное» суммирование bool-значений sum({boolean}::integer), или через подсчет количества записей, подходящих под условие count(*) FILTER(WHERE {boolean}). Воспользуемся вторым вариантом:

count(*) FILTER(WHERE chain_start) OVER(PARTITION BY client ORDER BY dt) grpid

dt         | client | balance | cond | days | chain_start | grpid
-----------------------------------------------------------------
2020-01-01 | Алиса  |  150.00 | t    |    1 | t           |     1
2020-01-02 | Алиса  |  100.00 | t    |    1 | f           |     1
2020-01-03 | Алиса  |  200.00 | t    |    2 | f           |     1
2020-01-06 | Алиса  |   50.00 | t    |    2 | t           |     2
2020-01-10 | Алиса  |    5.00 | t    |    2 | t           |     3
2020-01-01 | Боб    |  100.00 | t    |    1 | t           |     1
2020-01-02 | Боб    |  150.00 | t    |    6 | f           |     1
2020-01-08 | Боб    |  200.00 | t    |    1 | f           |     1

На этом шаге длину всех звеньев каждой цепочки мы уже знаем, «неинтересные» записи нам больше не нужны, поэтому просто отфильтруем их.

Шаг 5: Собираем цепочки

Чтобы вычислить среднее по всем дням в цепочке, нам потребуется суммарное количество дней и «интегральный» баланс:

SELECT
  client
, min(dt) chain_dt
, sum(days * balance) balance
, sum(days) days
FROM
  ...
GROUP BY
  1, grpid
ORDER BY
  1, grpid;

client | chain_dt   | balance | days
-------------------------------------
Алиса  | 2020-01-01 |  650.00 |    4
Алиса  | 2020-01-06 |  100.00 |    2
Алиса  | 2020-01-10 |   10.00 |    2
Боб    | 2020-01-01 | 1200.00 |    8

Шаг 6: Ищем прикладные максимумы

С помощью DISTINCT ON оставим единственную запись (с максимальным значением days) по каждому клиенту:

SELECT DISTINCT ON(client)
  *
FROM
  ...
ORDER BY
  client, days DESC;

client | chain_dt   | balance | days
-------------------------------------
Алиса  | 2020-01-01 |  650.00 |    4
Боб    | 2020-01-01 | 1200.00 |    8

Собственно, на этом — все, осталось только…

Объединяем и оптимизируем

Итоговый запрос

WITH step123 AS (
  SELECT
    *
  , CASE
      WHEN cond THEN
        lag(cond) OVER(w) IS DISTINCT FROM cond
    END chain_start
  , CASE
      WHEN cond THEN
        coalesce(lead(dt) OVER(w), '2020-01-12') - dt
    END days
  FROM
    tbl
  , LATERAL(SELECT balance > 0 cond) T
  WINDOW
    w AS (PARTITION BY client ORDER BY dt)
)
, step4 AS (
  SELECT
    *
  , count(*) FILTER(WHERE chain_start) OVER(PARTITION BY client ORDER BY dt) grpid
  FROM
    step123
  WHERE
    cond
)
SELECT DISTINCT ON(client)
    client
  , sum(days) OVER(w) days
  , min(dt) OVER(w) chain_dt
  , sum(days * balance) OVER(w) balance
FROM
  step4
WINDOW
  w AS (PARTITION BY client, grpid)
ORDER BY
  1, 2 DESC;

Здесь мы объединили и оптимизировали первые три шага:

• LATERAL-подзапрос позволил нам вычислить дополнительное поле без лишнего прохода по выборке и сразу использовать его в функции
• вынос общего определения под WINDOW помогает PostgreSQL не делать двойную сортировку для формирования «окна» и вычислить обе функции в одном WindowAgg-узле
• «ленивое» вычисление функции под CASE уменьшает количество производимых операций

Аналогично мы объединили и следующие два шага. Но порядок «окна» вычисления агрегатов (client, grpid) и уникализации (client, sum(days)) не совпал, поэтому Sort-узлов в последем блоке останется все-таки два — перед WindowAgg и перед Unique.


[посмотреть на explain.tensor.ru]

Замечу, что при нумерации цепочек сначала отрабатывает WHERE-условие, поэтому генерируемые оконной функцией номера оказываются последовательными.