blog/content/ru/python-speedup-with-rust.md

title, description, position, category

title	description	position	category
Ускоряем Python используя Rust.	Как встроить Rust в проект на Python.	2	Python

Описание проблемы

Каким бы Python удобным не был, скоростью он похвастаться никак не может. И для некоторых задач это достаточно критично.

Например, не так давно у меня была задача достать много информации из файлов логов. Проблема в том, что один лог-файл занимает от 3-4 ГБ. А файлов таких много и информацию достать требуется быстро. Изначальный вариант на Python был написан за минут 15-20, но скорость его работы занимала 30-40 минут на один файл. После переписывания одной функции на Rust, скрипт отработал за 1 минуту.

Давайте напишем свой проект со встроенной функцией на Rust.

Задача проекта будет следующей:

Дан файл лога запросов на некоторый сервер. Требуется найти количество запросов к определённому файлу и сумму переданных байт.

Для демонстрации мы напишем 2 функции, одна будет использовать наивное решение задачи на питоне, вторая будет выполнять парсинг на расте.

Сам файл лог имеет следующую структуру:

"$line_num" "-" "$method $url" "$current_time" "$bytes_sent"

где,

• $line_num - номер строчки;
• $method - HTTP метод;
• $url - URL до файла, содержащий один из сгенерированных идентификаторов;
• $current_time - текущее время (время выполнения запроса);
• $bytes_sent - сколько байт было отправлено.

Данный формат логов отчасти подражает формату G-Core labs.

Проект на Python

Сначала мы напишем весь функционал на Python. Для этого создадим проект python-rust по гайду из блога.

Для создания CLI я буду использовать typer.

Весь код проекта доступен в репозитории.

$ poetry add typer 

Генератор логов

Для тестов напишем генератор лог-файла.

Данный генератор должен работать следующим образом. Я указываю сколько строк лог-файла я хочу увидеть, сколько уникальных id файлов использовать и название файла, куда писать.

Функция же генерирует лог в заданном формате.

import secrets
import time
import uuid
from pathlib import Path
from typing import Any

import typer

tpr = typer.Typer()


def quote(somthing: Any) -> str:
    """
    Quote string.

    :param somthing: any string.
    :return: quoted string.
    """
    return f'"{somthing}"'


@tpr.command()
def generator(  # noqa: WPS210
    output: Path,
    lines: int = 2_000_000,  # noqa: WPS303
    ids: int = 1000,
) -> None:
    """
    Test log generator.

    :param ids: how many file id's to generate.
    :param output: output file path.
    :param lines: how many lines to write, defaults to 2_000_000
    """
    ids_pool = [uuid.uuid4().hex for _ in range(ids)]

    with open(output, "w") as out_file:
        for line_num in range(lines):
            item_id = secrets.choice(ids_pool)
            prefix = secrets.token_hex(60)
            url = f"GET /{prefix}/{item_id}.jpg"
            current_time = int(time.time())
            bytes_sent = secrets.randbelow(800)  # noqa: WPS432
            line = [
                quote(line_num),
                quote("-"),
                quote(url),
                quote(current_time),
                quote(bytes_sent),
            ]
            out_file.write(" ".join(line))
            out_file.write("\n")
    typer.secho("Log successfully generated.", fg=typer.colors.GREEN)


@tpr.command()
def parser(input_file: Path, rust: bool = False) -> None:
    """
    Parse given log file.

    :param input_file: path of input file.
    :param rust: use rust parser implementation.
    """
    typer.secho("Not implemented", err=True, fg=typer.colors.RED)


def main() -> None:
    """Main program entrypoint."""
    tpr()

Для удобства работы добавим в pyproject.toml информацию о командах.

[tool.poetry.scripts]
pyrust = "python_rust.main:main"

Теперь мы можем вызывать нашу программу.

$ pyrust --help
Usage: pyrust [OPTIONS] COMMAND [ARGS]...

Options:
  --install-completion [bash|zsh|fish|powershell|pwsh]
                                  Install completion for the specified shell.
  --show-completion [bash|zsh|fish|powershell|pwsh]
                                  Show completion for the specified shell, to
                                  copy it or customize the installation.
  --help                          Show this message and exit.

Commands:
  generator  Test log generator.
  parser     Parse given log file.

Работает отлично.

С помощью данного генератора сгенерируем файл на 8 милионов строк.

$ pyrust generator test.log --lines 8000000
Log successfully generated.

У нас получился файл на 1.5G. Для начального теста этого будет достаточно.

Реализация парсера на Python

Представленный формат разделяет ифнормацию кавычками, поэтому для сплита строк мы будем использовать встроенный в python модуль shlex.

Функция парсинга логов будет выглядить следующим образом:

import shlex
from pathlib import Path
from typing import Dict, Tuple


def parse_python(  # noqa: WPS210
    filename: Path,
) -> Dict[str, Tuple[int, int]]:
    """
    Parse log file with python.

    :param filename: log file.
    :return: parsed data.
    """
    parsed_data = {}
    with open(filename, "r") as input_file:
        for line in input_file:
            spl = shlex.split(line)
            # Splitting method and actual url.
            url = spl[2].split()[1]
            # Splitting url by /
            # This split will turn this "/test/aaa.png"
            # into "aaa".
            file_id = url.split("/")[-1].split(".")[0]
            file_info = parsed_data.get(file_id)
            # If information about file isn't found.
            if file_info is None:
                downloads, bytes_sent = 0, 0
            else:
                downloads, bytes_sent = file_info
            # Incrementing counters.
            downloads += 1
            bytes_sent += int(spl[4])
            # Saving back.
            parsed_data[file_id] = (downloads, bytes_sent)
    return parsed_data

Давайте импортируем её, модифицируем команду парсинга и замерим скорость работы.

@tpr.command()
def parser(input_file: Path, rust: bool = False) -> None:
    """
    Parse given log file.

    :param input_file: path of input file.
    :param rust: use rust parser implementation.
    """
    if rust:
        typer.secho("Not implemented", input_file, color=typer.colors.RED)
        return
    else:
        parsed_data = parse_python(input_file)

    typer.secho(
        f"Found {len(parsed_data)} files",  # noqa: WPS237
        color=typer.colors.CYAN,
    )

Как можно видеть из кода, мы просто подсчитываем итоговое количество найденных файлов.

Посмотрим сколько времени займёт парсинг нашего сгенерированного файла.

$ time pyrust parser "test.log"
Found 1000 files
pyrust parser test.log  2443.42s user 2.10s system 99% cpu 40:59.30 total

Обработка данного файла заняла 40 минут. Это довольно печальный результат. Самое время попробовать использовать Rust.

Интеграция Rust

Для интеграции Rust в python код я буду использовать PyO3 для биндингов и maturin для сборки проекта.

Начнём с инициализации проекта и установки сборщика. В корневой папке проекта создайте папку проекта с растом.

В теории можно использовать maturin как основной инструмент сборки проекта, но это лишает вас всех прелестей poetry. Поэтому удобнее использовать Rust в подпроекте и указать его как зависимость в списке зависимостей своего проекта.

$ cargo new --lib rusty_log_parser

Теперь создадим pyproject.toml и напишем туда описание нашего python пакета.

[tool.poetry]
name = "rusty_log_parser"
version = "0.1.0"
description = "Log file parser with Rust core"
authors = ["Pavel Kirilin <win10@list.ru>"]

[build-system]
requires = ["maturin>=0.12,<0.13"]
build-backend = "maturin"

Также поправим Cargo.toml.

[lib]
# Название модуля
name = "rusty_log_parser"
# Обязательный тип крейта, чтобы можно было использовать его
# в питоне.
crate-type = ["cdylib"]

[dependencies]
# Библиотека биндингов для питона.
pyo3 = { version = "0.15.1", features = ["extension-module"] }
# Библиотека для сплита. Повторяет функционал shlex.
shell-words = "1.0.0"

После этих нехитрых изменений попробуем собрать проект.

$ cd rusty_log_parser
$ cargo build

Теперь напишем сам парсер логов в Rust.

use std::collections::HashMap;
use std::fs::File;
use std::io;
use std::io::BufRead;

use pyo3::prelude::*;

/// Parse log file in rust.
/// This function parses log file and returns HashMap with Strings as keys and
/// Tuple of two u128 as values.
#[pyfunction]
fn parse_rust(filename: &str) -> PyResult<HashMap<String, (u128, u128)>> {
    let mut result_map = HashMap::new();
    // Iterating over file.
    for log in io::BufReader::new(File::open(filename)?).lines().flatten() {
        // Splitting log string.
        if let Ok(mut spl) = shell_words::split(&log) {
            // Getting file id.
            let file_id_opt = spl.get_mut(2).and_then(|http_req| {
                // Splitting method and URL.
                http_req.split(' ').into_iter().nth(1).and_then(|url| {
                    // Splitting by / and getting the last split.
                    url.split('/')
                        .into_iter()
                        .last()
                        // Split file id and extension.
                        .and_then(|item_url| item_url.split('.').into_iter().next())
                        // Turning &str into String.
                        .map(String::from)
                })
            });
            // Getting number of bytes sent.
            let bytes_sent_opt =
                spl.get_mut(4)
                    // Parsing string to u128
                    .and_then(|bytes_str| match bytes_str.parse::<u128>() {
                        Ok(bytes_sent) => Some(bytes_sent),
                        Err(_) => None,
                    });
            if file_id_opt.is_none() || bytes_sent_opt.is_none() {
                continue;
            }
            let file_id = file_id_opt.unwrap();
            let bytes_sent = bytes_sent_opt.unwrap();

            match result_map.get(&file_id) {
                Some(&(downloads, total_bytes_sent)) => {
                    result_map.insert(file_id, (downloads + 1, total_bytes_sent + bytes_sent));
                }
                None => {
                    result_map.insert(file_id, (1, bytes_sent));
                }
            }
        }
    }
    Ok(result_map)
}

/// A Python module implemented in Rust. The name of this function must match
/// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
/// import the module.
#[pymodule]
fn rusty_log_parser(_py: Python, m: &PyModule) -> PyResult<()> {
    // Adding function to the module.
    m.add_function(wrap_pyfunction!(parse_rust, m)?)?;
    Ok(())
}

Это и будет наша функция, которую мы будем вызывать. Теперь добавим python обёртку для нашего модуля.

Для этого в проекте rusty_log_parser надо создать папку с тем же названием что и проект. В данном случае это будет rusty_log_parser.

Внутри этой папки мы создадим __init__.py, в котором опишем и задокументируем доступные функции модуля.

from pathlib import Path
from typing import Dict, Tuple

from .rusty_log_parser import parse_rust as _parse_rust


def parse_rust(input_file: Path) -> Dict[str, Tuple[int, int]]:
    """
    Parse log file using Rust as a backend.

    :param input_file: log file to parse.
    :return: Parsed
    """
    return _parse_rust(str(input_file.expanduser()))

Осталось только добавить свеженаписанный пакет как зависимости нашего проекта.

Для этого в корневой pyproject.toml надо добавить следующую зависимость.

[tool.poetry.dependencies]
...
rusty_log_parser = { path = "./rusty_log_parser" }

Теперь мы можем использовать нашу функцию.

Перепишем нашу функцию парсера таким образом, чтобы она использовала rust, когда был передан соответствующий параметр.

from rusty_log_parser import parse_rust

...

@tpr.command()
def parser(input_file: Path, rust: bool = False) -> None:
    """
    Parse given log file.

    :param input_file: path of input file.
    :param rust: use rust parser implementation.
    """
    if rust:
        parsed_data = parse_rust(input_file)
    else:
        parsed_data = parse_python(input_file)

    typer.secho(
        f"Found {len(parsed_data)} files",  # noqa: WPS237
        color=typer.colors.CYAN,
    )

Теперь проведём итоговый замер.

$ time pyrust parser test.log --rust
Found 1000 files
pyrust parser test.log --rust  20.44s user 0.35s system 99% cpu 20.867 total

Итого виден небольшой выйигрыш во времени в 2423 секунды, из чего следует, что реализация на Rust быстрее всего в 122 раза.

Данной статьёй я не подталкиваю всех переписывать всё на Rust, а просто ознакамливаю с возможностями. Но если всё же хочется, то можно и попробовать.

До новых встреч.