Как сделать случайную выборку в Polars без переполнения памяти

Почему sample в Polars LazyFrame материализует широкий столбец и приводит к OOM, и как выбрать строки по индексам с потоковым выполнением. Пример кода.

Выборка из огромного столбца в Polars LazyFrame выглядит обманчиво простой, пока процесс не приводит к взрывному росту использования памяти. Если датасет велик и задействован один широкий столбец (например, сырой текст веб‑страниц), наивный вызов sample на LazyFrame всё равно может спровоцировать полную материализацию и аварийное завершение. Ниже — логика, стоящая за этим, и альтернатива, которая помогает в сценариях с жёсткими ограничениями по памяти.

Настройка и неудачный подход

Данные читаются из Parquet с помощью pl.scan_parquet(...) — полный DataFrame в память не помещается. Выборка одного большого столбца приводит к сбою конвейера даже при размере выборки 1, тогда как та же операция успешно отрабатывает на меньшем столбце.

import polars as pl
src_path = "path/to/data.parquet"
lz = pl.scan_parquet(src_path)
samp_n = 1  # примерный размер; даже это приводит к переполнению памяти, если столбец огромный
subset = lz.select(
    pl.col("content_col").sample(n=samp_n, seed=0)
)
# (1) запись за один проход в этом сценарии завершается ошибкой
subset.sink_parquet("subset.parquet")
# (2) сбор (collect) результирующей выборки тоже завершается ошибкой
subset.collect()

Что происходит на самом деле

В целом collect по умолчанию должен использовать потоковый движок, но sample, вероятно, не оптимизирован под стриминг в этом контексте и может заставить движок считать весь столбец целиком в память. Для очень большого текстоподобного столбца этого достаточно, чтобы выйти за пределы доступной ОЗУ, даже если итоговая выборка сама по себе мала.

Более экономный по памяти путь

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

import polars as pl
import numpy as np
plan = pl.LazyFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
k = 2
total_rows = plan.select(pl.len()).collect().item()
take_idx = sorted(np.random.choice(total_rows, size=k, replace=False))
plan.select(pl.col("x").gather(take_idx)).collect()

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

Почему это стоит запомнить

В «ленивых», ориентированных на стриминг движках поведение каждого выражения имеет значение. Две внешне похожие операции могут иметь радикально разные планы выполнения. Зная, что sample способен затянуть всё в память, тогда как выборка по индексам может этого избежать, вы получаете конкретную стратегию на случай ограничений по ОЗУ.

Итоги

Если выборка большого столбца в Polars LazyFrame приводит к ошибкам из‑за нехватки памяти, откажитесь от sample и переключитесь на выборку по индексам: посчитайте число строк, выберите случайные позиции и соберите эти строки. Такой путь ближе к потоковой модели исполнения и снижает риск материализации полного «тяжёлого» столбца. Проверьте подход в своей среде, особенно на широком столбце, который вызывал сбой, и держите логику выборки в русле операций, которые движок способен эффективно спланировать.