Как логически индексировать фрейм данных pandas с помощью функции

#python #pandas

#python #pandas

Вопрос:

Есть ли способ использовать методы для логической индексации в фрейме данных pandas?

Например:

 import pandas


def filter_func(v) -> bool:
    return v == 'asd'


def main():
    df_test = pandas.DataFrame(
        [
            ['sd'], ['asd'], ['sdf']
        ],
        columns=["col-a"]
    )
    #### ERROR: This next line calls filter_func with all contents of column 'col-a'
    result = df_test[df_test['col-a'] == filter_func(df_test['col-a'])]


if __name__ == '__main__':
    main()
  

В приведенном выше примере я хочу сохранить только те значения, для которых filter_func будут возвращены True . И a result должен содержать фрейм данных с одной строкой, но вместо этого я получаю пустой фрейм данных.

Я понимаю, что вместо выполнения filter_func для каждой строки он выполняется только один раз.

Есть ли способ вызвать его для каждой строки?

Должен ли я использовать apply или map для Series в этом случае?

Или есть какой-либо другой способ?

Ответ №1:

 import pandas as pd
import numpy as np
import random

# sample data
np.random.seed(365)
random.seed(365)
rows = 1100
data = {'a': np.random.randint(10, size=(rows)),
        'groups': [random.choice(['1-5', '6-25', '26-100', '100-500', '500-1000', '>1000']) for _ in range(rows)],
        'treatment': [random.choice(['Yes', 'No']) for _ in range(rows)],
        'date': pd.bdate_range(datetime.today(), freq='h', periods=rows).tolist()}
df = pd.DataFrame(data)

   a  groups treatment                date
0  2   >1000       Yes 2020-10-06 00:00:00
1  4  26-100        No 2020-10-06 01:00:00
2  1   >1000       Yes 2020-10-06 02:00:00
3  5    6-25       Yes 2020-10-06 03:00:00
4  2  26-100        No 2020-10-06 04:00:00

# filter function
def filter_func(v) -> bool:
    return v == '26-100'


# call function
filtered = df[filter_func(df.groups)]

# display(filtered)
    a  groups treatment                date
1   4  26-100        No 2020-10-06 01:00:00
4   2  26-100        No 2020-10-06 04:00:00
21  2  26-100       Yes 2020-10-06 21:00:00
24  9  26-100       Yes 2020-10-07 00:00:00
32  5  26-100        No 2020-10-07 08:00:00