Python 3.3 以上で使える Python/C API で文字列アクセスを高速化

試しに英語で Blog を書いてみたのですが、書くので精一杯で結局何が言いたいのか分からない感じになってしまったので今後は日本語 Blog 書いてから英訳しようと思います。

Python 3 は 3.2 まで、文字列を unicode に統一した関係で Python 2.7 に比べて遅くなったりメモリ効率が悪くなったりしてしまっていたのですが、 Python 3.3 で PEP 393 Flexible String Representation が導入されて改善されました。

PEP 393 は Python の内部だけではなく Python/C API にも変更を加えており、内部を理解しつつ新しい API を適切に使えば、バイト列と文字列の間の変換を行うような C 拡張を高速化することができます。そろそろ Python 3.2 のサポートを切れる時期なので、思い当たる人は目を通しておきましょう

PEP 393 の概要

PEP 393 では文字列オブジェクトの内部表現が幾つかに分かれています。

Compact (非 ASCII)

Compact は一番基本的な unicode オブジェクトです。

+--------+--------+------+-------+------+-------------+------+-------------+------+
| Header | length | hash | state | wstr | utf8_length | utf8 | wstr_length | DATA |
+--------+--------+------+-------+------+-------------+------+-------------+------+

Header:      Python オブジェクト共通のヘッダ
length:      文字列長 (codepoint数)
state:       各種フラグ
hash:        ハッシュ値キャッシュ
wstr:        wchar_t* 表現のキャッシュ
utf8_length: utf-8 にエンコードした時のバイト長
utf8:        utf-8 表現のキャッシュ
wstr_length: wstr の長さ
DATA:        NUL 終端したコードポイント列

DATA 内の各 codepoint を何バイトで表現するかは、その文字列が含む最大の codepoint によって 1byte, 2byte, 4byte のどれかが切り替わります。これを kind と呼んで、 state で管理しています。例えば文字列の codepoint が 0~255 の範囲に収まる場合、 1BYTE_KIND となり、DATA の大きさは NUL 終端を含めて文字列長+1 byte になります。

wstr と utf8 は、それぞれエンコードした結果のキャッシュです。これは PyUnicode_AsUTF8() などの API を呼び出した時に生成されます。例えば utf8 は次のようなコードで生成されています:

    if (PyUnicode_UTF8(unicode) == NULL) {
        assert(!PyUnicode_IS_COMPACT_ASCII(unicode));
        bytes = _PyUnicode_AsUTF8String(unicode, "strict");
        if (bytes == NULL)
            return NULL;
        _PyUnicode_UTF8(unicode) = PyObject_MALLOC(PyBytes_GET_SIZE(bytes) + 1);
        if (_PyUnicode_UTF8(unicode) == NULL) {
            PyErr_NoMemory();
            Py_DECREF(bytes);
            return NULL;
        }
        _PyUnicode_UTF8_LENGTH(unicode) = PyBytes_GET_SIZE(bytes);
        Py_MEMCPY(_PyUnicode_UTF8(unicode),
        PyBytes_AS_STRING(bytes),
        _PyUnicode_UTF8_LENGTH(unicode) + 1);
        Py_DECREF(bytes);
    }

    if (psize)
        *psize = PyUnicode_UTF8_LENGTH(unicode);
    return PyUnicode_UTF8(unicode);

このコードを読めば分かるように、一旦 PyUnicode_AsUTF8String() で utf-8 にエンコードした bytes オブジェクトを作ってから、その中身を malloc して確保した領域に memcpy しています。繰り返し利用される場合以外は PyUnicode_AsUTF8String() を使って自前で bytes オブジェクトの中を見たほうが速くて省メモリです。

Compact ASCII

文字列に含まれる文字が ASCII である場合、次のような等式が成り立ちます。

utf8_length = length
utf8 = DATA
wstr_length = length

そのため、 Compact から要らない要素を取り除いた特別な表現を利用します。

+--------+--------+------+-------+------+------+
| Header | length | hash | state | wstr | DATA |
+--------+--------+------+-------+------+------+

utf8 キャッシュが存在せず、 PyUnicode_AsUTF8() と PyUnicode_AsUTF8AndSize() は直接 DATA を返すので常に高速です。

Legacy

PEP 393 によって deprecated になった API のために用意された、古い構造です。 PyUnicode_Ready() で Compact 形式にインプレイス変換できます。

事例1: UltiraJSON

UltraJSON は高速な JSON エンコーダ／デコーダです。

典型的な JSON は多くの ASCII 文字列を含みます。そして Compact ASCII 表現の場合、 utf8 へのエンコードをスキップすることができます。それ以外の場合は、 PyUnicode_AsUTF8String() を使って bytes 型に変換してやります。

# https://github.com/esnme/ultrajson/pull/159
diff --git a/python/objToJSON.c b/python/objToJSON.c
index e56aa9b..a5b2f62 100644
--- a/python/objToJSON.c
+++ b/python/objToJSON.c
@@ -145,7 +145,17 @@ static void *PyStringToUTF8(JSOBJ _obj, JSONTypeContext *tc, void *outValue, siz
static void *PyUnicodeToUTF8(JSOBJ _obj, JSONTypeContext *tc, void *outValue, size_t *_outLen)
{
   PyObject *obj = (PyObject *) _obj;
-  PyObject *newObj = PyUnicode_EncodeUTF8 (PyUnicode_AS_UNICODE(obj), PyUnicode_GET_SIZE(obj), NULL);
+  PyObject *newObj;
+#if (PY_VERSION_HEX >= 0x03030000)
+  if(PyUnicode_IS_COMPACT_ASCII(obj))
+  {
+    Py_ssize_t len;
+    char *data = PyUnicode_AsUTF8AndSize(obj, &len);
+    *_outLen = len;
+    return data;
+  }
+#endif
+  newObj = PyUnicode_AsUTF8String(obj);
   if(!newObj)
   {
     return NULL;

このパッチの効果をマイクロベンチマークで確認してみます:

$ # Before
$ python3.4 -m timeit -n 10000 -s 'import ujson; x = ["a"*10]*100' 'ujson.dumps(x)'
10000 loops, best of 3: 15.8 usec per loop

$ # After
$ python3.4 -m timeit -n 10000 -s 'import ujson; x = ["a"*10]*100' 'ujson.dumps(x)'
10000 loops, best of 3: 7.14 usec per loop

事例2: Meinheld

Meinheld は高速な WSGI サーバーです。

HTTPヘッダのフィールド名は RFC で ASCII 文字列だと定義されていて、フィールド値はバイト列ですが PEP 3333 で latin-1 として扱うと定義されています。

HTTP ヘッダを作る場合は、 "HTTP_ACCEPT_ENCODING" のような文字列を作る必要があります。受信バッファには "accept-encoding" のような文字列が入っているので、インプレイスで大文字化と - から _ への変換を行っても、 HTTP_ をつけることができません。この場合、 PyUnicode_New(length, maxchar) の第二引数に 127 を与えると Compact ASCII 形式の空の unicode オブジェクトを作ってくれるので、そこに対して直接文字列を書き込む事で "HTTP_" をつけるためだけに受信バッファと別に中間バッファを挟む必要がありません。

一方 HTTP フィールド値の方は "HTTP_" のような prefix をつけたりしないので、 PyUnicode_New() を使わなくても直接受信バッファから unicode を生成できます。 PyUnicode_DecodeLatin1() を使えば、実際に受信バッファに ASCII 外の文字があるかどうかによって自動的に Compact ASCII か Compact かを選んで生成してくれます。

逆に unicode から HTTP ヘッダを生成する場合も、unicode が ASCII か Latin-1 でない場合はエラーにできるので、 DATA 部分を直接 writev(2) に渡すことができます。

static int
wsgi_to_bytes(PyObject *value, char **data, Py_ssize_t *len)
{
#ifdef PY3
    if (!PyUnicode_Check(value)) {
        PyErr_Format(PyExc_TypeError, "expected unicode object, value "
        "of type %.200s found", value->ob_type->tp_name);
        return -1;
    }
    if (PyUnicode_READY(value) == -1) {
        return -1;
    }
    if (PyUnicode_KIND(value) != PyUnicode_1BYTE_KIND) {
        PyErr_SetString(PyExc_ValueError,
        "unicode object contains non latin-1 characters");
        return -1;
    }
    *len = PyUnicode_GET_SIZE(value);
    *data = (char*)PyUnicode_1BYTE_DATA(value);
    return 0;
#else
    return PyBytes_AsStringAndSize(value, data, len);
#endif
}

meinheld の主に http ヘッダ生成側を最適化してベンチマークしたところ、だいたい 23.5k req/sec から 25k req/sec へパフォーマンスが向上しました。

ちなみに前回 picohttpparser の比較で Python 2.7 を利用したのは、ヘッダパース部分をリライトするときに新APIを使ってしまったため、 Python 3 では純粋に HTTP パーサーの性能差を見ることができなかったためです。今回は逆に picohttpparser 化前のソースコードを最適化したので、ヘッダパース側の最適化は手を抜いています。完全に最適化したらもう少し差が広がると思います。

最適化前:

$ ./wrk http://localhost:8000/
Running 10s test @ http://localhost:8000/
  2 threads and 10 connections
  Thread Stats   Avg      Stdev     Max   +/- Stdev
    Latency   385.68us   25.67us 843.00us   74.12%
    Req/Sec    12.42k   632.30    13.78k    59.64%
  234475 requests in 10.00s, 39.13MB read
Requests/sec:  23448.18
Transfer/sec:      3.91MB
$ ./wrk http://localhost:8000/
Running 10s test @ http://localhost:8000/
  2 threads and 10 connections
  Thread Stats   Avg      Stdev     Max   +/- Stdev
    Latency   379.78us   24.96us   0.90ms   73.86%
    Req/Sec    12.46k   639.91    14.00k    52.45%
  235688 requests in 10.00s, 39.33MB read
Requests/sec:  23569.38
Transfer/sec:      3.93MB
$ ./wrk http://localhost:8000/
Running 10s test @ http://localhost:8000/
  2 threads and 10 connections
  Thread Stats   Avg      Stdev     Max   +/- Stdev
    Latency   388.15us   24.65us 524.00us   73.31%
    Req/Sec    12.43k   623.66    13.78k    55.26%
  234899 requests in 10.00s, 39.20MB read
Requests/sec:  23490.25
Transfer/sec:      3.92MB

最適化後:

$ ./wrk http://localhost:8000/
Running 10s test @ http://localhost:8000/
  2 threads and 10 connections
  Thread Stats   Avg      Stdev     Max   +/- Stdev
    Latency   339.98us   65.87us   4.33ms   92.45%
    Req/Sec    13.56k     1.43k   15.44k    81.82%
  253189 requests in 10.00s, 42.26MB read
Requests/sec:  25319.67
Transfer/sec:      4.23MB
$ ./wrk http://localhost:8000/
Running 10s test @ http://localhost:8000/
  2 threads and 10 connections
  Thread Stats   Avg      Stdev     Max   +/- Stdev
    Latency   364.82us   78.85us   1.01ms   84.10%
    Req/Sec    12.99k     1.81k   15.44k    80.63%
  243685 requests in 10.00s, 40.67MB read
Requests/sec:  24368.90
Transfer/sec:      4.07MB
$ ./wrk http://localhost:8000/
Running 10s test @ http://localhost:8000/
  2 threads and 10 connections
  Thread Stats   Avg      Stdev     Max   +/- Stdev
    Latency   329.00us   22.40us 464.00us   73.99%
    Req/Sec    13.66k   760.36    15.44k    61.95%
  258730 requests in 10.00s, 43.18MB read
Requests/sec:  25873.60
Transfer/sec:      4.32MB