Python - векторный поиск регулярных выражений для классификации

Question

У меня есть функция идентификатора, которая проходит через все элементы в столбце DataFrame и присваивает им категорию. Код, как у меня сейчас, выглядит следующим образом;

def fruit_replace(x): 
    fruit_quantity = re.search(r'(\\d+)quantity', x) 
    if 'apple' in x: 
     return 'green' 
    elif 'pear' in x: 
     return 'green' 
    elif 'cherry' in x: 
     return 'red' 
    elif 'banana' in x: 
     return 'yellow' 
    elif fruit_quantity != None: 
     return fruit_quantity.group(0) 

Я применяю это в лямбда-функции на DataFrame и присваиваю результаты в новом столбце. К сожалению, это немного сложно из-за того, что поиск fruit_quantity отличается от других.

Процесс должен давать что-то подобное;

Оригинал DataFrame

pd.DataFrame({'fruit_type': ['big apple', 'small cherry', 'jerry 10quantity']}) 

В этот

pd.DataFrame({'fruit_type': ['big apple', 'small cherry', 'peach 10quantity'], 
       'category': ['green', 'red', 10]}) 

Мой вопрос, если этот код может быть улучшена в более вещий или панд образом, и, возможно, векторизация? Я должен применить это примерно к 5 миллионам строк, и это занимает некоторое время.

Большое спасибо!

Answer 1

вы можете использовать boolean indexing в сочетании с str.contains() методом:

df['category'] = np.nan 

df.loc[df.fruit_type.str.contains(r'\b(?:apple|pear)\b'), 'category'] = 'green' 
df.loc[df.fruit_type.str.contains(r'\b(?:cherry)\b'), 'category'] = 'red' 
df.loc[df.fruit_type.str.contains(r'\b(?:banana)\b'), 'category'] = 'yellow' 
df.loc[df['category'].isnull() & (df.fruit_type.str.contains(r'\d+q')), 'category'] = \ 
    df.fruit_type.str.extract(r'(\d+)q', expand=False) 

Результат:

In [270]: df 
Out[270]: 
     fruit_type category 
0   big apple green 
1  small cherry  red 
2 jerry 10quantity  10