Zenn
🐍

Pythonのコードだけでパッケージを100倍高速化する

に公開
Python
Codon
tech

はじめに

Pythonはシンプルで読みやすい文法と、豊富な標準ライブラリやサードパーティ製パッケージを持つことから、データ分析、Web開発、機械学習、自動化スクリプトなど幅広い分野で採用されています。また、大規模なコミュニティが存在し、ドキュメントや情報が充実している点も人気の理由です。

サーバーサイド向け言語として定着していますが、最近ではブラウザ上で動作する環境として PyScript なども登場しており、クライアントサイドのスクリプトとしての応用にも期待が高まっています。

ただし、インタプリタ型の言語であるため、大量データの処理などでは実行速度に制約が生じることがあります。

Pythonの高速化について

Pythonが遅いと感じたとき、いくつかのアプローチが考えられます。

PyPy

  • CPythonの代替実装として、JIT（Just-In-Time）コンパイラを備えたPyPyを利用することで、ループや反復処理の速度を大きく向上させることができます。
  • ただし、PyPIにビルド済みホイールが提供されていないパッケージも多く、依存関係の多いプロジェクトでは環境構築に手間がかかる場合があります。

numba

  • NumPyと組み合わせて利用できる @njit デコレーターによる関数単位のJITコンパイルで、数値計算ループをネイティブコード並みに高速化できます。
  • 文字列操作や動的な型付けが多い処理では、かえってCPythonより遅くなるケースもあるため、適切な箇所に限定して使うのがコツです。

Cython

  • Pythonのソースコードをほぼそのままに、型アノテーションを追加してC/C++にコンパイルし、拡張モジュールとしてインポートできます。
  • 事前コンパイルが必要で、型指定やメモリ管理に注意がいるものの、最適化次第ではC言語レベルの性能が得られます。

codon

https://docs.exaloop.io/codon

新しい高速化選択肢として注目されるのが Codon です。

まず、codonはPythonをコンパイルするプロジェクトではなく、Pythonライクでコンパイルが可能な言語です。そのため、現状のところ特に意識せずに書いた場合両者に互換性はありません。

Non-goals
❌ CPythonの代替品: CodonはCPythonの代替実装ではありません。

codon独自のパイプライン演算子

codonにはPythonにはない独自の演算子も存在します。

def add1(x):
    return x + 1

result = 2 |> add1
print(result)

https://docs.exaloop.io/codon/advanced/pipelines

なお、現時点では Windows をサポートしていません。
https://github.com/exaloop/codon/issues/69

高速化としてのcodon

Codonでコンパイルしたモジュールは、Python拡張モジュール（pyext）として配布できます。
純粋なPython環境でも動作するように、“モック” を用意しておく方法を以下に示します。

dataclass モック

Codonで定義したクラスをPythonから使う場合、@dataclass(python=True) デコレーターが必要になります。
Python標準には存在しないため、以下のように alt.py にモックを用意しておきます。
また、同じディレクトリに空のファイル alt.codon を置いておきます。

# alt.py
def dataclass(_cls=None, *, python=False):
    def wrap(cls):
        return cls

    if _cls is None:
        return wrap
    else:
        return wrap(_cls)

これを、モジュールの先頭で読み込むだけです。

from .alt import *

setup.py

ビルド時にCodonが利用可能かを判定し、可能な場合のみ拡張モジュールをコンパイルする設定例です。

from setuptools import setup
from exaloop.codon.build_ext import BuildCodonExt, CodonExtension
import os

codon_path = os.environ.get("CODON_PATH")

setup(
    name="fm_index",
    version="0.1",
    packages=["fm_index"] if not codon_path else [],
    ext_modules=(
        []
        if not codon_path
        else [
            CodonExtension(
                "fm_index",
                os.path.join("fm_index", "core.py"),
            ),
        ]
    ),
    cmdclass={"build_ext": BuildCodonExt},
)

jit デコレーター

実行時にコード中の関数を動的にコンパイルするデコレーターも利用可能です。ただし、この方式では実行環境にCodonが必要になります。
https://docs.exaloop.io/codon/interoperability/decorator

実際に計測する

リポジトリ: https://github.com/sk-uma/codon-fm-index
FM-Index の全文探索アルゴリズムを、純粋Python版とCodon版で比較した例です。

結果

環境 実行時間
pure python 31.83 s
with codon 0.91 s

約30倍の高速化が確認できました。（アルゴリズム最適化が不完全なため、標準検索より遅いケースもあります）

検証用コード
from random import choices, randint, seed
from fm_index import FMSimpleIndex
from time import time

seed(42)

start_time = time()

search_string = "".join(
    choices("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890", k=30000)
)

index = FMSimpleIndex(search_string)

start_idx = randint(0, len(search_string) - 1)
length = randint(1, randint(1, 10))
query = search_string[start_idx : start_idx + length]

end_time = time()

print(f"Time taken: {end_time - start_time} seconds")

おわりに

Python本来の手軽さを維持しつつ、PyPy、numba、Cython、Codonといった多様な手法を組み合わせることで、用途や規模に応じた最適化が可能です。特にCodonは成長中のプロジェクトであり、今後さらに機能が充実すれば、現実的な選択肢の一つになるかもしれません。

Discussion