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

    Похожие публикации

    Комментарии 3

      0

      Спасибо за статью. Так как я не пишу SQL профессионально, но люблю порешать задачки и поучить новое, было интересно.




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

      При попытке вставить данные с помощью COPY может внезапно оказаться, что 01.01.2020 — это 1 января, а 02.01.2020 — это 1 февраля.


      пруфы
      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

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

      Чтобы пересилить это дело, можно указать явно порядок (день-месяц-год) с помощью set datestyle to 'DMY' (подробнее здесь). Однако так как в статье в разных местах используется разный формат данных ('02.01.2020' и '2020-01-12'), с этим тоже могут возникнуть проблемы.

        0
        Однако так как в статье в разных местах используется разный формат данных ('02.01.2020' и '2020-01-12'), с этим тоже могут возникнуть проблемы.
        Это как раз не большая проблема, поскольку решается переключением datestyle сначала «туда» (DMY), как было правильно подмечено, а потом «обратно» (YMD).
        0
        не та ветка

        Только полноправные пользователи могут оставлять комментарии. Войдите, пожалуйста.

        Самое читаемое