Лучший способ добавить результат применения (несколько выходов) в существующий DataFrame с именами столбцов

Question

Я применяю функцию к строкам кадра данных в pandas. Эта функция возвращает четыре значения (то есть четыре значения в строке). На практике это означает, что возвращаемый объект из функции apply является Series, содержащим кортежи. Я хочу добавить их в свои столбцы. Я знаю, что могу преобразовать этот вывод в DataFrame, а затем объединить со старым DataFrame, например, так:

import pandas as pd


def some_func(i):
    return i+1, i+2, i+3, i+4

df = pd.DataFrame(range(10), columns=['start'])
res = df.apply(lambda row: some_func(row['start']), axis=1)

# convert to df and add column names
res_df = res.apply(pd.Series)
res_df.columns = ['label_1', 'label_2', 'label_3', 'label_4']

# concatenate with old df
df = pd.concat([df, res_df], axis=1)
print(df)

Мой вопрос: есть ли лучший способ сделать это? Особенно res.apply(pd.Series) кажется избыточным, но я не знаю лучшей альтернативы. Производительность - важный фактор для меня.

---

Как и требовалось, пример входного DataFrame может выглядеть следующим образом

   start
0      0
1      1
2      2
3      3
4      4
5      5
6      6
7      7
8      8
9      9

И ожидаемый результат с четырьмя добавленными столбцами:

   start  label_1  label_2  label_3  label_4
0      0        1        2        3        4
1      1        2        3        4        5
2      2        3        4        5        6
3      3        4        5        6        7
4      4        5        6        7        8
5      5        6        7        8        9
6      6        7        8        9       10
7      7        8        9       10       11
8      8        9       10       11       12
9      9       10       11       12       13

Answer 1

Непосредственное присвоение значений DataFrame будет быстрее, чем объединение.

Это один из способов присвоения

df = pd.DataFrame(range(10), columns=['start'])

df['label_1'], df['label_2'], df['label_3'], df['label_4'] = zip(*[some_func(x) for x in df['start']])

Это быстрее, чем res.apply(pd.Series).

См. добавление нескольких столбцов к pandas одновременно для получения дополнительных способов назначения нескольких столбцов.

Answer 2

Вот пара вещей, которые можно попробовать вместо метода apply дважды: использовать генератор и конструктор pandas.DataFrame, а также метод DataFrame.join:

import pandas as pd

def some_func(i):
    return i+1, i+2, i+3, i+4

df = pd.DataFrame(range(10), columns=['start'])

df = df.join(pd.DataFrame(some_func(x) for x in df['start'])
             .rename(columns=lambda x: f'label_{x+1}'))

[out]

   start  label_1  label_2  label_3  label_4
0      0        1        2        3        4
1      1        2        3        4        5
2      2        3        4        5        6
3      3        4        5        6        7
4      4        5        6        7        8
5      5        6        7        8        9
6      6        7        8        9       10
7      7        8        9       10       11
8      8        9       10       11       12
9      9       10       11       12       13

---

Время

Это должно предложить ~144x ускорение

Использование df = df = pd.DataFrame(range(100000), columns=['start']) для тестирования.

оригинальный метод:

22.2 s ± 246 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

это решение:

152 ms ± 3.31 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

@ Решение Andrea:

1.66 s ± 17.3 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

@ Решение Keval Dave: <- самое быстрое! </strong> ~230 x ускорение

95.2 ms ± 857 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Answer 3

Если вам нужен однострочник, вы можете попробовать:

df[['label_1', 'label_2', 'label_3', 'label_4']] = pd.DataFrame(
    df.apply(lambda row: some_func(row['start']), axis=1).tolist())

Это очень быстрый способ, как предлагается здесь

Answer 4

Если вы вернете pd.Series из своей функции, то Pandas превратит его элементы в столбцы результирующего кадра данных при вызове apply().

Если вы используете индекс для Серии, элементы в индексе станут именами столбцов результирующего DataFrame.

В вашем случае:

res_index = pd.Index([
    'label_1',
    'label_2',
    'label_3',
    'label_4',
])
res_df = df.apply(
    lambda row: pd.Series(some_func(row['start']), index=res_index),
    axis=1,
)
df = pd.concat([df, res_df], axis=1)
print(df)

Возможно, чуть чище для res_df - применить его только к серии :

res_df = df['start'].apply(
    lambda i: pd.Series(some_func(i), index=res_index),
)

Если вы готовы переписать свой some_func, чтобы вернуть pd.Series напрямую:

def some_func(i, index=None):
    return pd.Series(
        [i+1, i+2, i+3, i+4],
        index=index,
    )

res_index = pd.Index([
    'label_1',
    'label_2',
    'label_3',
    'label_4',
])
res_df = df['start'].apply(some_func, index=res_index)
df = pd.concat([df, res_df], axis=1)
print(df)

Все вышеперечисленное вернет ожидаемый результат:

   start  label_1  label_2  label_3  label_4
0      0        1        2        3        4
1      1        2        3        4        5
2      2        3        4        5        6
3      3        4        5        6        7
4      4        5        6        7        8
5      5        6        7        8        9
6      6        7        8        9       10
7      7        8        9       10       11
8      8        9       10       11       12
9      9       10       11       12       13