Pandas: crie duas novas colunas em um dataframe com valores calculados a partir de uma coluna pré-existente

Estou trabalhando com a biblioteca pandas e quero adicionar duas novas colunas a um dataframe dfcom n colunas (n> 0).
Essas novas colunas resultam da aplicação de uma função a uma das colunas no dataframe.

A função a ser aplicada é como:

def calculate(x):
    ...operate...
    return z, y

Um método para criar uma nova coluna para uma função que retorna apenas um valor é:

df['new_col']) = df['column_A'].map(a_function)

Então, o que eu quero, e tentei sem sucesso (*), é algo como:

(df['new_col_zetas'], df['new_col_ys']) = df['column_A'].map(calculate)

Qual poderia ser a melhor maneira de fazer isso? Eu examinei a documentação sem nenhuma pista.

** df['column_A'].map(calculate)retorna uma série pandas em que cada item consiste em uma tupla z, y. E tentar atribuir isso a duas colunas de dataframe produz um ValueError. *

Eu apenas usaria zip:

In [1]: from pandas import *

In [2]: def calculate(x):
   ...:     return x*2, x*3
   ...: 

In [3]: df = DataFrame({'a': [1,2,3], 'b': [2,3,4]})

In [4]: df
Out[4]: 
   a  b
0  1  2
1  2  3
2  3  4

In [5]: df["A1"], df["A2"] = zip(*df["a"].map(calculate))

In [6]: df
Out[6]: 
   a  b  A1  A2
0  1  2   2   3
1  2  3   4   6
2  3  4   6   9

A principal resposta é falha em minha opinião. Felizmente, ninguém está importando em massa todos os pandas para seu namespace com from pandas import *. Além disso, o mapmétodo deve ser reservado para aqueles momentos em que é passado um dicionário ou série. Pode ter uma função, mas é para isso que applyé usado.

Então, se você deve usar a abordagem acima, eu escreveria assim

df["A1"], df["A2"] = zip(*df["a"].apply(calculate))

Na verdade, não há razão para usar o zip aqui. Você pode simplesmente fazer isso:

df["A1"], df["A2"] = calculate(df['a'])

Este segundo método também é muito mais rápido em DataFrames maiores

df = pd.DataFrame({'a': [1,2,3] * 100000, 'b': [2,3,4] * 100000})

DataFrame criado com 300.000 linhas

%timeit df["A1"], df["A2"] = calculate(df['a'])
2.65 ms ± 92.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit df["A1"], df["A2"] = zip(*df["a"].apply(calculate))
159 ms ± 5.24 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

60x mais rápido do que zip

Em geral, evite usar aplicar

Aplicar geralmente não é muito mais rápido do que iterar em uma lista Python. Vamos testar o desempenho de um loop for para fazer a mesma coisa que acima

%%timeit
A1, A2 = [], []
for val in df['a']:
    A1.append(val**2)
    A2.append(val**3)

df['A1'] = A1
df['A2'] = A2

298 ms ± 7.14 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Portanto, é duas vezes mais lento, o que não é uma regressão de desempenho terrível, mas se citonizarmos o acima, obteremos um desempenho muito melhor. Supondo que você esteja usando ipython:

%load_ext cython

%%cython
cpdef power(vals):
    A1, A2 = [], []
    cdef double val
    for val in vals:
        A1.append(val**2)
        A2.append(val**3)

    return A1, A2

%timeit df['A1'], df['A2'] = power(df['a'])
72.7 ms ± 2.16 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Atribuir diretamente sem aplicar

Você pode obter melhorias de velocidade ainda maiores se usar as operações vetorizadas diretas.

%timeit df['A1'], df['A2'] = df['a'] ** 2, df['a'] ** 3
5.13 ms ± 320 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Isso aproveita as operações vetorizadas extremamente rápidas do NumPy em vez de nossos loops. Agora temos um aumento de 30x em relação ao original.

O teste de velocidade mais simples com apply

O exemplo acima deve mostrar claramente o quão lento applypode ser, mas só para ficar mais claro, vamos dar uma olhada no exemplo mais básico. Vamos elevar ao quadrado uma série de 10 milhões de números com e sem aplicação

s = pd.Series(np.random.rand(10000000))

%timeit s.apply(calc)
3.3 s ± 57.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Sem aplicar é 50x mais rápido

%timeit s ** 2
66 ms ± 2 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)