Pandas – szybsza alternatywa dla apply

TLDR: np.vectorize() jest nawet kilkaset razy szybszy od pd.apply()

Czasami podczas pracy z pandasowym DataFrame’m zachodzi potrzeba wykonania bardziej nietrywialnej operacji na zbiorze. Załóżmy, że analizujemy ceny nieruchomości w Kalifornii, korzystając z popularnego zbioru California Housing Prices.

import pandas as pd
import numpy as np

california_housing = pd.read_csv("./sample_data/california_housing_train.csv")

Nasz przełożony podpowiada nam, że cena domów, w których sypialnie zajmują dużą część budynku powinny być tańsze niż wskazuje na to mediana wartości budynku. W związku z tym, chcemy zmniejszyć cenę dla takich budynków, natomiast zwiększyć dla budynków, w których taka sytuacja nie występuje.

%%timeit
california_housing['median_house_value_scaled_1'] = \
    california_housing.apply(lambda x: x["median_house_value"]\
                             if x["total_rooms"]> 3. * x["total_bedrooms"]\
                             else 1.5 * x["median_house_value"], axis = 1)

Oczywiste wydaje się być użycie apply(). Tak napisany kod zajmuje Google Colabowi 594 ms. Nie wydaje się być to tragedią. Dataset ma jednak jedynie 17 000 rekordów. Załóżmy, że nasz dataset jest jednak 100 razy większy:

california_housing_long = pd.concat([california_housing]*100)

Wówczas wykonanie tej samej operacji zajmuje już Colabowi ponad minutę. Zakładając, że mamy takich operacji do wykonania kilka (albo, nie daj Boże, musimy to robić regularnie) jest to zauważalna strata czasu.

Z pomocą przychodzi nam np.vectorize(). Dzięki jego zastosowaniu, możemy zachować przejrzysty kod, który będzie liczył się o wiele szybciej – kod zapewniający tą samą funkcjonalność, dla zbioru o 1 700 000 rekordów liczy się 481 ms – ponad 130 razy szybciej niż analogiczny kod z wykorzystaniem apply()

def house_value_scale(median_house_value, total_rooms, total_bedrooms):
  if total_rooms > 5. * total_bedrooms:
    return median_house_value
  else:
    return median_house_value * 1.5

%%timeit
california_housing_long['median_house_value_scaled_2'] = np.vectorize(house_value_scale)\
    (california_housing_long['median_house_value'], 
     california_housing_long['total_rooms'], 
     california_housing_long['median_house_value'])

Korzystanie z niej wygląda nieco inaczej niż w przypadku apply(). Najlepiej najpierw zdefiniować funkcję operującą na pojedynczym rzędzie (bądź kolumnie).

Vectorize tworzy z niej funkcję operującą już na „kolumnach” – jako argumenty do utworzonej funkcji (druga para nawiasów) podajemy „kolumny”.