pandas: filtro complexo em linhas de DataFrame

86

Eu gostaria de filtrar linhas por uma função de cada linha, por exemplo

def f(row):
  return sin(row['velocity'])/np.prod(['masses']) > 5

df = pandas.DataFrame(...)
filtered = df[apply_to_all_rows(df, f)]

Ou para outro exemplo mais complexo e inventado,

def g(row):
  if row['col1'].method1() == 1:
    val = row['col1'].method2() / row['col1'].method3(row['col3'], row['col4'])
  else:
    val = row['col2'].method5(row['col6'])
  return np.sin(val)

df = pandas.DataFrame(...)
filtered = df[apply_to_all_rows(df, g)]

Como posso fazer isso?

python pandas Duckworthd
fonte

Respostas:

122

Você pode fazer isso usando DataFrame.apply, que aplica uma função ao longo de um determinado eixo,

In [3]: df = pandas.DataFrame(np.random.randn(5, 3), columns=['a', 'b', 'c'])

In [4]: df
Out[4]: 
          a         b         c
0 -0.001968 -1.877945 -1.515674
1 -0.540628  0.793913 -0.983315
2 -1.313574  1.946410  0.826350
3  0.015763 -0.267860 -2.228350
4  0.563111  1.195459  0.343168

In [6]: df[df.apply(lambda x: x['b'] > x['c'], axis=1)]
Out[6]: 
          a         b         c
1 -0.540628  0.793913 -0.983315
2 -1.313574  1.946410  0.826350
3  0.015763 -0.267860 -2.228350
4  0.563111  1.195459  0.343168
Duckworthd
fonte
16
Não há necessidade de applynesta situação. Um índice booleano regular funcionará perfeitamente. df[df['b] > df['c']]. Existem muito poucas situações que realmente exigem applye mesmo poucas que precisam comaxis=1
Ted Petrou
@TedPetrou E se você não tiver certeza de que cada elemento em seu dataframe é do tipo certo. Um índice booleano regular oferece suporte ao tratamento de exceções?
D. Ror.
13

Suponha que eu tenha um DataFrame como segue:

In [39]: df
Out[39]: 
      mass1     mass2  velocity
0  1.461711 -0.404452  0.722502
1 -2.169377  1.131037  0.232047
2  0.009450 -0.868753  0.598470
3  0.602463  0.299249  0.474564
4 -0.675339 -0.816702  0.799289

Posso usar sin e DataFrame.prod para criar uma máscara booleana:

In [40]: mask = (np.sin(df.velocity) / df.ix[:, 0:2].prod(axis=1)) > 0

In [41]: mask
Out[41]: 
0    False
1    False
2    False
3     True
4     True

Em seguida, use a máscara para selecionar no DataFrame:

In [42]: df[mask]
Out[42]: 
      mass1     mass2  velocity
3  0.602463  0.299249  0.474564
4 -0.675339 -0.816702  0.799289
Chang She
fonte
2
na verdade, este foi provavelmente um mau exemplo: np.sintransmite automaticamente para todos os elementos. E se eu o substituísse por uma função menos inteligente que só pudesse lidar com uma entrada de cada vez?
duckworthd
5

Não posso comentar a resposta de duckworthd , mas não está funcionando perfeitamente. Ele trava quando o dataframe está vazio:

df = pandas.DataFrame(columns=['a', 'b', 'c'])
df[df.apply(lambda x: x['b'] > x['c'], axis=1)]

Saídas: 

ValueError: Must pass DataFrame with boolean values only

Para mim, parece um bug no pandas, já que {} é definitivamente um conjunto válido de valores booleanos. Para uma solução, consulte a resposta de Roy Hyunjin Han .

cglacet
fonte
3

A melhor abordagem que encontrei é, em vez de usar reduce=Truepara evitar erros para df vazio (já que este argumento está obsoleto de qualquer maneira), apenas verifique se o tamanho do df é> 0 antes de aplicar o filtro:

def my_filter(row):
    if row.columnA == something:
        return True

    return False

if len(df.index) > 0:
    df[df.apply(my_filter, axis=1)]
user553965
fonte
0

Você pode usar o loc propriedade para fatiar seu dataframe.

Conforme documentação , locpode-se ter callable functioncomo argumento.

In [3]: df = pandas.DataFrame(np.random.randn(5, 3), columns=['a', 'b', 'c'])

In [4]: df
Out[4]: 
          a         b         c
0 -0.001968 -1.877945 -1.515674
1 -0.540628  0.793913 -0.983315
2 -1.313574  1.946410  0.826350
3  0.015763 -0.267860 -2.228350
4  0.563111  1.195459  0.343168

# define lambda function
In [5]: myfilter = lambda x: x['b'] > x['c']

# use my lambda in loc
In [6]: df1 = df.loc[fif]

se você deseja combinar sua função de filtro fifcom outros critérios de filtro

df1 = df.loc[fif].loc[(df.b >= 0.5)]
Pierock
fonte