読者です 読者をやめる 読者になる 読者になる

StatsFragments

Python, R, Rust, 統計, 機械学習とか

Python でパイプ演算子を使いたい <2>

pandas Python

ネタ記事です。/ This is a joke post which makes no practical sense.


過去にこんなエントリを書いた。

sinhrks.hatenablog.com

R では パイプ演算子 %>% を使って連続した処理を記述できる。式に含まれる x, y, z は非標準評価 (NSE) によって data.frame の列として解決される。

# R (magrittr + dplyr)
df %>% mutate(x = y + z) %>% group_by(x) %>% summarize_each(funs(sum))

Python (pandas) ではほぼ同じ処理をメソッドチェインを使って書ける。チェインとパイプ演算子でどちらが読みやすいかは好みの問題だと思うものの、式の中に 何度も df が出てくるのはちょっとすっきりしない。

# Python (pandas)
df.assign(x=df['y'] + df['z']).groupby('x').sum()

上の式をこんな風に書けるとうれしい。

df >> assign(x=y+z) >> groupby(x) >> sum()

前のエントリでは マジックメソッドと関数定義で近いことをやったが、以下のような課題がある。

  • パイプ演算子で処理したいメソッドを、それぞれ関数として別に定義しなければならない
  • R の非標準評価のようなことはできない (ある変数を別の環境で評価することができない)

前者は手間をかければなんとかなるが、後者については NameError が発生してしまうためどうしようもない。

df.assign(x=y+z)
# NameError: name 'y' is not defined

上のような表現を式に含めるには、 lambda を使って無名関数にすればよい。lambda を使うと、下のように定義されていない関数 / 変数を式に含めることができる。

lambda: df >> assign(x=y+z)
# <function <lambda> at 0x102310620>

当然、この無名関数をそのまま評価すると NameError が発生する。定義した無名関数の中身を調べて、妥当な処理に置き換えてやればよい。

Python 標準では ast モジュールを使って 抽象構文木 (AST) の探索や置換ができる。が、今回はそこまで複雑なことをやるわけではないので、バイトコードを直接書き換えられるとうれしい。

バイトコードの置換について簡単な方法がないか調べたところ CodeTransformer というパッケージを見つけた。今回はこれを使いたい。

github.com

CodeTransformerとは

バイトコードに対するパターンマッチ⇨置換が簡単にできる。詳しい使い方はドキュメントを。

まずは関数がバイトコードとしてどのように表現されているのか確かめたい。これは標準の dis モジュールを使うのが簡単。

import dis
dis.dis(lambda x: x + 2)
#   1           0 LOAD_FAST                0 (x)
#               3 LOAD_CONST               1 (2)
#               6 BINARY_ADD
#               7 RETURN_VALUE

CodeTransformer を使うと適当なパターンにマッチするバイトコードを置換することができる。ドキュメントの例にある通り、バイトコード中の加算 (BINARY_ADD) を乗算 (BINARY_MULTIPLY) に置き換える CodeTransformer は以下のようになる。

import codetransformer
codetransformer.__version__
# '0.6.0'

from codetransformer import CodeTransformer, pattern
from codetransformer.instructions import BINARY_ADD, BINARY_MULTIPLY

class Multiply(CodeTransformer):
    
    @pattern(BINARY_ADD)
    def _add2mul(self, add_instr):
        yield BINARY_MULTIPLY().steal(add_instr)

pattern デコレータで指定したバイトコードにマッチした時、対応したメソッドが呼び出される。呼び出されるメソッドから置換後のバイトコードを返せばよい。

作成した CodeTransformer を使って関数を置換する。

mul = Multiply()(lambda x: x + 2)
mul(5)
# 10

Multiply によって置き換えられた後のバイトコードを確認すると、BINARY_ADD が BINARY_MULTIPLY に置換されていることが確かめられる。

dis.dis(mul)
#   1           0 LOAD_FAST                0 (x)
#               3 LOAD_CONST               0 (2)
#               6 BINARY_MULTIPLY
#               7 RETURN_VALUE

パイプ演算子を使いたい

ここから、pandas と組み合わせて使ってみる。まずはデータを準備する。

import pandas as pd
df = pd.DataFrame({'A': [1, 2, 3, 4, 5, 6],
                   'B': list('ABABAB'),
                   'C': [1, 2, 3, 1, 2, 3]})
df
#    A  B  C
# 0  1  A  1
# 1  2  B  2
# 2  3  A  3
# 3  4  B  1
# 4  5  A  2
# 5  6  B  3

パイプ演算子を処理するには、lambda: df >> head()バイトコードlambda: df.head()バイトコードに置換すればよい。置換前、置換後のバイトコードをそれぞれ確認すると、

# 置換前
dis.dis(lambda: df >> head())
#   1           0 LOAD_GLOBAL              0 (df)
#               3 LOAD_GLOBAL              1 (head)
#               6 CALL_FUNCTION            0 (0 positional, 0 keyword pair)
#               9 BINARY_RSHIFT
#              10 RETURN_VALUE

# 置換後
dis.dis(lambda: df.head())
#   1           0 LOAD_GLOBAL              0 (df)
#               3 LOAD_ATTR                1 (head)
#               6 CALL_FUNCTION            0 (0 positional, 0 keyword pair)
#               9 RETURN_VALUE

比べてみると、以下のような変換を行えばよさそうだ。

  • 置換前のパターン: LOAD_GLOBAL → CALL_FUNCTION → BINARY_RSHIFT
  • 置換後のパターン: LOAD_ATTR → CALL_FUNCTION

連続したパターンとマッチさせるには、pattern デコレータと対応するメソッドに複数の引数を指定すればよい。

from codetransformer.instructions import BINARY_RSHIFT, LOAD_GLOBAL, LOAD_ATTR, CALL_FUNCTION

class Pipify(CodeTransformer):
    
    @pattern(LOAD_GLOBAL, CALL_FUNCTION, BINARY_RSHIFT)
    def _pipe(self, load, call, rshift):
        # LOAD_GLOBAL を LOAD_ATTR に置換
        yield LOAD_ATTR(load.arg)
        yield call
        # BINARY_SHIFT は無視

これにパイプ演算を含む無名関数を渡すと、期待通り動いているように見える。

Pipify()(lambda: df >> head())()
#    A  B  C
# 0  1  A  1
# 1  2  B  2
# 2  3  A  3
# 3  4  B  1
# 4  5  A  2

dis.dis(Pipify()(lambda: df >> head()))
#   1           0 LOAD_GLOBAL              0 (df)
#               3 LOAD_ATTR                1 (head)
#               6 CALL_FUNCTION            0 (0 positional, 0 keyword pair)
#               9 RETURN_VALUE

もっとも、上のクラスは非常に単純な条件でしか動かない。たとえば、パイプで接続される関数(メソッド)に引数がある場合、patternの定義とマッチしなくなるため置換が行われない。

# NG
Pipify()(lambda: df >> head(2))()
# NameError: name 'head' is not defined

dis.dis(lambda: df >> head(2))
#   1           0 LOAD_GLOBAL              0 (df)
#               3 LOAD_GLOBAL              1 (head)
#               6 LOAD_CONST               1 (2)
#               9 CALL_FUNCTION            1 (1 positional, 0 keyword pair)
#              12 BINARY_RSHIFT
#              13 RETURN_VALUE

もう少しマシな CodeTransformer の定義を考える。パイプ演算子の記法から、 LOAD_GLOBAL から次の BINARY_RSHIFT までを 1 単位とした pattern 定義でマッチすればよいような気がする。

CodeTransformer の定義は以下のようになる。(~BINARY_RSHIFT)[plus]BINARY_SHIFT 以外の任意のコードの繰り返しとマッチする。

from codetransformer import plus

class Pipify(CodeTransformer):
    
    @pattern(LOAD_GLOBAL, (~BINARY_RSHIFT)[plus], BINARY_RSHIFT)
    def _pipe(self, load, *args):
        
        # メソッド呼び出しの処理
        if hasattr(self, 'first_load'):
            yield LOAD_ATTR(load.arg)
        else:
            # 最初の LOAD_GLOBAL は置き換えない
            self.first_load = load
            yield load
            
        # メソッドへの引数の処理
        for arg in args[:-1]:
            if isinstance(arg, LOAD_GLOBAL):
                # 関数の一番最初の LOAD_GLOBAL から解決
                # この定義では、パイプ演算中に更新されたデータやグローバル変数には
                # アクセスできない
                yield self.first_load
                yield LOAD_ATTR(arg.arg)
            else:
                yield arg

        # 末尾は BINARY_RSHIFT なので無視

コメントで "メソッドへの引数の処理" とある部分は、 (記載のとおり限定的だが) R の非標準評価に近い処理になる。

この定義を使うと、上で失敗したメソッド呼び出しも処理できる。

Pipify()(lambda: df >> head(2))()
#    A  B  C
# 0  1  A  1
# 1  2  B  2

dis.dis(Pipify()(lambda: df >> head(2)))
#   1           0 LOAD_GLOBAL              0 (df)
#               3 LOAD_GLOBAL              0 (df)
#               6 LOAD_ATTR                1 (head)
#               9 LOAD_CONST               0 (2)
#              12 CALL_FUNCTION            1 (1 positional, 0 keyword pair)
#              15 RETURN_VALUE

いちいち Pipify()(lambda: df >> ...)() と書くのは面倒なので、パイプ演算を開始する関数 p を定義する。この関数を使うとパイプ演算は p(lambda: df >> ...) のように書ける。

def p(expr):
    return Pipify()(expr)()

以下、いくつかのパターンを試す。それぞれ、冒頭にコメントとして記載した処理と同じ結果になる。

# df.head()
p(lambda: df >> head())
#    A  B  C
# 0  1  A  1
# 1  2  B  2
# 2  3  A  3
# 3  4  B  1
# 4  5  A  2

# df.head(2)
p(lambda: df >> head(2))
#    A  B  C
# 0  1  A  1
# 1  2  B  2

パイプ演算中の BNameError とはならず df.B として解決される (非標準評価もどき)。

# df.groupby(df.B).sum()
p(lambda: df >> groupby(B) >> sum())
#     A  C
# B       
# A   9  6
# B  12  6

もう少し複雑な処理もできる。.assigndf.A + df.C の結果からなる X 列を新規で作成し、B 列の値ごとに集約して合計を取る。

# df.assign(X=df.A+df.C).groupby(df.B).sum()
p(lambda: df >> assign(X=A+C) >> groupby(B) >> sum())
#     A  C   X
# B           
# A   9  6  15
# B  12  6  18

改行を含む場合は括弧で囲む。

p(lambda: (df >>
           assign(X=A+C) >>
           groupby(B) >>
           sum()))
# 略

まとめ

  • CodeTransformerバイトコードを置換することで、パイプ演算のような処理が書ける。また、非標準評価に近いこともできる。
  • R を使ったほうがいい。