Как вычислить столбец в фрейме данных Polars, используя np.linspace

Как упоминалось в комментариях, добавление функции в стиле np.linspace в поляры — это открытый запрос на добавление функции. Пока это не реализовано, простая реализация с использованием собственного API выражений поляров может выглядеть следующим образом.

Во-первых, мы используем pl.int_range (спасибо @Dean MacGregor), чтобы создать диапазон целых чисел от 0 до num (эксклюзивный). Затем мы масштабируем и сдвигаем диапазон в соответствии с start, stop и num. Наконец, мы включаем столбец с помощью pl.Expr.implode, чтобы получить столбец со списком диапазона для каждой строки.

def pl_linspace(start: pl.Expr, stop: pl.Expr, num: pl.Expr) -> pl.Expr:
    grid = pl.int_range(num)
    _scale = (stop - start) / (num - 1)
    _offset = start
    return (grid * _scale + _offset).implode().over(pl.int_range(pl.len()))

df.with_columns(
    pl_linspace(
        start=pl.col("np_linspace_start"),
        stop=pl.col("np_linspace_stop"),
        num=pl.col("np_linspace_num"),
    ).alias("pl_linspace")
)

shape: (3, 4)
┌───────────────────┬──────────────────┬─────────────────┬────────────────────────────────┐
│ np_linspace_start ┆ np_linspace_stop ┆ np_linspace_num ┆ pl_linspace                    │
│ ---               ┆ ---              ┆ ---             ┆ ---                            │
│ i64               ┆ i64              ┆ i64             ┆ list[f64]                      │
╞═══════════════════╪══════════════════╪═════════════════╪════════════════════════════════╡
│ 0                 ┆ 8                ┆ 5               ┆ [0.0, 2.0, 4.0, 6.0, 8.0]      │
│ 0                 ┆ 6                ┆ 4               ┆ [0.0, 2.0, 4.0, 6.0]           │
│ 0                 ┆ 7                ┆ 4               ┆ [0.0, 2.333333, 4.666667, 7.0] │
└───────────────────┴──────────────────┴─────────────────┴────────────────────────────────┘

Примечание. Если num равно 1, то деление при вычислении _scale приведет к бесконечным значениям. Этого можно избежать, добавив следующее к pl_linspace.

_scale = pl.when(_scale.is_infinite()).then(pl.lit(0)).otherwise(_scale)

Мне очень нравится ответ @Hericks, но он поднимает проблему, которая возникает время от времени, когда нам нужно обойти невозможность заставить списки выполнять операции с другими списками, explodeing, а затем implodeing с овером. Это over действительно дорого, и если мы привлечем нашего старого приятеля pyarrow, чтобы помочь нам, мы сможем этого избежать.

Вот альтернатива pl_linspace, использующая pyarrow.

import pyarrow as pa
def pals_linspace(
    start: pl.Expr | str, stop: pl.Expr | str, num: pl.Expr | str
) -> pl.Expr:
    if isinstance(start, str):
        start = pl.col(start)
    if isinstance(stop, str):
        stop = pl.col(stop)
    if isinstance(num, str):
        num = pl.col(num)
        
    grid = pl.int_ranges(num)
    # Note the repeat_by here to broadcast the scalers to the same shape as grid
    scale = ((stop - start) / (num - 1)).repeat_by(grid.list.len())
    offset = start.repeat_by(grid.list.len())
    
    # Need to make a struct so all the columns we need are in a Series
    all_struct = pl.struct(
        grid.alias("grid"), scale.alias("scale"), offset.alias("offset")
    )
    return all_struct.map_batches(
        lambda s: (
            pl.from_arrow(
                # This is the key to avoiding over, its first argument is the offsets
                # of the list we want to end up with. Since we want to end up with the 
                # same shape as `grid` we can just use that as-is
                
                # The second argument is just @Herick's formula
                pa.LargeListArray.from_arrays(
                    s.struct.field("grid").to_arrow().offsets,
                    (
                        s.struct.field("grid").explode()
                        * s.struct.field("scale").explode()
                        + s.struct.field("offset").explode()
                    ).to_arrow(),
                )
            )
        )
    )

Производительность

Если я сделаю %%timeit из pl_linspace и pals_linspace с заданным df, то pl_linspace будет быстрее. Однако если мы сделаем

df = pl.DataFrame(
    data = {
        "np_linspace_start": [0, 0, 0] * 100,
        "np_linspace_stop": [8, 6, 7] * 100,
        "np_linspace_num": [5, 4, 4] * 100,
    }
)

затем

%%timeit 
df.with_columns(
    ls=pals_linspace("np_linspace_start", "np_linspace_stop", "np_linspace_num")
)
2.95 ms ± 259 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit
df.with_columns(
    pl_linspace(
        start=pl.col("np_linspace_start"),
        stop=pl.col("np_linspace_stop"),
        num=pl.col("np_linspace_num"),
    ).alias("pl_linspace")
)
9.9 ms ± 722 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Я думаю, что оба ответа хороши, но когда я читал ответ @Dean MacGregor, мне стало интересно, сможем ли мы сделать это быстрее.

На данный момент я не думаю, что мы можем сделать чистое полярное решение, которое будет быстрее, но я обнаружил, что обычно комбинация Polars + DuckDB работает очень хорошо.

Итак, вот duckdb решение, которое использует функции generate_series() и list_transform():

duckdb.sql("""
    select
        np_linspace_start,
        np_linspace_stop,
        np_linspace_num,
        list_transform(
            generate_series(0, np_linspace_num - 1),
            x -> x * np_linspace_stop / (np_linspace_num - 1)
        ) as pl_linspace
    from df
""")

┌───────────────────┬──────────────────┬─────────────────┬────────────────────────────────┐
│ np_linspace_start ┆ np_linspace_stop ┆ np_linspace_num ┆ pl_linspace                    │
│ ---               ┆ ---              ┆ ---             ┆ ---                            │
│ i64               ┆ i64              ┆ i64             ┆ list[f64]                      │
╞═══════════════════╪══════════════════╪═════════════════╪════════════════════════════════╡
│ 0                 ┆ 8                ┆ 5               ┆ [0.0, 2.0, 4.0, 6.0, 8.0]      │
│ 0                 ┆ 6                ┆ 4               ┆ [0.0, 2.0, 4.0, 6.0]           │
│ 0                 ┆ 7                ┆ 4               ┆ [0.0, 2.333333, 4.666667, 7.0] │
└───────────────────┴──────────────────┴─────────────────┴────────────────────────────────┘

В моих тестах это работало примерно в 5 раз быстрее, чем полярное решение с explode(), и примерно в 2 раза быстрее, чем pyarrow решение:

%%timeit

duckdb.sql("""
    select
        np_linspace_start,
        np_linspace_stop,
        np_linspace_num,
        list_transform(
            generate_series(0, np_linspace_num - 1),
            x -> x * np_linspace_stop / (np_linspace_num - 1)
        ) as pl_linspace
    from df
""")

472 µs ± 47 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

%%timeit

df.with_columns(
    pl_linspace(
        start=pl.col("np_linspace_start"),
        stop=pl.col("np_linspace_stop"),
        num=pl.col("np_linspace_num"),
    ).alias("pl_linspace")
)

2.1 ms ± 341 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit

df.with_columns(
    ls=pals_linspace("np_linspace_start", "np_linspace_stop", "np_linspace_num")
)

826 µs ± 69.9 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

обновление Как отметил Дин МакГрегор в комментариях, код duckdb не включает метод .pl(), который бы преобразовывал результаты обратно в поляры. Я не делал этого, потому что, по моему опыту, преобразование в поляры не масштабируется линейно с размером кадра данных. Наш тестовый фрейм данных состоит всего из 300 строк, поэтому преобразование в поляры может занять значительную часть времени выполнения.

Я попробовал все три метода на более крупном фрейме данных (3M строк), это, конечно, более гибко, потому что я не запускал его в цикле. Похоже, метод pyarrow побеждает, а duckdb занимает второе место. В 30M строках мне не хотелось ждать достаточно долго, пока решение explode() завершится, но оба решения pyarrow и duckdb сработали в течение 20 секунд.

%%time

a = df.with_columns(
    pl_linspace(
        start=pl.col("np_linspace_start"),
        stop=pl.col("np_linspace_stop"),
        num=pl.col("np_linspace_num"),
    ).alias("pl_linspace")
)

CPU times: total: 6.36 s
Wall time: 18 s

%%time

a = duckdb.sql("""
    select
        np_linspace_start,
        np_linspace_stop,
        np_linspace_num,
        list_transform(generate_series(0, np_linspace_num - 1), x -> x * np_linspace_stop / (np_linspace_num - 1)) as pl_linspace
    from df
""").pl()

CPU times: total: 2.08 s
Wall time: 1.03 s

%%time

a = df.with_columns(
    ls=pals_linspace("np_linspace_start", "np_linspace_stop", "np_linspace_num")
)

CPU times: total: 781 ms
Wall time: 1.29 s

Итоговые списки должны выглядеть примерно так: ._ranges(start, pl.col.stop + 1, (stop - start) / (num - 1). Проблема, с которой мы сейчас столкнулись, заключается в том, что (stop - start) / (num - 1) не является целым числом, поэтому мы не можем использовать int_ranges().

Вопрос в том, можем ли мы как-то модифицировать решения, которые мы могли бы создать int_ranges(), а затем преобразовать их в результирующие float_ranges()?

Мы не можем умножать элементы списка на другой столбец (пока). Но что мы можем сделать, так это умножить или разделить элементы списка по переменной.

А что, если бы у нас было такое число p, что выражение p * (stop - start) / (num - 1) для всех строк приводило бы к целому числу? Тогда мы могли бы создавать списки int_ranges(), например списки .int_ranges(p * start, p * pl.col.stop + 1, p * (stop - start) / (num - 1), а затем просто делить все элементы на p.

Для этого нам нужно найти наименьшее общее кратное всех (num - 1) чисел и разделить его на наибольший общий делитель числа (stop - start). Я не думаю, что функция Polars для этого существует, но мы можем использовать numpy.lcm() и numpy.gcd().

Тогда решение будет выглядеть так:

n1 = list(df.select(pl.col.np_linspace_num - 1).unique().to_series())
n2 = list(df.select(pl.col.np_linspace_stop - pl.col.np_linspace_start).unique().to_series())
p = np.lcm.reduce(n1) / np.gcd.reduce(n2)

df.with_columns(
    pl
    .int_ranges(
        p * pl.col.np_linspace_start,
        p * (pl.col.np_linspace_stop + 1),
        p * (pl.col.np_linspace_stop - pl.col.np_linspace_start) / (pl.col.np_linspace_num - 1)
    )
    .list.eval(pl.element() / p)
    .alias("ls")
)

┌───────────────────┬──────────────────┬─────────────────┬────────────────────────────────┐
│ np_linspace_start ┆ np_linspace_stop ┆ np_linspace_num ┆ ls                             │
│ ---               ┆ ---              ┆ ---             ┆ ---                            │
│ i64               ┆ i64              ┆ i64             ┆ list[f64]                      │
╞═══════════════════╪══════════════════╪═════════════════╪════════════════════════════════╡
│ 0                 ┆ 8                ┆ 5               ┆ [0.0, 2.0, 4.0, 6.0, 8.0]      │
│ 0                 ┆ 6                ┆ 4               ┆ [0.0, 2.0, 4.0, 6.0]           │
│ 0                 ┆ 7                ┆ 4               ┆ [0.0, 2.333333, 4.666667, 7.0] │
└───────────────────┴──────────────────┴─────────────────┴────────────────────────────────┘

Это также кажется довольно быстрым в 3M строках:

%%time

n1 = list(df.select(pl.col.np_linspace_num - 1).unique().to_series())
n2 = list(df.select(pl.col.np_linspace_stop - pl.col.np_linspace_start).unique().to_series())
p = np.lcm.reduce(n1) / np.gcd.reduce(n2)

a = df.with_columns(
    pl
    .int_ranges(
        p * pl.col.np_linspace_start,
        p * (pl.col.np_linspace_stop + 1),
        p * (pl.col.np_linspace_stop - pl.col.np_linspace_start) / (pl.col.np_linspace_num - 1)
    )
    .list.eval(pl.element() / p)
    .alias("ls")
)

CPU times: total: 750 ms
Wall time: 908 ms