Поиск дубликатов в Polars: почему нет раннего выхода и что работает быстрее

Почему связка is_duplicated + any в Polars LazyFrame не дает раннего выхода, как проверить и какие быстрые альтернативы выбрать для одной и нескольких колонок.

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

Постановка задачи

Рассмотрим вспомогательную функцию, которая пытается вернуть true при первом дубликате в подмножестве колонок. Идея — выйти раньше, комбинируя is_duplicated с any.

def try_fast_dupe_flag(lf: pl.LazyFrame, keys: list[str]) -> bool:
    """Попытаться завершить раньше при первом дубликате"""
    return (
        lf.select(
            pl.struct(keys).is_duplicated().any()
        )
        .collect()
        .item()
    )

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

Почему раннего выхода не происходит

Даже если дубликат обнаружен сразу, выражение не завершится раньше. Главная причина в том, что is_duplicated не знает о последующем вызове any в дереве выражений. У него нет информации, что достаточно одного true, чтобы остановиться, поэтому он продолжает вычисления на всем входе. Это означает, что положение дубликатов — в начале или в конце — не меняет общее время выполнения.

Это можно увидеть, замерив время для двух крупных конвейеров LazyFrame: с дубликатами в начале и с дубликатами в конце. Результаты практически одинаковы.

lead = (
    pl.select(x=pl.concat_list(0, pl.int_ranges(0, 100_000_000)))
      .explode('x')
      .lazy()
)
trail = (
    pl.select(x=pl.concat_list(pl.int_ranges(0, 100_000_000), 0))
      .explode('x')
      .lazy()
)
%%timeit
try_fast_dupe_flag(lead, 'x')
# 5.12 с ± 290 мс на цикл (среднее ± стандартное отклонение из 7 запусков, по 1 циклу)
%%timeit
try_fast_dupe_flag(trail, 'x')
# 5.14 с ± 157 мс на цикл (среднее ± стандартное отклонение из 7 запусков, по 1 циклу)

Поскольку вычисления не могут прерываться внутри графа выражений, стоимость в обоих случаях примерно одинаковая. Чтобы по‑настоящему делать ранний выход, нужен либо специальный плагин, который умеет останавливаться, либо map_batches с ручным циклом на Python — обычно он будет медленнее, за исключением крайних сценариев, где первые элементы сразу совпадают.

Ранний выход на Python через map_batches

Если хочется увидеть ранний выход в действии, можно сделать Python‑реализацию с map_batches. Однако такой подход не годится для продакшена из‑за циклов на Python и накладных расходов.

import time
def python_dupe_probe(s: pl.Series) -> pl.Series:
    start = time.time()
    for ii, lval in enumerate(s):
        for jj, rval in enumerate(s[ii + 1:]):
            if time.time() - start > 10:
                print(f"took too long, at i={ii} j={jj}")
                return pl.Series([None], dtype=pl.Boolean)
            if lval == rval:
                return pl.Series([True])
    return pl.Series([False])
lead.select(pl.struct('x').map_batches(python_dupe_probe)).collect()
# это занимает 0.0 с
trail.select(pl.struct('x').map_batches(python_dupe_probe)).collect()
# слишком долго; дошли до i=0 j=27725000

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

Быстрые альтернативы без раннего выхода

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

lead.select(pl.col('x').unique().len() != pl.col('x').len()).collect()
trail.select(pl.col('x').unique().len() != pl.col('x').len()).collect()
# каждый запрос занимает примерно 0.7 с

Но переход к unique на основе struct для одной колонки радикально замедляет вычисления.

lead.select(pl.struct('x').unique().len() != pl.struct('x').len()).collect()
trail.select(pl.struct('x').unique().len() != pl.struct('x').len()).collect()
# каждый запрос занимает примерно 1 мин 8 с

Для нескольких колонок хорошо работает прием с сортировкой и rle_id. Он быстрее, чем is_duplicated, хотя все равно уступает методу unique против len для одной колонки.

lead.select(
    pl.struct('x').sort().rle_id().max() + 1 != pl.struct('x').len()
).collect()
trail.select(
    pl.struct('x').sort().rle_id().max() + 1 != pl.struct('x').len()
).collect()
# каждый запрос занимает примерно 3 с

Идея проста: после сортировки равные подряд строки образуют «прогоны». Если все строки уникальны, прогонов нет, и максимальный rle_id ровно на единицу меньше общей длины, потому что индексация начинается с нуля.

Практическая функция

Объединив все вместе, можно сделать удобный хелпер, который переключается между очень быстрым вариантом для одной колонки и приемом с rle_id для нескольких. Раннего выхода он не делает, но на практике работает быстрее исходной связки is_duplicated + any.

def dupe_exists(lazy_ds: pl.LazyFrame, keys: list[str]) -> bool:
    """Раннего выхода нет, но на практике быстро"""
    if len(keys) == 1:
        return (
            lazy_ds.select(
                pl.col(keys[0]).unique().len() != pl.col(keys[0]).len()
            )
            .collect()
            .item()
        )
    else:
        return (
            lazy_ds.select(
                pl.struct(keys).sort().rle_id().max() + 1 != pl.struct(keys).len()
            )
            .collect()
            .item()
        )

Почему это важно

Работая с выражениями Polars LazyFrame, легко предположить, что такие операторы, как any, заставят последующие узлы прерывать вычисления. В реальности узлы вроде is_duplicated не осведомлены о дальнейших агрегациях в цепочке и поэтому считают весь вход. Понимание этого помогает избежать ловушек производительности и понять, когда вообще стоит задумываться о ручном раннем выходе. Если дубликаты появляются в начале только изредка, цикл на Python через map_batches вряд ли окупится. Когда же совпадения стабильно находятся в самом начале, может понадобиться специализированный плагин или собственная логика.

Выводы

Ранний выход звучит привлекательно, но выражения Polars не будут прерываться только потому, что где‑то дальше в конвейере стоит any. Для одной колонки самый простой и эффективный способ найти дубликаты — сравнить длину множества уникальных значений с общей длиной. Для нескольких колонок хороший вариант — отсортировать и применить rle_id, не опускаясь до циклов на Python. Если требуется настоящий ранний выход, остаются два пути: специализированный плагин или ручная реализация через map_batches, причем второй вариант имеет смысл лишь в крайних случаях.