Собираем «цепочки» с помощью 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-условие, поэтому генерируемые оконной функцией номера оказываются последовательными.