Polarsに関するTips#

Polarsは、データフレーム操作を効率的に行うための強力なライブラリですが、さらに効果的に活用するためにはいくつかの便利なテクニックやコツを知っておくことが重要です。この章では、Polarsをよりスムーズに、かつ効率的に使用するための実践的なTipsを紹介します。これらのテクニックを使うことで、データの操作や処理を一層高速化し、作業の生産性を向上させることができます。

import panel as pn
pn.extension()
import polars as pl
from helper.jupyter import row

条件分岐#

Polarsでは条件分岐を記述する際、pl.when().then().when().then().otherwise()のようなメソッドチェーンを使用します。しかし、この方法は記述が煩雑で、可読性も低くなりがちです。本書では、このようなメソッドチェーンを簡潔に記述するために、match()関数を利用します。

次のコード例では、rank_expr1pl.when()チェーンを使用してランク付けのロジックを記述しています。一方、rank_expr2ではmatch()関数を使用して、同じ処理をより簡潔に記述しています。match()関数の引数については:

  • match()奇数番目の引数には条件式を指定します。

  • 偶数番目の引数には条件が一致した場合の結果式を指定します。

  • 引数の数が奇数の場合、最後の引数はotherwise()に渡されるデフォルト値として扱われます。

from helper.polars import match

df = pl.DataFrame({
    "name": ["Alice", "Bob", "Charlie", "Diana"],
    "score": [95, 85, 70, 50],
})

rank_expr1 = (
    pl.when(pl.col("score").is_between(90, 100)).then(pl.lit("A"))
    .when(pl.col("score").is_between(80, 89)).then(pl.lit("B"))
    .when(pl.col("score").is_between(70, 79)).then(pl.lit("C"))
    .otherwise(pl.lit("F"))                                    
)

rank_expr2 = match(
    pl.col("score").is_between(90, 100), pl.lit("A"),
    pl.col("score").is_between(80, 89), pl.lit("B"),
    pl.col("score").is_between(70, 79), pl.lit("C"),
    pl.lit("F")
)

df2 = df.with_columns(
    rank_expr1.alias('rank1'),
    rank_expr2.alias("rank2")
)
df2
shape: (4, 4)
namescorerank1rank2
stri64strstr
"Alice"95"A""A"
"Bob"85"B""B"
"Charlie"70"C""C"
"Diana"50"F""F"

このように、match()を使うことで可読性が向上し、コードの冗長さを軽減できます。

メソッドチェーン#

Polarsでは、メソッドチェーンを活用することでデータ操作を効率的に行うことができますが、メソッドチェーンの途中でどのようにデータが変化するのかを把握するのは難しい場合があります。そこで、デバッグのために次の2つの方法を活用できます。

  • helper.polars.DataCapturer: メソッドチェーンの任意の点でデータをキャプチャし、状態を確認できます。

  • helper.polars.PipeLogger: メソッドチェーン全体の入出力と引数をキャプチャし、可視化します。

DataCapturer#

DataCapturerを使用すると、メソッドチェーン内の特定のポイントでデータの状態をキャプチャできます。以下は、DataCapturerの使い方を示す例です。

from helper.polars import DataCapturer

df = pl.DataFrame({
    "name": ["Alice", "Bob", "Charlie", "David", "Eve"],
    "age": [25, 30, 35, 40, 45],
    "score": [85, 90, 95, 80, 70],
    "department": ["HR", "IT", "HR", "IT", "HR"]
})

cap = DataCapturer() #❶

result = (
df
.filter(pl.col("age") > 30)  
.with_columns(
  (pl.col("score") * 1.1).alias("adjusted_score")
)
.pipe(cap.before_group) #❷
.group_by("department")
.agg([
  pl.col("adjusted_score").mean().alias("average_score"),
  pl.col("age").max().alias("max_age")
])
.pipe(cap.before_sort) #❷
.sort("average_score", descending=True)
)

row(cap.before_group, cap.before_sort) #❸
shape: (3, 5)
nameagescoredepartmentadjusted_score
stri64i64strf64
"Charlie"3595"HR"104.5
"David"4080"IT"88.0
"Eve"4570"HR"77.0
shape: (2, 3)
departmentaverage_scoremax_age
strf64i64
"HR"90.7545
"IT"88.040

このコードでは、cap.before_groupでグループ化の前のデータ状態をキャプチャし、cap.before_sortでソート前のデータ状態をキャプチャしています。

❶メソッドチェーンを開始する前に、DataCapturer()のインスタンスcapを作成します。
❷メソッドチェーンの特定の場所で、.pipe(cap.name)を挿入し、その時点のデータをnameとしてキャプチャします。
❸キャプチャしたデータはcap.nameでアクセスできます。

PipeLogger#

PipeLoggerを使用すると、メソッドチェーン内の各ステップで発生するデータの入出力を視覚的に確認できます。これにより、データがどのように変化しているかを逐一確認しながら処理を進めることができ、デバッグが容易になります。

以下のコード例では、PipeLoggerを使用してデータフレームを操作し、処理中の入出力をJupyterLabで可視化する方法を示します。

from helper.polars import PipeLogger

result = (
PipeLogger(df) #❶
.filter(pl.col("age") > 30)
.with_columns(
  (pl.col("score") * 1.1).alias("adjusted_score")
)
.group_by("department")
.agg([
  pl.col("adjusted_score").mean().alias("average_score"),
  pl.col("age").max().alias("max_age")
])
.sort("average_score", descending=True)
)

result #❷

PipeLoggerを使うことで、メソッドチェーン内のデータの変化を視覚的に追跡し、デバッグがしやすくなります。

❶最初のオブジェクトであるdfPipeLogger()でラップします。これにより、結果はDataFrameではなく、すべての入出力をキャプチャするPipeLoggerオブジェクトになります。PipeLoggerは、メソッドチェーンの各ステップでデータの状態を保存します。

❷最後に、PipeLoggerオブジェクトをセルの最後に配置し、インタラクティブなウィジェットとして結果を表示します。これにより、各メソッドの入出力を可視化することができます。

ExprCapturer#

演算式のメソッドチェーンに流れているデータを観察したい場合は、Expr.inspect()を使用することができます。このメソッドは、データをそのまま出力します。例えば、

df.select(
    pl.col('name', 'department').str.slice(0, 1).str.to_lowercase().value_counts().inspect().implode()
);
shape: (5,)
Series: 'name' [struct[2]]
[
	{"e",1}
	{"a",1}
	{"b",1}
	{"c",1}
	{"d",1}
]shape: (2,)
Series: 'department' [struct[2]]
[
	{"h",3}
	{"i",2}
]

しかし、演算式のチェーン内を流れるデータをキャプチャして、後で解析したい場合には、本書が提供するExprCapturerを使用する方法をおすすめします。以下のコードでは、ExprCapturerのインスタンスcapを作成し、.pipe(cap.counts)を使用してデータをcountsという名前でキャプチャしています:

from helper.polars import ExprCapturer

cap = ExprCapturer()
df_res = df.select(
    pl.col('name', 'department').str.slice(0, 1).str.to_lowercase().value_counts().pipe(cap.counts).implode()
)
df_res
shape: (1, 2)
namedepartment
list[struct[2]]list[struct[2]]
[{"d",1}, {"e",1}, … {"c",1}][{"h",3}, {"i",2}]

複数の列に対して計算を行うため、キャプチャしたデータは複数のSeriesオブジェクトのリストとして保存されます。例えば、次のようにしてキャプチャしたデータを確認できます:

row(cap.counts[0], cap.counts[1])
shape: (5,)
name
struct[2]
{"d",1}
{"e",1}
{"b",1}
{"a",1}
{"c",1}
shape: (2,)
department
struct[2]
{"h",3}
{"i",2}

複雑な演算式#

同じ中間結果を複数回利用する場合、メソッドチェーンのメリットが薄れることがあります。このような場合は、中間結果を変数に保存して後で利用するのが一般的です。たとえば、3次元ベクトルを正規化(normalize)するとき、ベクトルの長さを複数回使用します。一見すると、lengthの式が複数回実行されるように見えますが、Polarsでは演算結果をキャッシュするため、同じ計算を2回実行することはありません。

df = pl.DataFrame(
    dict(x=[1.0, 2.0, 3.0], y=[3.0, 2.0, 1.0], z=[2.0, 1.0, 3.0])
)

x, y, z = pl.col('x'), pl.col('y'), pl.col('z')
length = (x**2 + y**2 + z**2).sqrt()

df1 = df.select(
    x = x / length,
    y = y / length,
    z = z / length
)

row(df, df1)
shape: (3, 3)
xyz
f64f64f64
1.03.02.0
2.02.01.0
3.01.03.0
shape: (3, 3)
xyz
f64f64f64
0.2672610.8017840.534522
0.6666670.6666670.333333
0.6882470.2294160.688247

次のコードでは、LazyDataFrame.explain()を使用して実行プランを表示します。この結果、ベクトルの長さを計算する式が__POLARS_CSER_0x30925ccb05afdfe7という名前の列にキャッシュされていることが分かります。

print(
    df
    .lazy()
    .with_columns(
        x = x / length,
        y = y / length,
        z = z / length
    )
    .explain()
)
simple π 3/4 ["x", "y", "z"]
   WITH_COLUMNS:
   [[(col("x")) / (col("__POLARS_CSER_0x30925ccb05afdfe7"))].alias("x"), [(col("y")) / (col("__POLARS_CSER_0x30925ccb05afdfe7"))].alias("y"), [(col("z")) / (col("__POLARS_CSER_0x30925ccb05afdfe7"))].alias("z")] 
     WITH_COLUMNS:
     [[([(col("x").pow([dyn int: 2])) + (col("y").pow([dyn int: 2]))]) + (col("z").pow([dyn int: 2]))].sqrt().alias("__POLARS_CSER_0x30925ccb05afdfe7")] 
      DF ["x", "y", "z"]; PROJECT */3 COLUMNS; SELECTION: None

直接演算式で複雑な計算を書くのは非常に手間がかかります。例えば、平方根を求める場合、sqrt(x)のように記述することはできず、x.sqrt()と書く必要があります。より数学的な記法で演算式を記述できるようにするため、本書では次のpolars_exprsデコレータを提供します。このデコレータを使用することで、通常の関数をPolarsの演算式に変換することができます。この関数内では、演算式の利点をそのまま活用することが可能です。また、.list.first()のようなサブネームスペースの操作も、list_first()のような関数として利用できます。さらに、c_で始まる変数は列を表す演算式として解釈されます。例えば、c_xpl.col('x')と同じ意味を持ちます。

from helper.polars import polars_exprs

@polars_exprs
def norm():
    length = sqrt(c_x**2 + c_y**2 + c_z**2)
    return dict(x=x / length, y=y / length, z=z / length)

df2 = df.select(**norm())
row(df, df2)
shape: (3, 3)
xyz
f64f64f64
1.03.02.0
2.02.01.0
3.01.03.0
shape: (3, 3)
xyz
f64f64f64
0.2672610.8017840.534522
0.6666670.6666670.333333
0.6882470.2294160.688247

c_で始まる変数名を使用しない場合、キーワード引数を用いて列名を指定することができます。例えば、次のコードではその例を示しています:

@polars_exprs
def norm2():
    length = sqrt(x**2 + y**2 + z**2)
    return dict(x=x / length, y=y / length, z=z / length)

df = pl.DataFrame(
    dict(a=[1.0, 2.0, 3.0], b=[3.0, 2.0, 1.0], c=[2.0, 1.0, 3.0])
)
df1 = df.select(**norm2(x='a', y='b', z='c'))

row(df, df1)
shape: (3, 3)
abc
f64f64f64
1.03.02.0
2.02.01.0
3.01.03.0
shape: (3, 3)
xyz
f64f64f64
0.2672610.8017840.534522
0.6666670.6666670.333333
0.6882470.2294160.688247

polars_exprsでデコレートされた関数が実行される際、以下の順序でシンボルが解釈されます。

  1. キーワード引数: キーワード引数に指定されたシンボルは、そのままキーワード引数として解釈されます。引数の値が文字列の場合、自動的にpl.col()を用いて列を表す式に置き換えられます。

  2. c_で始まるシンボル: c_で始まるシンボルはpl.col()に変換されます。例えば、c_namepl.col('name')となります。

  3. Polarsライブラリ内の演算式関連の関数やメソッド

  4. 関数定義時のグローバル変数

  5. Pythonのビルトイン関数

以下の例では、それぞれのシンボルが次のように解釈されます:

  • Ap:キーワード引数で指定した列。pl.col('Amp')およびpl.col('p')に対応します。

  • c_f:列pl.col('f')に変換されます。

  • math:グローバル変数として解釈されます。

  • sin():Polarsの演算式メソッドsin()を使用します。

  • print():Pythonのビルトイン関数print()を使用します。

import math

@polars_exprs
def my_expr():
    phase = 2 * math.pi * c_f + p
    print(phase)
    value = A * sin(phase)
    return value

df = pl.DataFrame(dict(
    Amp=[1, 2, 3],
    f=[10, 20, 30],
    p=[0.1, 0.2, 0.3]
))

df.select(my_expr(A="Amp", p="p"))
[([(dyn float: 6.283185) * (col("f"))]) + (col("p"))]
shape: (3, 1)
Amp
f64
0.099833
0.397339
0.886561

polars_exprsデコレータを使用することで得られる主なメリットは以下の通りです:

  1. より直感的な記法で計算式を記述できる

    • 通常のPython関数のように計算式を書けるため、コードが簡潔で分かりやすくなります。

    • 例えば、sqrt(x)のような数学的な記法で演算式を記述可能です(通常のPolarsではx.sqrt()と記述する必要があります)。

  2. 列参照が簡単になる

    • c_で始まる変数(例: c_x)を使えば、pl.col('x')のように書かずに列を参照できます。

    • キーワード引数を使うことで、柔軟に列名を指定可能です。例えば、x='column1'のように明示的に列を指定できます。

  3. Polars演算式のメリットをそのまま活用

    • 関数内で使用される演算はPolarsの遅延評価(lazy execution)や高速なクエリ実行エンジンの恩恵を受けられるため、大量データの処理でも効率的です。

    • 演算式の関数やメソッド(例: sin()sqrt())が直接使えるため、記述が直感的になります。

  4. サブネームスペースの関数が簡単に使える

    • Polarsのサブネームスペース(例: .list.first())に対応したカスタム関数(例: list_first())を使えるため、柔軟性が向上します。

  5. コードの再利用性が向上

    • 演算式をPython関数として抽象化できるため、複雑な計算や処理を繰り返し使用する場合にコードを再利用しやすくなります。

    • 同じ計算式を異なる列に適用したい場合も、キーワード引数を用いて簡単に対応可能です。

  6. Polarsの機能とPythonの標準機能の統合

    • Polars演算式に加え、標準Pythonのビルトイン関数(例: print())やグローバル変数(例: math.pi)も利用できるため、柔軟な処理が可能です。

  7. 読みやすく保守性の高いコード

    • 数学的な記法や柔軟な列指定によって、コードが読みやすくなるため、保守性が向上します。

    • チームでの共有や将来的な修正が容易になります。

配列を処理するユーザー関数#

次のコードは、a 列が b 列より小さい値を指定された値に置き換えます。Polars の演算式でこの計算を実装するのは難しいため、pyarrownumpy の機能を使用すれば簡単に実現でき、map_batches() を活用しています。

import numpy as np

df = pl.DataFrame({
    "a":[1, 2, 3, 10],
    "b":[1, 3, 4, 5],
})

def func(args, values):
    a, b = [s.to_numpy() for s in args]
    c = np.copy(a)
    c[a < b] = values
    return pl.Series(c)

df.with_columns(
    pl_result=pl.map_batches(['a', 'b'], lambda args:func(args, [100, 200]))
)    
shape: (4, 3)
abpl_result
i64i64i64
111
23100
34200
10510

map_batches() には次の三つの問題があります:

  1. 複数の列を一つのリストにまとめてユーザー関数に渡します。

  2. ユーザー関数に他の引数を渡すことができません。

  3. ユーザー関数の入力および出力が Series オブジェクトでない場合、結果が正しく動作しません。

これらの問題を解決するために、本書では series_batchpyarrow_batch、および numpy_batch の三つのデコレータを提供します。以下のコードでは、これらのデコレータを使用して se_funcpa_func、および np_funcmap_batches() に渡せる形式の関数に変換します。変換後の関数は、引数として列データを受け取ります。列以外の引数を渡したい場合は、func(values=values) のように記述することが可能です。

from helper.polars import pyarrow_batch, numpy_batch, series_batch

@series_batch
def se_func(a, b, values):
    return a.scatter((a < b).arg_true(), values)

@pyarrow_batch
def pa_func(a, b, values):
    from pyarrow.compute import replace_with_mask, less
    return replace_with_mask(a, less(a, b), values)

@numpy_batch
def np_func(a, b, values):
    c = a.copy()
    c[a < b] = values
    return c

values = [100, 200]

df.with_columns(
    se_result=pl.map_batches(['a', 'b'], se_func(values=values)),
    pa_result=pl.map_batches(['a', 'b'], pa_func(values=values)),
    np_result=pl.map_batches(['a', 'b'], np_func(values=values)),
)
shape: (4, 5)
abse_resultpa_resultnp_result
i64i64i64i64i64
11111
23100100100
34200200200
105101010

演算式の置き換え#

以下のcalc_rotate_expr()は、三次元座標の回転を行う演算式を生成します。x, y, zは三次元座標、nx, ny, nzは回転軸を表すベクトル、angleは回転の角度(度数)を表します。

def calc_rotate_expr():
    x, y, z, nx, ny, nz, angle = [pl.col(name) for name in 'x y z nx ny nz angle'.split()]
    magnitude = (nx**2 + ny**2 + nz**2).sqrt()
    nx /= magnitude
    ny /= magnitude
    nz /= magnitude

    # Convert angle to radians
    angle_rad = angle.radians()

    # Calculate cosine and sine of the angle
    cos_theta = angle_rad.cos()
    sin_theta = angle_rad.sin()

    # Calculate the rotated coordinates using Rodrigues' rotation formula
    dot_product = nx * x + ny * y + nz * z
    cross_product_x = ny * z - nz * y
    cross_product_y = nz * x - nx * z
    cross_product_z = nx * y - ny * x

    rx = (cos_theta * x + 
          sin_theta * cross_product_x + 
          (1 - cos_theta) * dot_product * nx)

    ry = (cos_theta * y + 
          sin_theta * cross_product_y + 
          (1 - cos_theta) * dot_product * ny)

    rz = (cos_theta * z + 
          sin_theta * cross_product_z + 
          (1 - cos_theta) * dot_product * nz)

    return pl.struct(rx.alias('rx'), ry.alias('ry'), rz.alias('rz'))

この関数の生成する演算式を、次のように使用できます。

df = pl.DataFrame(dict(
    x=[1, 2, 3, 4],
    y=[2, 1, 3, 2],
    z=[1, 0, 0, 5],
    nx=[1, 1, 1, 1],
    ny=[0, 0, 0, 0],
    nz=[0, 0, 0, 0],
    angle=[10, 20, 30, 40],
))

rot_expr = calc_rotate_expr()
df_rot = df.select(rot_expr.struct.unnest())
row(df, df_rot)
shape: (4, 7)
xyznxnynzangle
i64i64i64i64i64i64i64
12110010
21010020
33010030
42510040
shape: (4, 3)
rxryrz
f64f64f64
1.01.7959671.332104
2.00.9396930.34202
3.02.5980761.5
4.0-1.6818495.115797

しかし、この演算式には以下の問題があります:

  • 列名が演算式内に固定されており、異なる列名を持つデータフレームを直接処理することができない。

  • 一部の列を定数として扱いたい場合、処理が複雑になる。

このような問題を解決するため、本書が提供するexpression_replace()を使用すると、演算式内の列名を簡単に置き換えることができ、さらに、列を定数などの他の演算式に置き換えることも可能です。

expression_replace(expr, mapper=None, **kw)

この関数は、演算式expr内の列名をmapperまたはキーワード引数で置き換えます。

以下のコードでは、rot_expr内のnx, ny, nzおよびangle列を定数に置き換えた後、df2を使って計算を行います。

from helper.polars import expression_replace
rot_expr2 = expression_replace(rot_expr, nx=pl.lit(1), ny=pl.lit(0), nz=pl.lit(0), angle=pl.lit(45))
df2 = df.select('x', 'y', 'z')
df2_rot = df2.select(rot_expr2.struct.unnest())
row(df2, df2_rot)
shape: (4, 3)
xyz
i64i64i64
121
210
330
425
shape: (4, 3)
rxryrz
f64f64f64
1.00.7071072.12132
2.00.7071070.707107
3.02.121322.12132
4.0-2.121324.949747

次のコードは、列名を大文字に対応するように置き換えます。

rot_expr3 = expression_replace(
    rot_expr, 
    x='X', y='Y', z='Z', 
    angle=pl.lit(30)
)

df3 = df.select(pl.col('x', 'y', 'z').name.to_uppercase(), pl.col('nx', 'ny', 'nz'))
df3_rot = df3.select(rot_expr3.struct.unnest())
row(df3, df3_rot)
shape: (4, 6)
XYZnxnynz
i64i64i64i64i64i64
121100
210100
330100
425100
shape: (4, 3)
rxryrz
f64f64f64
1.01.2320511.866025
2.00.8660250.5
3.02.5980761.5
4.0-0.7679495.330127