Безопасная выборка из списков в Polars: 2 с начала, 2 из середины и 2 с конца

Polars на практике: берём 2 из начала, 2 из середины (случайно) и 2 с конца списка. Объясняем сбой sample и показываем, как избежать ShapeError надёжно.

Выборка из столбцов-списков в Polars обычно проста, пока не требуется совместить фиксированные срезы со случайной подвыборкой и при этом не «сломаться» на коротких списках. В этом руководстве показано, как взять первые два значения, последние два значения и два случайных значения из середины каждого списка, а затем оставить только уникальные элементы. Также объясняется, почему наивный подход даёт сбой и как реализовать устойчивое решение, работающее при любой длине списка.

Воспроизводим пример и сбой

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

import polars as pl
data = pl.DataFrame(
    {
        "key": ["a", "b"],
        "nums": [[1, 2, 3, 4, 5], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]],
    }
)
broken = (
    data.select(
        top=pl.col("nums").list.head(2),
        mid=pl.col("nums")
            .list.set_difference(pl.col("nums").list.head(2))
            .list.set_difference(pl.col("nums").list.tail(2))
            .list.sample(2, seed=1234),
        bottom=pl.col("nums").list.tail(2),
    )
    .select(pick=pl.concat_list(["top", "mid", "bottom"]).list.unique())
)

Эта попытка приводит к исключению, потому что в коротких списках средний срез может содержать меньше двух элементов:

ShapeError: cannot take a larger sample than the total population when `with_replacement=false`

Почему это ломается

Средняя часть вычисляется путём удаления из списка первых двух и последних двух значений. Для списков длиной пять и меньше остаток оказывается меньше двух. Выборка без возвращения при n, превышающем доступную «популяцию», недопустима — именно это и вызывает ShapeError.

Идиоматичный способ выбрать два из середины

Вместо того чтобы просить list.sample выдать фиксированное количество, которое может оказаться недостижимым, перемешайте средний сегмент и возьмите первые два из перемешанного результата. При fraction=1 и shuffle=True средний список перемешивается, но не меняет размер, а head(2) безопасно ограничивает результат двумя или меньшим числом значений в зависимости от доступности.

data.select(
    pl.col("nums")
      .list.head(pl.col("nums").list.len() - 2)
      .list.slice(2)
      .list.sample(fraction=1, shuffle=True)
      .list.head(2)
      .alias("mid_two")
)

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

Собираем итоговую выборку с проверкой длины

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

result = data.with_columns(
    pl.when(pl.col("nums").list.len() > 5)
      .then(
          pl.concat_list(
              pl.col("nums").list.head(2),
              pl.col("nums")
                .list.head(pl.col("nums").list.len() - 2)
                .list.slice(2)
                .list.sample(fraction=1, shuffle=True)
                .list.head(2),
              pl.col("nums").list.tail(2),
          )
      )
      .otherwise(pl.col("nums"))
      .list.unique()
      .alias("pick")
)

Для списков длиннее пяти это даёт два элемента с начала, два случайных из среднего сегмента и два с конца. Для более коротких списков возвращается исходный список без изменений. Применение list.unique гарантирует, что в итоговом списке останутся уникальные значения.

Альтернатива: сделать размер выборки условным

Другой вариант — не перемешивать среднюю часть и условно задавать, сколько элементов выбирать. Если передать в n выражение, то при недостаточной длине выборка вернёт ноль элементов, что делает всю конструкцию безопасной для любых размеров.

alt = data.with_columns(
    pl.when(pl.col("nums").list.len() > 5)
      .then(
          pl.concat_list(
              pl.col("nums").list.head(2),
              pl.col("nums")
                .list.head(pl.col("nums").list.len() - 2)
                .list.slice(2)
                .list.sample(
                    n=pl.when(pl.col("nums").list.len() > 5).then(2).otherwise(0)
                ),
              pl.col("nums").list.tail(2),
          )
      )
      .otherwise(pl.col("nums"))
      .list.unique()
      .alias("pick")
)

Этот вариант полностью избегает перемешивания и использует «ограничитель» размера выборки, чтобы предотвратить ошибки. Он даёт ту же форму результата и соблюдает требование уникальности.

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

Сочетание детерминированных срезов со случайными выборками — частая задача при обработке данных. Связка «перемешать с fraction=1 и обрезать head» надёжна и избавляет от краевых случаев, когда требуемый размер выборки превышает доступную «популяцию». Переключение между преобразованным и исходным списком через when/then делает конвейер предсказуемым и удобным в сопровождении.

Итоги

Делая выборку из столбцов-списков в Polars, избегайте фиксированного n, которое может оказаться больше доступной «популяции». Либо перемешивайте и обрезайте, либо делайте n условным. Разбейте список с помощью list.head, list.tail и list.slice, используйте list.sample с fraction=1, shuffle=True, чтобы безопасно перемешать сегмент, и переключайтесь между исходным и собранным результатом через when/then. Завершайте list.unique, когда нужно удалить дубликаты в итоговой выборке.