Как создать последовательность чисел с плавающей точкой в Polars типа List[f64]

На чтение 3 мин Опубликовано 11.10.2024

Содержание

Вопрос или проблема
Ответ или решение
Пояснение:

Вопрос или проблема

У меня есть список полярных значений List[f64], столбец “a”. Я хочу создать новый список полярных значений List[f64], столбец “b”, который будет последовательностью от минимума до максимума списка в строке столбца a с шагом 0.5, включая оба конца. Так, для строки со списком в столбце “a” [0.0, 3.0, 2.0, 6.0, 2.0], значение в столбце b должно быть [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0].

Это мое решение, но в нем есть ошибка.

df = df.with_columns(
        pl.col("a").list.eval(
            pl.arange(pl.element().min(), pl.element().max(), 1)
            .append(pl.arange(pl.element().min(), pl.element().max(), 1) + 0.5)
            .append(pl.element().max())
            .append(pl.element().max() - 0.5)
            .unique()
            .sort(),
            parallel=True,
        )
        .alias("b")
    )

Это решение не работает для крайних случаев, когда столбец a содержит только 1 уникальное значение в своем списке. Поскольку у полярных значений, похоже, есть только целочисленная функция arange(), когда я создаю второй список и прибавляю 0.5, если имеется только одно уникальное значение, это приводит к появлению 2 значений в выходных данных: самого значения и самого значения – 0.5.

Вот немного тестовых данных. Столбец “a” содержит списки, минимальные и максимальные значения которых должны использоваться для определения границ последовательности, которую представляет столбец “b”.

pl.DataFrame([
    pl.Series('a', [[4.0, 5.0, 3.0, 7.0, 0.0, 1.0, 6.0, 2.0], [2.0, 4.0, 3.0, 0.0, 1.0], [1.0, 2.0, 3.0, 0.0, 4.0], [1.0, 3.0, 2.0, 0.0], [1.0, 0.0]], dtype=pl.List(pl.Float64)),        
    pl.Series('b', [[0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5, 7.0], [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0], [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0], [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0], [0.0, 0.5, 1.0]], dtype=pl.List(pl.Float64))
])

Скорость здесь довольно важна, я переписываю на Polars с этой целью. Спасибо.

Это относительно просто, создавая включительный целочисленный диапазон от 2*min до 2*max и деля его на 2:

df.with_columns(b = pl.col.a.list.eval(
    pl.arange(2*pl.element().min(), 2*pl.element().max() + 1) / 2
))

Ответ или решение

Для создания нового столбца b, который будет представлять собой последовательность с шагом 0.5 от минимального до максимального значения в списках из столбца a, можно воспользоваться библиотекой Polars. Ваша задача заключается в том, чтобы учесть различные сценарии, включая крайние случаи, такие как списки, содержащие только одно уникальное значение.

Вот как можно это реализовать:

Используем метод list.eval для обработки каждого списка в столбце a.
Для получения последовательности от минимального до максимального значения, можно создать целочисленный диапазон, умножив минимальное и максимальное значение на 2, чтобы избежать проблем с типами при использовании метода arange().
Затем, делим все значения на 2, чтобы получить нужные значения с шагом 0.5.
Обратите внимание, как это решение обрабатывает случай, когда в списке всего одно уникальное значение.

Вот окончательный код для создания столбца b:

import polars as pl

# Создание исходного DataFrame
df = pl.DataFrame([
    pl.Series('a', [[4.0, 5.0, 3.0, 7.0, 0.0, 1.0, 6.0, 2.0], 
                     [2.0, 4.0, 3.0, 0.0, 1.0], 
                     [1.0, 2.0, 3.0, 0.0, 4.0], 
                     [1.0, 3.0, 2.0, 0.0], 
                     [1.0, 0.0]], dtype=pl.List(pl.Float64)))
])

# Создание нового столбца b
df = df.with_columns(
    b=pl.col("a").list.eval(
        # Создаем диапазон от 2 * min до 2 * max + 1, затем делим на 2
        pl.arange(2 * pl.element().min(), 2 * pl.element().max() + 1) / 2
    )
)

print(df)

Пояснение:

pl.element().min() и pl.element().max() находят минимальное и максимальное значения в каждом списке.
pl.arange(2 * min, 2 * max + 1) / 2 создает последовательность чисел с шагом 0.5 от минимального до максимального значения. Мы умножаем на 2, чтобы использовать arange(), который работает с целыми числами, а затем делим результат на 2, чтобы вернуться к плавающим числам.
Так как мы используем + 1 в arange(), мы включаем максимальное значение в результирующий диапазон.

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