Python 関数で処理

Polars の組み込み関数と演算子で、よくある処理のほとんどは書けます。それでも表現できないものが出てきます。外部の表を引く、API を呼ぶ、既存の Python アルゴリズムをそのまま使う、といった処理です。そんなときは自前の関数を演算式の連鎖や集約の中に差し込みます。グループ化も並列化もメモリ管理も Polars が受け持ち、書くのは業務ロジックの部分だけです。

この章では「介入の度合い」が低い順に手段を紹介します。pipe はデータに触らず関数を鎖に差し込むだけ、map_batches / map_groups は Series のまとまりを関数に渡し、map_elements は要素ごとに呼び、rolling_map はウィンドウを渡します。後ろほど柔軟で、後ろほど遅くなります。

import polars as pl
import numpy as np
from helper.jupyter import row

pipe

Polars では、コードの可読性を上げたり、関数を柔軟に適用したりするために pipe メソッドを使えます。主に次の 2つの場面で使います。

1. DataFrame.pipe()

2. Expr.pipe()

それぞれの使い方を例で確認します。

DataFrame.pipe

DataFrame.pipe メソッドは、関数を処理の連鎖の中に差し込むためのものです。複雑な処理を段階的に書けるので、コードの可読性と保守性が上がります。

pipe() の最初の引数には、適用する関数を渡します。その関数の最初の引数として DataFrame 自体が渡されます。残りの位置引数とキーワード引数は、そのまま関数に渡されます。

df = pl.DataFrame({
    "A": [1, 2, 3],
    "B": [4, 5, 6],
    "C": [7, 8, 9]
})

def add_column(df, col_name, value):
    return df.with_columns(pl.lit(value).alias(col_name))

def multiply_column(df, col_name, factor):
    return df.with_columns((pl.col(col_name) * factor).alias(col_name))

df_res = (
    df.pipe(add_column, col_name="D", value=10)
      .pipe(multiply_column, col_name="A", factor=2)
)

row(df, df_res)

shape: (3, 3) A B C i64 i64 i64 1 4 7 2 5 8 3 6 9

shape: (3, 4) A B C D i64 i64 i64 i32 2 4 7 10 4 5 8 10 6 6 9 10

Expr.pipe

演算式の pipe メソッドを使うと、演算式に独自の変換を差し込めます。次の例では、2次式を計算する関数 quadratic() を定義し、pipe で複数の演算式に適用します。

pipe() の最初の引数には、Expr を受け取る関数を指定します。その関数の最初の引数として演算式が渡され、残りの引数もそのまま引き継がれます。

次のコードでは、.pipe(quadratic, ...) で次の計算を行います。

• A 列に対して \( x^2 + 2 \cdot x + 3 \) を計算し、列名に "_2" を追加します。

• B 列と C 列の合計に対して \( x^2 + 2 \cdot x \) を計算し、列名に "_q" を追加します。

def quadratic(x, a=0, b=0, c=0, suffix='_2'):
    return (a * x ** 2 + b * x + c).name.suffix(suffix)

df_res = (
    df.with_columns(
        pl.col("A").pipe(quadratic, 1, 2, 3),
        (pl.col("B") + pl.col("C")).pipe(quadratic, a=1, b=2, suffix='_q')
    )
)

row(df, df_res)

shape: (3, 3) A B C i64 i64 i64 1 4 7 2 5 8 3 6 9

shape: (3, 5) A B C literal_2 literal_q i64 i64 i64 i64 i64 1 4 7 6 143 2 5 8 11 195 3 6 9 18 255

map_batches

pl.map_batches は、複数の列をまとめて処理したいときに使う関数です。指定した列をユーザー定義関数に渡し、その結果を新しい列として追加できます。

pl.map_batches(
    column_names, 
    f: Callable[[Sequence[Series]], Series | Any], 
    return_dtype: PolarsDataType | pl.DataTypeExpr | None = None
)

• column_names：ユーザー関数に渡す列名または式のリスト。

• f：データを処理するユーザー定義関数。この関数は、指定した列の pl.Series のリストを受け取り、新しい pl.Series またはその他の値を返します。

• return_dtype：ユーザー関数の戻り値の要素のデータ型。既定値が None の場合は、データ型が自動で推定されます。

次の例では、a 列と b 列にユーザー定義関数 hypot を適用し、結果を新しい列 c として追加します。

df = pl.DataFrame(
    {
        "a": [3, 3, 3, 4],
        "b": [4, 12, 6, 7],
        "g": ['A', 'B', 'A', 'B']
    }
)

def hypot(args:list[pl.Series]):
    a, b = args
    return (a**2 + b**2)**0.5
    
df.with_columns(
    pl.map_batches(['a', 'b'], hypot, return_dtype=pl.Float64).alias('c')
)

shape: (4, 4)

a	b	g	c
i64	i64	str	f64
3	4	"A"	5.0
3	12	"B"	12.369317
3	6	"A"	6.708204
4	7	"B"	8.062258

column_names 引数には演算式も渡せます。次のコードでは、a 列の値に 1を加えた結果と b 列の値で計算します。

df.with_columns(
    pl.map_batches([pl.col('a') + 1, 'b'], hypot).alias('c')

)

shape: (4, 4)

a	b	g	c
i64	i64	str	f64
3	4	"A"	5.656854
3	12	"B"	12.649111
3	6	"A"	7.211103
4	7	"B"	8.602325

演算式とユーザー関数以外に、pl.map_batches() には次の引数があります。

• return_dtype（既定: None）: ユーザー関数が返す Series の要素のデータ型 を指定します。 指定しない場合、Polars がユーザー関数の出力型を推定します。

• is_elementwise（既定: False）: ユーザー関数が 「要素ごとの処理（elementwise）」 であることを Polars に伝えるフラグです。 True にすると、グループ化の有無に関係なく、入力データ全体がまとめてユーザー関数に渡されます。

• returns_scalar（既定: False）: ユーザー関数が スカラー値を返す 場合に指定します。 指定しないと、Polars は Series を返す関数として扱います。

次の f1()・f2()・f3() は、それぞれの引数の働きを示す例です。

• f1()：要素同士の計算のみを行い、Series を返す

• f2()：要素同士の計算ではないが、Series を返す

• f3()：Series ではなく スカラー値 を返す

import threading

lock = threading.Lock()

def f1(args):
    with lock:
        a, b = args
        print(f"{a.to_list()} {b.to_list()}")
        return a + b

def f2(args):
    with lock:
        a, b = args
        print(f"{a.to_list()} {b.to_list()}")
        return a - b.min()

def f3(args):
    with lock:
        a, b = args
        print(f"{a.to_list()} {b.to_list()}")
        return a.max() - b.min()

return_dtype を指定しない場合、Polars は出力型を推定するために、ダミーの入力（[1, 1]）で関数を一度実行します。

df.select(
    pl.map_batches(['a', 'b'], f1)
);

[1, 1] [1, 1]
[3, 3, 3, 4] [4, 12, 6, 7]

return_dtype を指定すると、このダミー実行がなくなります。

df.select(
    pl.map_batches(['a', 'b'], f1, return_dtype=pl.Int64)
);

[3, 3, 3, 4] [4, 12, 6, 7]

agg() の中で使うと、is_elementwise の意味がはっきりします。False の場合は、各グループのデータがグループごとにユーザー関数へ渡されます。常に正しい結果が得られますが、関数が複数回呼ばれるぶん遅くなります。

df.group_by('g', maintain_order=True).agg(
    pl.map_batches(['a', 'b'], f1, return_dtype=pl.Int64, is_elementwise=False).alias('f1'),
    pl.map_batches(['a', 'b'], f2, return_dtype=pl.Int64, is_elementwise=False).alias('f2'),
)

[3, 3] [4, 6]
[3, 3] [4, 6]
[3, 4] [12, 7]
[3, 4] [12, 7]

shape: (2, 3)

g	f1	f2
str	list[i64]	list[i64]
"A"	[7, 9]	[-1, -1]
"B"	[15, 11]	[-4, -3]

True の場合、グループとは無関係に、すべてのデータが一度に ユーザー関数に渡されます。

• f1() のような「要素同士の計算」であれば正しい結果

• f2() のように「要素ごとの計算ではない」関数は 誤った結果になる

df.group_by('g', maintain_order=True).agg(
    pl.map_batches(['a', 'b'], f1, return_dtype=pl.Int64, is_elementwise=True).alias('f1'),
    pl.map_batches(['a', 'b'], f2, return_dtype=pl.Int64, is_elementwise=True).alias('f2'),
)

[3, 3, 3, 4] [4, 12, 6, 7]
[3, 3, 3, 4] [4, 12, 6, 7]

shape: (2, 3)

g	f1	f2
str	list[i64]	list[i64]
"A"	[7, 9]	[-1, -1]
"B"	[15, 11]	[-1, 0]

f3() はスカラー値（1つの数値）を返すため、returns_scalar=True を指定しないとエラーになります。

df.group_by('g', maintain_order=True).agg(
    pl.map_batches(['a', 'b'], f3, return_dtype=pl.Int64, returns_scalar=True).alias('f3'),
)

[3, 3] [4, 6]
[3, 4] [12, 7]

shape: (2, 2)

g	f3
str	i64
"A"	-1
"B"	-3

演算式にも map_batches() メソッドがあります。次の例は、Expr.map_batches() で演算式の計算結果をユーザー関数に渡すものです。2列のデータは 2回に分けて square() 関数に渡されます。

def square(s):
    return s**2

df.select(
    pl.col('a', 'b').map_batches(square)
)

shape: (4, 2)

a	b
i64	i64
9	16
9	144
9	36
16	49

pl.struct() で複数の列を 1つの構造体列にまとめ、そのあと map_batches() で独自の関数を適用できます。次のコードは、複数の列をまとめて処理する例です。

def square2(s):
    return s.struct.field('a')**2 + s.struct.field('b')**2

df.select(
    pl.struct('a', 'b').map_batches(square2)
)

shape: (4, 1)

a
i64
25
153
45
65

map_groups

pl.map_groups() を使うと、各グループのデータをユーザー定義関数で処理し、結果をまとめて受け取れます。実装は map_groups() も map_batches() も同じ仕組みですが、agg() の中で使うなら map_groups() のほうが意図が伝わり、読みやすくなります。

次の例では、各グループごとに a 列の最大値と b 列の最大値の比を計算しています。

df = pl.DataFrame(
    {
        "a": [3, 3, 3, 4],
        "b": [4, 12, 6, 7],
        "g": ['A', 'B', 'A', 'B']
    }
)

def ratio_max(args):
    a, b = args
    return a.max() / b.max()

df.group_by('g').agg(
    pl.map_groups(['a', 'b'], ratio_max, returns_scalar=True)
)

shape: (2, 2)

g	a
str	f64
"A"	0.5
"B"	0.333333

map_elements

Expr.map_elements() を使うと、演算式の各値を 1つずつカスタム関数に渡し、その結果を新しい列として計算できます。次の例では、関数 f で各値を処理します。

def f(x):
    if x > 10:
        return 'large'
    elif x > 5:
        return 'middle'
    else:
        return 'small'

df.with_columns(
    pl.col.b.map_elements(f, return_dtype=pl.String).alias('b_category')
)

shape: (4, 4)

a	b	g	b_category
i64	i64	str	str
3	4	"A"	"small"
3	12	"B"	"large"
3	6	"A"	"middle"
4	7	"B"	"middle"

Expr.map_elements() は各値を 1つずつ処理するため、データ量が多いと遅くなります。そのため、ベクトル演算や Polars の組み込み関数では書けない場合にだけ使います。例えば、上の例は次のようにベクトル化した条件分岐でも書けます。

df.with_columns(
    pl.when(pl.col('b') > 10).then(pl.lit('large'))
      .when(pl.col('b') > 5).then(pl.lit('middle'))
      .otherwise(pl.lit('small'))
      .alias('b_category')
)

shape: (4, 4)

a	b	g	b_category
i64	i64	str	str
3	4	"A"	"small"
3	12	"B"	"large"
3	6	"A"	"middle"
4	7	"B"	"middle"

agg() や over() の中で map_elements() を使っても、ユーザー関数にはスカラー値が渡されます。

次の例では、各要素に対して関数 f が呼ばれ、値が順番に渡されることが確認できます。

def f(x):
    print("---")
    print(x)
    return x

g = df.group_by('g', maintain_order=True)
g.agg(pl.col('a').map_elements(f, return_dtype=pl.Int64))

---
3
---
3
---
3
---
4

shape: (2, 2)

g	a
str	list[i64]
"A"	[3, 3]
"B"	[3, 4]

グループごとのデータをユーザー関数に渡したいときは、map_batches() か .implode().map_elements() を使う必要があります。map_elements() の return_dtype 引数は、map_batches() と少し意味が異なります。

• map_batches() では、return_dtype は ユーザー関数の戻り値の要素型 を指定します。

• map_elements() では、return_dtype は ユーザー関数の戻り値そのもののデータ型 を指定します。

例えば .implode().map_elements() の場合、ユーザー関数の戻り値は Series になるため、return_dtype には pl.List(...) を指定します。

g.agg(pl.col('a').implode().map_elements(f, return_dtype=pl.List(pl.Int64)))

---
shape: (2,)
Series: '' [i64]
[
	3
	3
]
---
shape: (2,)
Series: '' [i64]
[
	3
	4
]

shape: (2, 2)

g	a
str	list[i64]
"A"	[3, 3]
"B"	[3, 4]

g.agg(pl.col('a').map_batches(f, return_dtype=pl.Int64))

---
shape: (2,)
Series: '' [i64]
[
	3
	3
]
---
shape: (2,)
Series: '' [i64]
[
	3
	4
]

shape: (2, 2)

g	a
str	list[i64]
"A"	[3, 3]
"B"	[3, 4]

rolling_map

pl.rolling_mean() などの移動ウィンドウ関数では対応できないときは、rolling_map() でウィンドウ内のデータを独自の関数に渡せます。次のコードは、各ウィンドウのデータをユーザー関数に渡し、結果を新しい列として追加します。

Tip

pl.Series._s：Series オブジェクトの内部データにアクセスするための属性です。実際のデータや、そのメモリの ID を確認できます。

arguments リストには、各ウィンドウのデータと、その Series オブジェクトのメモリ ID が記録されます。メモリを効率よく使うため、同じ Series オブジェクトが複数回渡されていることが分かります。

df = pl.DataFrame(
    {
        "a": [2, 4, 0, 3, 1, 6]
    }
)

arguments = []
def f(s):
    arguments.append((id(s._s), s.to_list()))
    return s.mean()

df.select(
    pl.col('a').rolling_map(f, window_size=3, center=True, min_samples=1)
)

for id_, arg in arguments:
    print(id_, arg)

1758068553648 [2, 4]
1758068551056 [2, 4, 0]
1758068544432 [4, 0, 3]
1758068553648 [0, 3, 1]
1758068551056 [3, 1, 6]
1758068544432 [1, 6]

ufunc

NumPy の ufunc 関数を直接使って、列あるいは演算式を計算できます。例えば、np.hypot は 2つの値の平方和の平方根を計算する関数で、Polars では次のように使えます。

df = pl.DataFrame(
    {
        "a": [3, 3, 3, 4],
        "b": [4, 12, 6, 7],
        "g": ['A', 'B', 'A', 'B']
    }
)

df.with_columns(
    np.hypot(pl.col("a"), pl.col("b")).alias("c")
)

shape: (4, 4)

a	b	g	c
i64	i64	str	f64
3	4	"A"	5.0
3	12	"B"	12.369317
3	6	"A"	6.708204
4	7	"B"	8.062258

np.hypot() で列を計算するとき、Polars の内部で何が起きているかを確かめるために、演算式を見てみます。この演算式は a 列を構造体列に変換し、b 列を "argument_1" として追加します。その構造体列を python_udf() で処理し、この関数の中で NumPy の np.hypot() を呼び出します。

np.hypot(pl.col("a"), pl.col("b"))

as_struct("a", col("b").alias("argument_1")).python_udf()

列名とフィールド名変更

Expr.name.map()

列名をユーザー関数で変更します。

Expr.name.map_fields()

Struct の field 名をユーザー関数で変更します。

def rename(n):
    return n.upper() + n
    
df.select(pl.all().name.map(rename))

shape: (4, 3)

Aa	Bb	Gg
i64	i64	str
3	4	"A"
3	12	"B"
3	6	"A"
4	7	"B"

DataFrame.map_rows()

行を Tuple としてユーザー関数に渡します。

def f(row):
    a, b, g = row
    return f"{g}:{a + b}", f"{g}:{a - b}"
    
df.map_rows(f).rename({"column_0": "add", "column_1": "sub"})

shape: (4, 2)

add	sub
str	str
"A:7"	"A:-1"
"B:15"	"B:-9"
"A:9"	"A:-3"
"B:11"	"B:-3"

行を辞書として受け取りたいときは、DataFrame.iter_rows() を使います。

for row in df.iter_rows(named=True):
    print(row)

{'a': 3, 'b': 4, 'g': 'A'}
{'a': 3, 'b': 12, 'g': 'B'}
{'a': 3, 'b': 6, 'g': 'A'}
{'a': 4, 'b': 7, 'g': 'B'}