Как ускорить apply(). Часть 2

В примере выше использовалась простая арифметическая функция, но можно ли ускорить apply с произвольной сложной функцией? Ответ — да, можно! Но с помощью чуть более сложной техники — numba.

Numba — это JIT (Just-In-Time) компилятор для Python, который работает по следующей схеме: в первый раз видит функцию, анализирует типы и компилирует в машинный код. А потом просто запускает скомпилированный машинный код, и скорость выполнения близка близка к C/C++.

import pandas as pd
import numpy as np
import time
from numba import jit

# Создаём датасет
np.random.seed(42)
df = pd.DataFrame({
    'col1': np.random.rand(1_000_000),
    'col2': np.random.rand(1_000_000),
    'col3': np.random.rand(1_000_000),
})

def complex_logic(row):
    result = 0
    # Имитируем "сложную логику" с несколькими условиями
    if row['col1'] > 0.5:
        result += row['col2'] * row['col3'] * 2
    else:
        result += row['col2'] * row['col3'] * 0.5
    
    if row['col2'] > row['col3']:
        result += row['col1'] * 10
    else:
        result += row['col1'] * 5
    return result

# Медленный apply()
start = time.time()
df['result_slow'] = df.apply(complex_logic, axis=1)
slow_time = time.time() - start

print(f"Время выполнения: {slow_time:.2f} секунд")

# Быстрый способ через numba
@jit(nopython=True) # numba пытается скомпилировать в машинный код
def complex_logic_fast(col1, col2, col3): # Переписываем функцию из pandas на numpy
    result = np.empty(len(col1))
    
    for i in range(len(col1)):
        res = 0
        if col1[i] > 0.5:
            res += col2[i] * col3[i] * 2
        else:
            res += col2[i] * col3[i] * 0.5

        if col2[i] > col3[i]:
            res += col1[i] * 10
        else:
            res += col1[i] * 5
        result[i] = res
    return result

# Первый вызов: компиляция
df['result_fast'] = complex_logic_fast(df['col1'].values, df['col2'].values, df['col3'].values)

# Второй вызов: уже скомпилировано (очень быстро)
start = time.time()
result_fast_2 = complex_logic_fast(df['col1'].values, df['col2'].values, df['col3'].values)
fast_time = time.time() - start
print(f"Время выполнения (уже скомпилировано): {fast_time:.4f} секунд")

Итого, numba = Python + компиляция в машинный код = скорость C с синтаксисом Python.

Ставьте ❤️, если было полезно!

📊 Simulative