演算式#

演算式(expression)は、データフレームの操作や変換の中心になる要素です。Polars の演算式は「列に作用する計算」を定義するもので、それ自体は結果を作りません。select() や with_columns() に渡した時点で初めて評価されます。

計算が後回しにされるからこそ、Polars はクエリ全体をまとめて最適化できます。スキャンの統合、不要な列の刈り込み、重複計算の再利用がその例です。この章では、演算式の使い方から、列の選び方、出力名の決まり方、出力の長さ、そしてよく使う演算式の一覧までを順に見ていきます。

import polars as pl
import numpy as np
from polars import selectors as cs
from helper.jupyter import row, capture_except

演算式を使用するメソッド#

ここでは、特によく使う3つのメソッドを簡単に説明します。select は出力する列を決め、with_columns は元の表に何を足すかを決め、filter は残す行を決めます。違いは「結果に対して何をするか」だけで、演算式そのものの書き方は同じです。

select#

DataFrame.select() は、データフレームから特定の列を選んだり、新しい列を計算したりするメソッドです。select() に渡した各演算式が、結果のデータフレームの列になります。演算式はリスト、位置引数、またはキーワード引数で指定できます。

select() と select_seq() はどちらも、元のデータフレームのデータを使って演算式を計算します。ただし、select() は並列で演算式を計算するのに対し、select_seq() は演算式を順番に計算します。1つの演算式あたりの計算が軽いときは、並列化の調整コストが利点を上回るため select_seq() が向いています。順番に計算すると中間結果が同時に残らないので、ピークメモリも低くなります。

どちらも入力のデータフレームに対して評価する点は同じで、同じ呼び出しの中で他の式の結果を参照することはできません。前後関係が必要なときは、select() や with_columns() を 2回に分けて書きます。

次の例では、select() および select_seq() の使用例を示します。

df = pl.DataFrame({
    "a": [1, 2, 3],
    "b": [4, 5, 6],
    "c": [7, 8, 9]
})

df1 = df.select(["a", pl.col("c") * 2])
df2 = df.select("a", pl.col("c") * 2)
df3 = df.select_seq(
    a = pl.col("a") * 2,         # "a" を2倍
    a_plus_c = pl.col("a") + pl.col("c") #❶
) 
row(df, df1, df2, df3)
shape: (3, 3)
abc
i64i64i64
147
258
369
shape: (3, 2)
ac
i64i64
114
216
318
shape: (3, 2)
ac
i64i64
114
216
318
shape: (3, 2)
aa_plus_c
i64i64
28
410
612

❶ 上の例では、select_seq() で a_plus_c の計算に使われる pl.col("a") は、元のデータの a 列を参照します。同じ呼び出しの中で a 列を2倍にしていても、その結果ではなく元の値が使われます。

with_columns#

DataFrame.with_columns() は DataFrame.select() と似ていますが、計算で作った新しい列を元の列と一緒に返します。既存のデータフレームに列を追加する操作に使えます。

次の例では、with_columns() を使った列の追加方法を示しています。df1 に新しい列 a_plus_b と const が追加され、df2 では既存の a 列が更新され、新しい列 a_plus_c が追加されています。

df1 = df.with_columns(
    a_plus_b = pl.col("a") + pl.col("b"),  # "a" + "b" の列を追加
    const = 10  # 定数列を追加
)

df2 = df.with_columns(
    a = pl.col("a") * 2,
    a_plus_c = pl.col("a") + pl.col("c")
)

row(df1, df2)
shape: (3, 5)
abca_plus_bconst
i64i64i64i64i32
147510
258710
369910
shape: (3, 4)
abca_plus_c
i64i64i64i64
2478
45810
66912

filter#

DataFrame.filter() は、演算式が返すブール値の条件に合う行だけを返します。次は filter() を使った例です。

df1 = df.filter(pl.col("a") > 1)  # "a" > 1 の行を選択
df2 = df.filter((pl.col("a") > 1) & (pl.col("b") < 6)) # 条件のAND
df3 = df.filter((pl.col("a") >= 3) |  (pl.col("b") <= 4)) # 条件のOR
df4 = df.filter(pl.col("a") > 1, pl.col("b") < 6) # 複数の式の場合は &で結合と同じ
row(df, df1, df2, df3, df4)
shape: (3, 3)
abc
i64i64i64
147
258
369
shape: (2, 3)
abc
i64i64i64
258
369
shape: (1, 3)
abc
i64i64i64
258
shape: (2, 3)
abc
i64i64i64
147
369
shape: (1, 3)
abc
i64i64i64
258

列の選択#

ほとんどの演算式は列の選択から始まります。列を選ぶ方法はいくつかあり、それぞれ用途が異なります。以下で順に説明します。

  • pl.col() — 列名やデータ型で列を選びます。

  • pl.nth() — N 番目の列を選びます。

  • pl.selectors — 条件を組み合わせて列を選びます。

pl.col()#

pl.col() は、列名やデータ型を指定して列を選ぶ式を作ります。

  • 列名での選択 — 列名を文字列で指定します。

df = pl.DataFrame({
    "a": [1], "b": [4.0], "c": [7], "x1":[10.0], "y1":[20.0], "x2":[4], "y2":[10], "xyz":[10.0]})
selected_df = df.select(pl.col("a"))
print(selected_df.columns)
['a']
  • 複数の列名での選択 — リストで複数の列名を指定できます。

selected_df = df.select(pl.col(["a", "b"]))
print(selected_df.columns)
['a', 'b']
  • データ型での選択 — データ型を指定すると、その型を持つ列をすべて選べます。

selected_df = df.select(pl.col(pl.Int64))
print(selected_df.columns)
['a', 'c', 'x2', 'y2']

pl.nth()#

pl.nth() は、N 番目の列を選びます。インデックスは0から始まります。

  • N 番目の列の選択 — インデックスを指定して列を選びます。

selected_df = df.select(pl.nth(0))
print(selected_df.columns)
['a']
  • 複数のインデックスでの選択 — インデックスをリストで指定し、対応する列を選びます。

selected_df = df.select(pl.nth([0, 2]))
print(selected_df.columns)
['a', 'c']

pl.all()#

pl.all() はデータフレームのすべての列を選びます。データフレーム全体をまとめて操作したいときや、他の選択条件と組み合わせたいときに便利です。

selected_df = df.select(pl.all())
print(selected_df.columns)
['a', 'b', 'c', 'x1', 'y1', 'x2', 'y2', 'xyz']

pl.exclude()#

pl.exclude() は指定した列を除外します。特定の列を除いたデータフレームを簡単に作れます。

  • 特定の列を除外

selected_df = df.select(pl.exclude("b"))
print(selected_df.columns)
['a', 'c', 'x1', 'y1', 'x2', 'y2', 'xyz']
  • 複数の列を除外

selected_df = df.select(pl.exclude(["a", "c"]))
print(selected_df.columns)
['b', 'x1', 'y1', 'x2', 'y2', 'xyz']

selectors#

polars.selectors モジュール(cs という別名でインポートします)は、条件に応じて列を選ぶための機能をまとめたものです。以下、主な関数を説明します。

データ型で列を選択#

cs.by_dtype() で、指定したデータ型の列を選びます。

selected_df = df.select(cs.by_dtype(pl.Int64))
print(selected_df.columns)
['a', 'c', 'x2', 'y2']

cs.integer() は、整数型の列をすべて選びます。Polars の整数型には Int8、Int16、Int32、Int64、UInt8、UInt16、UInt32、UInt64 などがあります。

cs.float() は、浮動小数点型の列をすべて選びます。Polars の浮動小数点型には Float32 と Float64 があります。

selected_df = df.select(cs.integer())
print(selected_df.columns)
['a', 'c', 'x2', 'y2']
selected_df = df.select(cs.float())
print(selected_df.columns)
['b', 'x1', 'y1', 'xyz']

cs.matches()#

正規表現で列名を選びます。Polars 2.0 の cs.by_name() は完全一致しかないため、パターンで選ぶときは cs.matches() を使います。

# x で始まる列、または数字で終わる列
selected_df = df.select(cs.matches(r"^x|y\d$"))
print(selected_df.columns)
['x1', 'y1', 'x2', 'y2', 'xyz']

cs.starts_with() と cs.ends_with()#

接頭辞と接尾辞で列を選びます。

selected_df = df.select(cs.starts_with("x"))
print(selected_df.columns)
['x1', 'x2', 'xyz']
selected_df = df.select(cs.ends_with("1"))
print(selected_df.columns)
['x1', 'y1']

集合演算#

上で説明した関数を組み合わせて、次のような集合演算を行えます。

  • 結合(|)— 複数の列選択を結合し、選択範囲を広げます。たとえば、特定の列とすべての列を組み合わせて選べます。

  • 差集合(-)— 特定の列を除外します。たとえば、すべての列から特定の列を除けます。

  • 交差(&)— 複数の条件に一致する列を選びます。たとえば、特定のパターンに一致し、特定のデータ型を持つ列を選べます。

  • 補集合(~)— 特定の条件に一致しない列を選びます。たとえば、特定のパターンに一致しない列を選べます。

次の選択文は、x で始まる列を選び、その中から 2 で終わる列を除外します。

selected_df = df.select(cs.starts_with('x') - cs.ends_with('2'))
print(selected_df.columns)
['x1', 'xyz']

次の選択文は、整数型の列と 1 で終わる列を選びます。

selected_df = df.select(cs.integer() | cs.ends_with('1'))
print(selected_df.columns)
['a', 'c', 'x1', 'y1', 'x2', 'y2']

列同士の演算式#

次の 3つの式は、それぞれ違う形の列演算です。

  • e1 — 列と列の演算。A と B を要素ごとに足します。

  • e2 — 2列と 1列の演算。実際には A + B と B + B の 2つが計算されます。

  • e3 — セレクタと式の組み合わせ。数値型の列すべてと B を足すので、式の数はデータの列構成で決まります。

e1 = pl.col("A") + pl.col("B")
e2 = pl.col("A", "B") + pl.col("B")
e3 = cs.numeric() + pl.col("B")  # セレクタと式はそのまま合成できる

print(e1)
print(e2)
print(e3)
[(col("A")) + (col("B"))]
[(cs.by_name('A', 'B', require_all=true)) + (col("B"))]
[(cs.numeric()) + (col("B"))]

次のプログラムは、LazyFrame.explain() を使って、式の演算の中身を表示します。

df = pl.DataFrame(dict(A=[1, 2], B=[10, 20], C=[100, 200], X=[1, 1], Y=[2, 2]))
ldf = df.lazy()
query = ldf.with_columns(C=e1)
print(query.explain())
simple π 5/5 ["A", "B", "C", "X", "Y"]
   WITH_COLUMNS:
   [(col("A") + col("B")).alias("C")] 
    DF ["A", "B", "C", "X", ...]; PROJECT["A", "B", "X", "Y"] 4/5 COLUMNS

複数列のセレクタがある場合、複数の式が作成されます。

query = ldf.select(e2)
print(query.explain())
SELECT [(col("A") + col("B")), (col("B") + col("B"))]
  DF ["A", "B", "C", "X", ...]; PROJECT["A", "B"] 2/5 COLUMNS
query = ldf.select(e3)
print(query.explain())
SELECT [(col("A") + col("B")), (col("B") + col("B")), (col("C") + col("B")), (col("X") + col("B")), (col("Y") + col("B"))]
  DF ["A", "B", "C", "X", ...]; PROJECT */5 COLUMNS

式の中に複数の列セレクタが現れると、Polars は位置で 1対 1に対応させます。左の 1列目と右の 1列目、という順です。結果の列名は左側のセレクタの列名になります。次のコードは pl.col('A', 'B') - pl.col('B', 'A') で A - B と B - A を計算し、A、B の 2列を得ます。長さが違うときは短い側がブロードキャストされ、例えば pl.col('A', 'B') - pl.col('B') は同じ B 列で A と B をそれぞれ引きます。

query = ldf.select(pl.col('A', 'B') - pl.col('B', 'A'))
print(query.explain())
SELECT [(col("A") - col("B")), (col("B") - col("A"))]
  DF ["A", "B", "C", "X", ...]; PROJECT["A", "B"] 2/5 COLUMNS

この挙動については、GitHub に次の提案があります。

現時点では、次のように pl.struct() で複数の列を1つの struct 列にまとめてから計算すれば、列同士の演算ができます。

df.with_columns(
    (pl.struct('A', 'B') - pl.struct('B', 'A'))
       .struct.rename_fields(['AB', 'BA'])
       .struct.field('*')
)
shape: (2, 7)
ABCXYABBA
i64i64i64i64i64i64i64
11010012-99
22020012-1818

式の出力名#

DataFrame に列を追加したり、列の計算結果を取り出したりするときは、結果に名前を付けられます。この名前を「式の出力名」と呼びます。出力名を指定しておくと、結果を後から参照したり、他の操作で使ったりしやすくなります。

ルート名と出力名#

出力名を指定しない場合、式で最初に現れる列の名前が出力名になります。式が参照している列名の一覧は「ルート名」と呼ばれ、.meta.root_names() と .meta.output_name() で確認できます。

def show_names(expr):
    print(f'root_names={expr.meta.root_names()}, output_name={expr.meta.output_name()}')
expr1 = pl.col("A") * pl.col("B") + 1
show_names(expr1)
root_names=['A', 'B'], output_name=A
expr2 = 1 + pl.col("B") * pl.col("A")
show_names(expr2)
root_names=['B', 'A'], output_name=literal
expr3 = (pl.col("A") * pl.col("B") + 1).alias("Z")

output_name が literal になるのは、式の先頭が数値など、列を参照しない項のときです。.name.keep() を使うと、ルート名の最初の列名に戻せます。

row(df.select(expr1), df.select(expr2), df.select(expr2.name.keep()), df.select(expr3))
shape: (2, 1)
A
i64
11
41
shape: (2, 1)
literal
i64
11
41
shape: (2, 1)
B
i64
11
41
shape: (2, 1)
Z
i64
11
41

複数の式が同じ出力名になると、計算時にエラーになります。次の 10 + pl.col('A', 'B') は A と B の両方について計算しますが、左端が数値であるため出力名がどちらも literal になり、衝突します。

%%capture_except
expr = 10 + pl.col('A', 'B')
df.select(expr)
DuplicateError: projections contained duplicate output name 'literal'. It's possible that multiple expressions are returning the same default column name. If this is the case, try renaming the columns with `.alias("new_name")` to avoid duplicate column names.

列の順番を入れ替えて列を先頭にするか、次のように .name.keep() で出力名を root_names[0] に戻します。

df.select(expr.name.keep())
shape: (2, 2)
AB
i64i64
1120
1230

.name.suffix() や .name.prefix() を使えば、ルート名をそのまま残した別名を付けられます。数値が左端にある式では名前が literal にそろうので、列の側を左端にした式に対して使います。

# 列を左端にすると、A_calc / B_calc という名前になる
df.select((pl.col("A", "B") + 10).name.suffix("_calc"))
shape: (2, 2)
A_calcB_calc
i64i64
1120
1230

ルート名の順番を変更#

次の pl.when().then().otherwise() の式では、root_names に並ぶ列名の順序は then()、otherwise()、when() の順になります。たとえば、次のような式では、expr.meta.root_names は ['x', 'y', 'j'] となります。

expr = pl.when(pl.col('j') > 50).then('x').otherwise('y')
print(expr.meta.root_names())
['x', 'y', 'j']

when() の中に複数の列がある場合、すべての式の出力名が x になるため、計算時に名前の重複で DuplicateError になります。

df2 = pl.DataFrame(dict(
    j=np.random.randint(10, 99, 2),
    k=np.random.randint(10, 99, 2),
    l=np.random.randint(10, 99, 2),
    x=10,
    y=20,
))
%%capture_except
cols = pl.col(['j', 'k', 'l'])
expr = pl.when(cols > 50).then('x').otherwise('y')
df2.select(expr)
DuplicateError: projections contained duplicate output name 'x'. It's possible that multiple expressions are returning the same default column name. If this is the case, try renaming the columns with `.alias("new_name")` to avoid duplicate column names.

この場合は、計算結果に影響せず、ルート名の順序だけを変える式を先頭に足すと解決できます。

df2.select(
    cols * 0 + pl.when(cols > 50).then('x').otherwise('y')
)
shape: (2, 3)
jkl
i32i32i32
102020
202010

出力名を指定する#

alias()#

alias() を使うと、計算結果に名前を付けられます。

result = df.select((pl.col('A') + pl.col('B')).alias("A_plus_B"))
print(result.columns)
['A_plus_B']

キーワード引数#

select() や with_columns() に新しい列を渡すとき、キーワード引数を使うとその列に名前を付けられます。

result = df.select(A_plus_B=pl.col('A') + pl.col('B'))
print(result.columns)
['A_plus_B']

.name の下には、出力名を操作するメソッドがいくつかあります。よく使う3つを紹介します。

map

ルート名に関数を適用して式の出力名を変更します。

result_df = df.select(
    (pl.col('A', 'B', 'C') + pl.col('B')).name.map(lambda name: name * 3)
)
print(result_df.columns)
['AAA', 'BBB', 'CCC']

prefix と suffix

ルート名の先頭と末尾に接頭辞・接尾辞を付けます。

result_df = df.select(
    (pl.col('A', 'B', 'C') + pl.col('B')).name.suffix("_res")
)
print(result_df.columns)
['A_res', 'B_res', 'C_res']

式の出力値#

演算式は「1つの列」に対して働きますが、その出力は 1行分だったり、元の列と同じ長さだったり、任意の長さだったりします。ここを押さえておくと、over() や agg() の一見奇妙な結果も説明できます。

式は独立して実行される#

Polars の式が独立して実行されることを、次のプログラムで確認します。

df = pl.DataFrame(
    dict(x=range(10), y=range(10))
)

result = df.select(pl.all().sample(4))
result
shape: (4, 2)
xy
i64i64
62
07
96
79

このプログラムでは、pl.all() で df のすべての列に sample(4) を適用しています。式は列ごとに独立して評価されるため、それぞれの列から別々に4行が選ばれ、選ばれる行は列によって一致しません。

seed を指定すると、すべての列で同じ乱数の系列が使われます。列ごとに独立して計算しても、同じ行を取り出せます。

df.select(pl.all().sample(4, seed=10))
shape: (4, 2)
xy
i64i64
11
66
00
77

なお、ランダムな行を取り出すだけなら DataFrame.sample() を使うのがもっとも簡単です。

df.sample(4)
shape: (4, 2)
xy
i64i64
22
99
44
33

.sort() は各列の値を独立に並べ替えます。

df = pl.DataFrame(
    dict(x=[1, 3, 2], y=[9, 8, 7])
)

複数の列を対象とする式は、列ごとの式に展開されてから独立に実行されます。次の例では x 列と y 列をそれぞれ並べ替えます。

df.select(pl.col('x', 'y').sort())
shape: (3, 2)
xy
i64i64
17
28
39

行の対応を保ったまま並べ替えるには、DataFrame.sort() を使います。

row(df.sort('x', 'y'), df.sort('y', 'x'))
shape: (3, 2)
xy
i64i64
19
27
38
shape: (3, 2)
xy
i64i64
27
38
19

式で同じことをするには、struct を使います。

df.select(
    pl.struct('y', 'x').sort().struct.field('*')
)
shape: (3, 2)
yx
i64i64
72
83
91

sort_by() を使う方法もあります。並べ替えの基準になる列を指定すれば、各列が独立に計算されても、行の対応は保たれます。

df.select(
    pl.all().sort_by("y", "x")
)
shape: (3, 2)
xy
i64i64
27
38
19

式の出力データ数#

長さ N の入力に対して、演算式の出力長は1、N、X の3通りです。

  • 長さ 1:集約系の演算式で、列全体が 1つの値にまとまります。

  • 長さ N:要素ごとの演算で、行数は変わりません。

  • 長さ X:フィルタやサンプリングなど、行数がデータ次第で決まるものです。

同じ select() の中で長さの違うものを混ぜるとエラーになるので、「集約結果と元の列を並べたい」ときは over() を使うか、いったん集約してから結合します。

それぞれの長さを、例で確認します。

出力長さが1#

出力長さが1の場合、式はデータフレーム全体に対して集計操作を行い、単一の結果を返します。例えば、列の合計や平均などです。次のコードでは、x 列の合計が計算され、単一の値 45 が返されます。

df = pl.DataFrame({
    "x": range(5),
})

result = df.select(pl.col("x").sum())
result
shape: (1, 1)
x
i64
10

出力長さが N#

出力長さが N の場合、式は各行に対して独立に操作を行い、元のデータフレームと同じ長さの結果を返します。たとえば、各要素に対するスカラー演算です。次のコードでは、x 列の各要素に 1 が加えられ、元のデータフレームと同じ長さの結果が返されます。

result = df.select(pl.col("x") + 1)
result
shape: (5, 1)
x
i64
1
2
3
4
5

出力長さが X の場合、特定の操作により元の長さ N と異なる長さの結果が返されます。たとえば、サンプリングやフィルタリングです。次のコードでは、x 列の3以上の値を取り出します。

result = df.select(pl.col('x').filter(pl.col('x') >=3))
result
shape: (2, 1)
x
i64
3
4

異なる長さの出力がある場合の動作#

長さが1の列は、その値が繰り返されて他の列と同じ長さに拡張されます。これにより、データフレームの行数を揃えられます。

df = pl.DataFrame(
    dict(x=[1, 3, 2], y=[9, 8, 7])
)

df.select(
    x_mean = pl.col('x').mean(),
    x_plus_y = pl.col('x') + pl.col('y')
)
shape: (3, 2)
x_meanx_plus_y
f64i64
2.010
2.011
2.09

次の2つの式の出力長さは1と2です。

df.select(
    x_mean = pl.col('x').mean(),
    x_top2 = pl.col('x').top_k(2)
)
shape: (2, 2)
x_meanx_top2
f64i64
2.02
2.03

出力長さの異なる複数の列を含む操作を行うと、Polars はエラーになります。

%%capture_except
df.select(
    x_sq = pl.col('x')**2,
    x_top2 = pl.col('x').top_k(2)
)
ShapeError: Series length 2 doesn't match the DataFrame height of 3

同じ式でも、列の中身によって出力長さは変わります。次の式では x 列の出力長さが2、y 列が3になるため、エラーになります。

%%capture_except
df.select(
    pl.col('x', 'y').filter(pl.col('x', 'y') >= 2)
)
ShapeError: Series length 2 doesn't match the DataFrame height of 3

Expr.implode() で複数の値を1つのリストにまとめれば、出力長さがすべて1になり、エラーを防げます。

df.select(
    pl.col('x', 'y').filter(pl.col('x', 'y') >= 2).implode()
)
shape: (1, 2)
xy
list[i64]list[i64]
[3, 2][9, 8, 7]

式のコンテキスト#

同じ演算式でも、置かれるメソッドによって渡されるデータが変わります。これを「コンテキスト」と呼びます。

select() と with_columns() は列全体を、group_by().agg() はグループごとの列を、.list.eval() は各リストを演算式に渡します。書き始める前に「いま自分はどのコンテキストにいるのか」を意識しておくと、後からデバッグする手間が減ります。

まず、次のデータフレームを用意します。この後のコンテキストの説明では、このデータフレームを使います。

df = pl.DataFrame(dict(x=[1, 2, 1, 2, 3, 3], y=[10, 20, 30, 40, 50, 60]))

agg コンテキスト#

次の .agg() の中の式は、x 列で分けられた3つのグループそれぞれのデータを処理します。出力が複数の値になる場合はリストにまとめられます。

df.group_by('x').agg(
    y = pl.col('y'),
    y_mult_x = pl.col('y') * pl.col('x'),
    sum_y_mult_x = (pl.col('y') * pl.col('x')).sum()
)
shape: (3, 4)
xyy_mult_xsum_y_mult_x
i64list[i64]list[i64]i64
1[10, 30][10, 30]40
2[20, 40][40, 80]120
3[50, 60][150, 180]330

list コンテキスト#

.list.eval() の式では、処理対象のリストが1つずつ渡されます。次のプログラムでは、x 列の各リストの平均値を計算し、その値を各要素から引きます。

  • 式❶ では、list.mean() で各リストの平均値を計算し、その平均値を各要素から引きます。

  • 式❷ では、list.eval() で x 列の各リストを処理し、式❸を適用します。

  • 式❸ では、pl.element() が現在処理中のリストの各要素を表し、その要素からリスト全体の平均値を引きます。

  • 式❹ では、list.agg() が x 列の各リストに対して集約式を実行します。結果の列はリスト型ではなく、スカラー値になります。この例では、各リストの最大値と最小値の差を計算します。

df = pl.DataFrame(dict(x=[[1, 2], [3, 4, 5, 6], [7, 8, 9]]))
df.with_columns(
    x1 = pl.col('x') - pl.col('x').list.mean(), #❶
    x2 = pl.col('x').list.eval(                 #❷
        pl.element() - pl.element().mean()      #❸
    ),
    x3 = pl.col('x').list.agg(                  #❹
        pl.element().max() - pl.element().min()
    )
)
shape: (3, 4)
xx1x2x3
list[i64]list[f64]list[f64]i64
[1, 2][-0.5, 0.5][-0.5, 0.5]1
[3, 4, … 6][-1.5, -0.5, … 1.5][-1.5, -0.5, … 1.5]3
[7, 8, 9][-1.0, 0.0, 1.0][-1.0, 0.0, 1.0]2

次のコードは、各リストの要素の中から、そのリストの平均値より大きい値を取り出します。この処理は list.eval() を使うと簡潔に書けます。使わない場合は、リスト全体に対する要素ごとの条件処理を書くのが難しくなります。

  • 式❺ では、pl.element() がリスト内の各要素を表します。

  • pl.element().mean() は現在処理中のリスト全体の平均値を計算します。

  • filter メソッドにより、各要素がリストの平均値より大きいかどうかを判定し、条件を満たす要素だけを取り出します。

df.with_columns(
    x_large = pl.col('x').list.eval(
        pl.element().filter(pl.element() > pl.element().mean()) #❺
    )    
)
shape: (3, 2)
xx_large
list[i64]list[i64]
[1, 2][2]
[3, 4, … 6][5, 6]
[7, 8, 9][9]

式❺の pl.element() を具体例に置き換えると、この計算の仕組みが分かりやすくなります。たとえば、リスト [7, 8, 9] に対して次のように計算します。

s = pl.Series([7, 8, 9])
s.filter(s > s.mean()).to_list()
[9]
  • この場合、s.mean() はリストの平均値を計算します (8.0)。

  • s > s.mean() は各要素が平均値より大きいかどうかを判定し、ブール値の配列 [False, False, True] を作ります。

  • filter メソッドを使うと、このブール条件を満たす要素(9)だけが抽出されます。

field コンテキスト#

.struct.with_fields() の中の式は、struct 列の各フィールドに対して計算を行います。

df = pl.DataFrame(
    {
        "coords": [{"x": 1, "y": 4}, {"x": 4, "y": 9}, {"x": 9, "y": 16}],
        "multiply": [10, 2, 3],
    }
)
df
shape: (3, 2)
coordsmultiply
struct[2]i64
{1,4}10
{4,9}2
{9,16}3

次の .struct.with_fields() の中の式では、pl.field(name) が struct 列の name フィールドの値を表します。

df.select(
    pl.col('coords').struct.with_fields(
        dist = (pl.field('x') ** 2 + pl.field('y') ** 2).sqrt() * pl.col('multiply')      
    ).struct.field('*')
)
shape: (3, 3)
xydist
i64i64f64
1441.231056
4919.697716
91655.072679

よく使う演算式#

ここでは、よく使う式をまとめて説明します。演算子、条件分岐、値の取り出し、連結、並べ替え、横方向の演算のいずれも、1〜2個のメソッドの組み合わせです。一度目を通して全体像をつかみ、必要になったときに戻ってくる使い方でかまいません。

演算子#

+、-、*、/ などの演算子には、対応する演算式のメソッドも用意されています。たとえば、次のコードでは + 演算子と add() メソッドで x 列と y 列の和を求めています。>= 演算子と ge() メソッドでは、列と数値を比較しています。

df = pl.DataFrame(dict(
    x=[1, 2, 3, 4],
    y=[40, 30, 20, 10]
))

df.with_columns(
    add1=pl.col('x') + pl.col('y'),
    add2=pl.col('x').add(pl.col('y')),
    cond1=pl.col('x') >= 3,
    cond2=pl.col('x').ge(3)
)
shape: (4, 6)
xyadd1add2cond1cond2
i64i64i64i64boolbool
1404141falsefalse
2303232falsefalse
3202323truetrue
4101414truetrue

& と | には、それぞれ and_ と or_ メソッドが対応します。次のコードでは、| 演算子と or_() メソッドで2つの条件の論理和(OR)を計算しています。

df.with_columns(
    cond1=(pl.col('x') > 3) | (pl.col('y') < 10),
    cond2=pl.col('x').gt(3).or_(pl.col('y').lt(10))
)
shape: (4, 4)
xycond1cond2
i64i64boolbool
140falsefalse
230falsefalse
320falsefalse
410truetrue

値を演算式にラップ#

pl.lit() は値を演算式にラップするための関数です。

Series や NumPy 配列を渡した場合は複数の値を持つ列として解釈され、その長さが結果の行数になります。一方、それ以外のデータ型はすべてスカラー値として扱われ、結果の行数に合わせて同じ値が繰り返されます。

そのため、次の例では Series と NumPy 配列の長さである3行のデータフレームが生成され、int、str、dict は各行で同じ値が繰り返されています。

また、Python の list や tuple は「列方向の複数値」ではなく、「1要素としてのリスト(list 型)」と解釈されるため、各行に同じリストが入ります。dict は struct 型として扱われる点にも注意が必要です。

この結果、series と array は行方向に展開され、list と tuple は「各行に同じリストを持つ列」として保持される、という違いが表に現れます。

pl.select(
    pl.lit(1).alias("int"),
    pl.lit("hello").alias("str"),
    pl.lit(dict(x=1, y=2)).alias("dict"),
    pl.lit([1, 2, 3]).alias("list"),
    pl.lit((1, 2, 3)).alias("tuple"),
    pl.lit(pl.Series([1, 2, 3])).alias("series"),
    pl.lit(np.array([1, 2, 3])).alias("array"),
)
shape: (3, 7)
intstrdictlisttupleseriesarray
i32strstruct[2]list[i64]list[i64]i64i64
1"hello"{1,2}[1, 2, 3][1, 2, 3]11
1"hello"{1,2}[1, 2, 3][1, 2, 3]22
1"hello"{1,2}[1, 2, 3][1, 2, 3]33

1つの pl.lit() から複数の列を作りたい場合は、主に次の2つの方法があります。

  • 辞書のリストを使う方法 辞書のリストを構造体(struct)型の Series に変換してから pl.lit() に渡します。その後、.struct.unnest() を呼び出すことで、構造体の各フィールドを独立した列として展開できます。

  • リストのリストを使う方法 リストのリストを list 型の Series に変換し、.list.to_struct(fields=...) で構造体の列に変換します。最後に .struct.unnest() を使うことで、複数の列を一度に作成できます。

このように struct 型を経由して .struct.unnest() を使えば、pl.lit() から複数の列を作れます。外部データや計算結果をまとめて列として追加したいときは、この形が便利です。

pl.select(
    # 辞書のリスト → struct Series → 複数列
    pl.lit(
        pl.Series(
            [{"x": 1, "y": 2}, {"x": 2, "y": 3}, {"x": 3, "y": 4}]
        )
    )
    .struct.unnest(),

    # リストのリスト → list Series → struct Series → 複数列
    pl.lit(
        pl.Series([[1, 2], [2, 3], [3, 4]])
    )
    .list.to_struct(fields=["a", "b"])
    .struct.unnest(),
)
shape: (3, 4)
xyab
i64i64i64i64
1212
2323
3434

数学関数#

sin や sqrt といった単項関数は、対応するメソッドで計算します。一方、arctan2 のような二項関数は、グローバル関数で計算します。

df.with_columns(
    sin=pl.col('x').sin(),
    sqrt=pl.col('y').sqrt(),
    arctan=pl.arctan2(pl.col('y'), pl.col('x')),
)
shape: (4, 5)
xysinsqrtarctan
i64i64f64f64f64
1400.8414716.3245551.545802
2300.9092975.4772261.504228
3200.141124.4721361.421906
410-0.7568023.1622781.19029

条件分岐は、pl.when(cond) を起点とし、その後に .then() や .when()、.otherwise() などのメソッドを組み合わせて表現します。たとえば、次のコードはこの条件分岐を式で表したものです。

if x > 4:
    result = 'very good'
elif x > 3:
    result = 'good'
elif x > 2:
    result = 'fair'
elif x > 1:
    result = 'bad'
else:
    result = 'very bad'
df.with_columns(
    pl.when(pl.col('x') > 4)
      .then(pl.lit('very good'))
      .when(pl.col('x') > 3)
      .then(pl.lit('good'))
      .when(pl.col('x') > 2)    
      .then(pl.lit('fair'))
      .when(pl.col('x') > 1)    
      .then(pl.lit('bad'))    
      .otherwise(pl.lit('very bad'))
      .alias('level')
)
shape: (4, 3)
xylevel
i64i64str
140"very bad"
230"bad"
320"fair"
410"good"

when-then-otherwise は、条件を満たす行についてだけ then の枝を計算し、満たさない行は otherwise を通ります。次のコードは "D" で始まる文字列の後ろを整数に変換します。"HAB" や "HFF" は条件を満たさず then の枝に入らないので、変換されずに null になります。

逆に、条件が変換できない行まで覆う場合、たとえば条件を pl.lit(True) にすると、同じ式が ComputeError: strict integer parsing failed で止まります。

df_s = pl.DataFrame(dict(val=["D12", "D32", "HAB", "HFF", "D100"]))

df_s.select(
    pl.when(pl.col("val").str.starts_with("D"))
    .then(pl.col("val").str.slice(1).str.to_integer())
    .otherwise(None)
)
shape: (5, 1)
val
i64
12
32
null
null
100

str.to_integer() に strict=False を渡すと、変換できない値がエラーではなく null になります。条件が変換できない行まで覆う可能性があるときは、条件そのものに頼るより確実です。

df_s1 = df_s.select(
    pl.when(pl.col("val").str.starts_with("D"))
    .then(pl.col("val").str.slice(1).str.to_integer(strict=False))
    .otherwise(None)
)

あるいは、when-then-otherwise で条件を満たさない値を null にしてから、変換を行います。

df_s2 = df_s.select(
    pl.when(pl.col("val").str.starts_with("D"))
    .then("val")
    .otherwise(None)
    .str.slice(1)
    .str.to_integer()
)

row(df_s1, df_s2)
shape: (5, 1)
val
i64
12
32
null
null
100
shape: (5, 1)
val
i64
12
32
null
null
100

.cut() は、境界値でデータを区間に分けてラベルを付ける式です。Polars 2.0 では非推奨になり、bin_intervals() が同じ役割を果たします。

  • breaks(bin_intervals では第1引数)— N 個の境界値

  • labels — N+1個のラベル

  • 区間の閉じ方 — cut() の left_closed=False は右閉区間、bin_intervals() の既定(左閉)に対して right_closed=True が対応します

次のコードは x 列にラベルを付ける例です。

  • right_closed=True(cut(left_closed=False) 相当)— 右側が閉区間

    • x <= 1 → very bad

    • 1 < x <= 2 → bad

    • 2 < x <= 3 → fair

    • 3 < x <= 4 → good

    • 4 < x → very good

  • 既定(左閉)— 左側が閉区間

    • x < 1 → very bad

    • 1 <= x < 2 → bad

    • 2 <= x < 3 → fair

    • 3 <= x < 4 → good

    • 4 <= x → very good

labels = ["very bad", "bad", "fair", "good", "very good"]
df.with_columns(
    level_right=pl.col("x").bin_intervals([1, 2, 3, 4], labels=labels, right_closed=True),
    level_left=pl.col("x").bin_intervals([1, 2, 3, 4], labels=labels),
)
shape: (4, 4)
xylevel_rightlevel_left
i64i64enumenum
140"very bad""bad"
230"bad""fair"
320"fair""good"
410"good""very good"

.clip() は、指定した最小値や最大値を境界として、範囲を外れた値を境界値に置き換える関数です。次の例では、y 列をいくつかの境界条件でクリップしています。

  • y1 列: y 列の値を 20〜30の範囲内に収めます。

  • y2 列: 下限値のみを 30に設定し、それより小さい値を 30にクリップします。

  • y3 列: 上限値のみを 20に設定し、それより大きい値を 20にクリップします。

  • y4 列: 上限値と下限値を式で指定します。この場合、各要素ごとに境界値が計算され、それを基準にクリップされます。

df.with_columns(
    y1=pl.col('y').clip(20, 30),
    y2=pl.col('y').clip(lower_bound=30),
    y3=pl.col('y').clip(upper_bound=20),
    y4=pl.col('y').clip(pl.col('x') * 10, pl.col('x') * 20)
)
shape: (4, 6)
xyy1y2y3y4
i64i64i64i64i64i64
14030402020
23030302030
32020302030
41020301040

.replace() と .replace_strict() を使うと、要素を置き換えられます。置き換え前後の値は、2つのリストまたは1つの辞書として渡せます。両者の違いは次のとおりです。

  • .replace(): 一部の要素を置き換えます。結果は元のデータ型のままになります。

  • .replace_strict(): すべての要素を置き換えます。結果は置き換え後のデータ型に変わります。

次のコードでは、x 列の値は置き換え後に文字列型(x2 列)に変わり、y 列の値は元のデータ型のまま(y2 列)です。

df.with_columns(
    x2=pl.col('x').replace_strict({1:'very bad', 2:'bad', 3:'fair'}, default='unknown'),
    y2=pl.col('y').replace([10, 20], [100, 200]),
)
shape: (4, 4)
xyx2y2
i64i64stri64
140"very bad"40
230"bad"30
320"fair"200
410"unknown"100

データの選択#

式のメソッドを使うと、条件に応じた取り出しや加工ができます。先頭・末尾を取る、位置で取る、条件で絞る、といった操作がここに含まれます。

よく使うメソッドとその用途を表にまとめます。

メソッド

説明

.head(n)

最初の n 行を取得します。

.tail(n)

最後の n 行を取得します。

.slice(offset, n)

指定したオフセットから n 行を取得します。

.get(index)

指定したインデックスの値を取得します。

.gather(indices)

指定したインデックスリストに基づいて行を取得します。

.gather_every(n, offset=0)

指定したステップ間隔で行を取得します。

.filter(condition)

条件を満たす行を取得します。

.top_k(k)

値が上位 k 個の行を取得します。

.top_k_by(col, k)

指定した列に基づいて値が上位 k 個の行を取得します。

.bottom_k(k)

値が下位 k 個の行を取得します。

.bottom_k_by(col, k)

指定した列に基づいて値が下位 k 個の行を取得します。

.sample(n)

ランダムに n 行を取得します。

df_res = df.select(
    head=pl.col("x").head(2),                       # 最初の2行
    tail=pl.col("x").tail(2),                       # 最後の2行
    slice=pl.col("x").slice(1, 2),                  # オフセット1から2行
    get=pl.col("x").get(3),                         # インデックス3の値
    gather=pl.col("x").gather([0, 3]),              # インデックス0と3の行
    gather_every=pl.col("x").gather_every(2, offset=1), # 1オフセットで2ステップごと
    filter=pl.col("x").filter(pl.col("y") // 10 % 2 == 0), # 条件でフィルタ
    top_k=pl.col("x").top_k(2),                     # 上位2件
    top_k_by=pl.col("x").top_k_by(pl.col("y"), 2),  # y列に基づく上位2件
    bottom_k=pl.col("x").bottom_k(2),               # 下位2件
    bottom_k_by=pl.col("x").bottom_k_by(pl.col("y"), 2), # y列に基づく下位2件
    sample=pl.col("x").sample(2),                   # ランダム2行
)
row(df, df_res)
shape: (4, 2)
xy
i64i64
140
230
320
410
shape: (2, 12)
headtailslicegetgathergather_everyfiltertop_ktop_k_bybottom_kbottom_k_bysample
i64i64i64i64i64i64i64i64i64i64i64i64
132412141143
243444332234

Polars の多くのメソッドでは、単なる値だけでなく、演算式も引数として渡せます。たとえば .get() に式を渡すと、複数の要素をまとめて取り出せます。

次のコードでは、

  • pl.len() でデータフレームの行数を取得します。

  • (pl.col("x") + 1) % pl.len() で、列 x の各要素に1を足し、データフレームの行数で割った余りをインデックスとします。

  • .get() で上の式で計算したインデックスを使用し、要素を取得します。

df.with_columns(
    get=pl.col('y').get((pl.col('x') + 1) % pl.len())
)
shape: (4, 3)
xyget
i64i64i64
14020
23010
32040
41030

データの結合#

concat() や .append() を使用すると、複数の演算結果を縦方向に結合できます。一方、.extend_constant() を使用すると、定数を複数回追加できます。

次のプログラムは、x 列と y 列の一部または定数を結合する例を示しています。

df2 = df.select(
    append=pl.col('x').append(pl.col('y').slice(1, 2)),
    concat=pl.concat([pl.col('x'), pl.col('y').slice(1, 2)]),
    extend=pl.col('x').extend_constant(99, 2),
)
row(df, df2)
shape: (4, 2)
xy
i64i64
140
230
320
410
shape: (6, 3)
appendconcatextend
i64i64i64
111
222
333
444
303099
202099

concat_list()、concat_arr()、struct() を使って、複数の演算結果を横方向に結合します。それぞれ、List 型、Array 型、Struct 型の列を作成します。次のプログラムは、x 列、y 列、そして x 列と y 列の商を横方向に結合する例です。

  • List 型と Array 型では、すべての要素のデータ型が一致している必要があります。

    • concat_list() は自動的に上位のデータ型にキャストします。

    • concat_arr() を使用する場合は、手動で同じデータ型にキャストする必要があります。

  • Struct 型では、各フィールドが独自のデータ型を持てます。

df3 = df.select(
    list=pl.concat_list('x', pl.col('y') / pl.col('x')),
    arr=pl.concat_arr(pl.col('x').cast(pl.Float64), pl.col('y') / pl.col('x')),
    struct=pl.struct('x', pl.col('y') / pl.col('x')),
)

row(df, df3)
shape: (4, 2)
xy
i64i64
140
230
320
410
shape: (4, 3)
listarrstruct
list[f64]array[f64, 2]struct[2]
[1.0, 40.0][1.0, 40.0]{1,40.0}
[2.0, 15.0][2.0, 15.0]{2,15.0}
[3.0, 6.666667][3.0, 6.666667]{3,6.666667}
[4.0, 2.5][4.0, 2.5]{4,2.5}

.explode() を使用すると、List 型および Array 型の列の要素を行優先で縦方向に結合できます。また、.list.explode() と .arr.explode() は同じ処理を行います。

次のプログラムは、上記の結果で作成した list 列と arr 列の要素を縦方向に結合する例を示しています。Struct 型の列を行優先で結合する場合は、まず concat_list() で .struct.unnest() の結果を List 型に変換し、そのあと .explode() で縦方向に結合します。列優先で結合する場合は、.struct.unnest() の結果を concat() で結合します。

df4 = df3.select(
    pl.col('list').explode(), 
    pl.col('arr').explode(),
    struct_row=pl.concat_list(pl.col('struct').struct.unnest()).explode(),
    struct_col=pl.concat([pl.col('struct').struct.unnest()])
)
row(df3, df4)
shape: (4, 3)
listarrstruct
list[f64]array[f64, 2]struct[2]
[1.0, 40.0][1.0, 40.0]{1,40.0}
[2.0, 15.0][2.0, 15.0]{2,15.0}
[3.0, 6.666667][3.0, 6.666667]{3,6.666667}
[4.0, 2.5][4.0, 2.5]{4,2.5}
shape: (8, 4)
listarrstruct_rowstruct_col
f64f64f64f64
1.01.01.01.0
40.040.040.02.0
2.02.02.03.0
15.015.015.04.0
3.03.03.040.0
6.6666676.6666676.66666715.0
4.04.04.06.666667
2.52.52.52.5

順番替え#

順序に関わる操作(ずらし、反転、ソート)は、時系列処理やウィンドウ計算でよく使います。shift() は「1期前」の値を作り、sort() は集約の前後関係を決め、reverse() は累積計算と組み合わせることが多いメソッドです。

.shift() は、列の値を上向きまたは下向きに N 個ずらす関数です。N が正の数なら下向き、負の数なら上向きにシフトします。空いた位置は fill_value で指定した値で埋め、既定では NULL になります。循環シフトにしたい場合は、.append() で先頭または末尾の値を結合してからシフトします。

df1 = df.with_columns(
    x_shift1=pl.col("x").shift(1),
    x_shift2=pl.col("x").shift(-2),
)

df2 = df.select(
    x_shift1=pl.col('x').tail(1).append(pl.col('x')).shift(1).tail(pl.len()),
    x_shift2=pl.col('x').append(pl.col('x').head(2)).shift(-2).head(pl.len())
)

row(df1, df2)
shape: (4, 4)
xyx_shift1x_shift2
i64i64i64i64
140null3
23014
3202null
4103null
shape: (4, 2)
x_shift1x_shift2
i64i64
43
14
21
32

.reverse() は、列の値を逆向きに並べ替える関数です。

df.with_columns(x_reverse=pl.col('x').reverse())
shape: (4, 3)
xyx_reverse
i64i64i64
1404
2303
3202
4101

.sort() と .sort_by() は、列の値を大小関係で並べ替える関数です。descending=True なら降順、既定では昇順(False)になります。

  • .sort(): 指定した列を昇順または降順で並べ替えます。

  • .sort_by(): 並べ替えの基準となる演算式を指定します。複数指定することもできます。

次のコードでは、x 列を降順で並べ替えた結果を x_sort 列として追加し、y 列を昇順で並べ替えた結果を y_sort 列として追加しています。また、x と y 列を、y 列と 24の距離を基準に昇順で並べ替えた結果を x_sort_by と y_sort_by 列として追加します。

df.with_columns(
    pl.col('x').sort(descending=True).alias('x_sort'),
    pl.col('y').sort().alias('y_sort'),
    pl.all().sort_by((pl.col('y') - 24).abs()).name.suffix('_sort_by')
)
shape: (4, 6)
xyx_sorty_sortx_sort_byy_sort_by
i64i64i64i64i64i64
140410320
230320230
320230410
410140140

NULL 処理#

NULL に関する式は NULL 処理 で詳しく説明します。

横向き演算#

関数名に horizontal が含まれる関数は、複数の列間で横方向の演算を行います。次のコードは、これらの関数をすべて出力します。

for name in dir(pl):
    if 'horizontal' in name:
        print(name)
all_horizontal
any_horizontal
cum_sum_horizontal
max_horizontal
mean_horizontal
min_horizontal
sum_horizontal

次のコードは、tag 列以外のすべての列の総和を計算し、その結果を新しい列 sum として追加します。

df = pl.DataFrame(
    {
     'tag': ['a', 'b', 'a', 'b', 'c', 'e'],        
     'A': [5, 1, 4, 5, 4, 6],
     'B': [3, 1, 1, 6, 4, 9],
     'C': [5, 6, 3, 7, 0, 9],
     'D': [0, 8, 0, 1, 0, 5],
     'E': [7, 7, 6, 2, 2, 8],
    }
)

df.with_columns(
    pl.sum_horizontal(pl.exclude("tag")).alias("sum")
)
shape: (6, 7)
tagABCDEsum
stri64i64i64i64i64i64
"a"5350720
"b"1168723
"a"4130614
"b"5671221
"c"4400210
"e"6995837

horizontal 関数が使えない場合は、concat_list() で複数の列を1つのリスト列に変換し、そのあと list.xxx 関数で計算する方法もあります。次の例では、すべての列の中央値を計算し、新しい列 median として追加しています(tag 列を除く)。

df.with_columns(
    pl.concat_list(pl.exclude("tag"))
    .list.median()
    .alias("median")
)
shape: (6, 7)
tagABCDEmedian
stri64i64i64i64i64f64
"a"535075.0
"b"116876.0
"a"413063.0
"b"567125.0
"c"440022.0
"e"699588.0

list.xxx で表せない計算は list.eval() で書けます。次のコードは、各行の値のうち大きいほうから 2つを取り、その合計を sum_top_2 に入れます。

list.eval() の中の .sum() はリスト全体の合計を 1要素のリストとして返すので、.list.first() で取り出します。

df.with_columns(
    pl.concat_list(pl.exclude("tag"))
    .list.eval(pl.element().sort(descending=True).head(2).sum())
    .list.first()
    .alias("sum_top_2")
)
shape: (6, 7)
tagABCDEsum_top_2
stri64i64i64i64i64i64
"a"5350712
"b"1168715
"a"4130610
"b"5671213
"c"440028
"e"6995818