ユーザー定義関数（User-defined functions）

Polars のエクスプレッションは非常に強力で柔軟であるため、他のライブラリよりもカスタム Python 関数の必要性ははるかに少ないとここまでで納得していただけたかと思います。

それでもなお、エクスプレッションの状態をサードパーティのライブラリに渡す、または Polars でデータに対してブラックボックス関数を適用する機能が必要です。

このために、以下のエクスプレッションを提供しています:

map_batches
map_elements

`map_batches` か、それとも `map_elements` か。

これらの関数には、どのように動作するか、そしてその結果としてユーザーにどのようなデータを渡すかという重要な違いがあります。

map_batches は、そのままの Series を expression に渡します。

map_batches は、select と group_by の両式で同じルールに従います。これは、Series が DataFrame のカラムを表すことを意味します。 group_by 式では、そのカラムはまだ集約されていないことに注意してください！

map_batches の使用例としては、例えばエクスプレッションのカラムをサードパーティのライブラリに渡すことが挙げられます。以下に、ニューラルネットワークモデルにエクスプレッションカラムを渡す方法を示します。

Python

map_batches

df.with_columns([
    pl.col("features").map_batches(lambda s: MyNeuralNetwork.forward(s.to_numpy())).alias("activations")
])

Rust

df.with_columns([
    col("features").map(|s| Ok(my_nn.forward(s))).alias("activations")
])

group_by 式で map_batches を使用するケースは限られています。それはパフォーマンス上の理由からのみ使用され、簡単に誤った結果をもたらす可能性があります。その理由を説明しましょう。

Python Rust

df = pl.DataFrame(
    {
        "keys": ["a", "a", "b"],
        "values": [10, 7, 1],
    }
)
print(df)

let df = df!(
    "keys" => &["a", "a", "b"],
    "values" => &[10, 7, 1],
)?;
println!("{}", df);

shape: (3, 2)
┌──────┬────────┐
│ keys ┆ values │
│ ---  ┆ ---    │
│ str  ┆ i64    │
╞══════╪════════╡
│ a    ┆ 10     │
│ a    ┆ 7      │
│ b    ┆ 1      │
└──────┴────────┘

上のスニペットでは、"keys" カラムでグループ化します。つまり、次のようなグループがあります：

"a" -> [10, 7]
"b" -> [1]

その後、右への shift 操作を適用すると、次のようになるでしょう：

"a" -> [null, 10]
"b" -> [null]

それを試してみて、何が得られるかを見てみましょう：

Python Rust

out = df.group_by("keys", maintain_order=True).agg(
    pl.col("values")
    .map_batches(lambda s: s.shift(), is_elementwise=True)
    .alias("shift_map_batches"),
    pl.col("values").shift().alias("shift_expression"),
)
print(out)

let out = df
    .clone()
    .lazy()
    .group_by(["keys"])
    .agg([
        col("values")
            .map(|s| Ok(Some(s.shift(1))), GetOutput::default())
            // note: the `'shift_map_batches'` alias is just there to show how you
            // get the same output as in the Python API example.
            .alias("shift_map_batches"),
        col("values").shift(lit(1)).alias("shift_expression"),
    ])
    .collect()?;

println!("{}", out);

shape: (2, 3)
┌──────┬───────────────────┬──────────────────┐
│ keys ┆ shift_map_batches ┆ shift_expression │
│ ---  ┆ ---               ┆ ---              │
│ str  ┆ list[i64]         ┆ list[i64]        │
╞══════╪═══════════════════╪══════════════════╡
│ a    ┆ [null, 10]        ┆ [null, 10]       │
│ b    ┆ [7]               ┆ [null]           │
└──────┴───────────────────┴──────────────────┘

あちゃー、明らかに間違った結果が出ましたね。グループ "b" はグループ "a" から値を得てしまいました 😵。

これは、集約する前に map_batches が関数を適用するため、ひどく間違ってしまったのです。それはつまり、全カラム [10, 7, 1] がシフトされて [null, 10, 7] になり、その後で集約されたということです。

だから私のアドバイスは、map_batches を group_by 式で使用しないことです。それが必要であり、何をしているかを知っている場合を除きます。

`map_elements` を使う

幸い、前の例は map_elements で修正できます。 map_elements は、その操作のための最小論理要素で動作します。

つまり：

select context -> 単一要素
group by context -> 単一グループ

したがって、map_elements を使えば、私たちの例を修正できるはずです：

Python

map_elements

Python Rust

out = df.group_by("keys", maintain_order=True).agg(
    pl.col("values")
    .map_elements(lambda s: s.shift(), return_dtype=pl.List(int))
    .alias("shift_map_elements"),
    pl.col("values").shift().alias("shift_expression"),
)
print(out)

let out = df
    .clone()
    .lazy()
    .group_by([col("keys")])
    .agg([
        col("values")
            .apply(|s| Ok(Some(s.shift(1))), GetOutput::default())
            // note: the `'shift_map_elements'` alias is just there to show how you
            // get the same output as in the Python API example.
            .alias("shift_map_elements"),
        col("values").shift(lit(1)).alias("shift_expression"),
    ])
    .collect()?;
println!("{}", out);

shape: (2, 3)
┌──────┬────────────────────┬──────────────────┐
│ keys ┆ shift_map_elements ┆ shift_expression │
│ ---  ┆ ---                ┆ ---              │
│ str  ┆ list[i64]          ┆ list[i64]        │
╞══════╪════════════════════╪══════════════════╡
│ a    ┆ [null, 10]         ┆ [null, 10]       │
│ b    ┆ [null]             ┆ [null]           │
└──────┴────────────────────┴──────────────────┘

そして観察すると、有効な結果が得られます！ 🎉

`map_elements` の `select` 式での使用

select 式では、map_elements エクスプレッションはカラムの要素を Python 関数に渡します。

注意してください。これは Python を実行しているので、遅くなります。

このセクションの最初に定義した DataFrame で続けて、 map_elements 関数の例と同じ目標を達成するためにエクスプレッション API を使用する反例を見てみましょう。

カウンターの追加

この例では、グローバルな counter を作成し、処理される各要素に整数 1 を加えます。各反復で、インクリメントの結果が要素値に追加されます。

注：この例は Rust では提供されていません。その理由は、グローバルな counter 値が並行評価されるときにデータ競合を引き起こす可能性があるためです。それを Mutex でラップして変数を保護することは可能ですが、それは例のポイントを曖昧にすることになります。この場合、Python Global Interpreter Lock のパフォーマンスのトレードオフがいくつかの安全保障を提供します。

Python Rust

counter = 0


def add_counter(val: int) -> int:
    global counter
    counter += 1
    return counter + val


out = df.select(
    pl.col("values").map_elements(add_counter).alias("solution_map_elements"),
    (pl.col("values") + pl.int_range(1, pl.len() + 1)).alias("solution_expr"),
)
print(out)

shape: (3, 2)
┌───────────────────────┬───────────────┐
│ solution_map_elements ┆ solution_expr │
│ ---                   ┆ ---           │
│ i64                   ┆ i64           │
╞═══════════════════════╪═══════════════╡
│ 11                    ┆ 11            │
│ 9                     ┆ 9             │
│ 4                     ┆ 4             │
└───────────────────────┴───────────────┘

複数のカラム値の組み合わせ

単一の map_elements 関数コールで異なるカラムの値にアクセスしたい場合、 struct データタイプを作成することができます。このデータタイプは、struct 内のフィールドとしてそれらのカラムを収集します。したがって、"keys" と "values" のカラムから struct を作成すると、次のような struct 要素が得られます：

[
    {"keys": "a", "values": 10},
    {"keys": "a", "values": 7},
    {"keys": "b", "values": 1},
]

Python では、これらは呼び出し元の Python 関数に dict として渡され、field: str によってインデックスされることができます。Rust では、Struct タイプの Series を取得します。struct のフィールドはインデックス化され、ダウンキャストすることができます。

Python Rust

out = df.select(
    pl.struct(["keys", "values"])
    .map_elements(lambda x: len(x["keys"]) + x["values"])
    .alias("solution_map_elements"),
    (pl.col("keys").str.len_bytes() + pl.col("values")).alias("solution_expr"),
)
print(out)

let out = df
    .lazy()
    .select([
        // pack to struct to get access to multiple fields in a custom `apply/map`
        as_struct(vec![col("keys"), col("values")])
            // we will compute the len(a) + b
            .apply(
                |s| {
                    // downcast to struct
                    let ca = s.struct_()?;

                    // get the fields as Series
                    let s_a = &ca.fields()[0];
                    let s_b = &ca.fields()[1];

                    // downcast the `Series` to their known type
                    let ca_a = s_a.str()?;
                    let ca_b = s_b.i32()?;

                    // iterate both `ChunkedArrays`
                    let out: Int32Chunked = ca_a
                        .into_iter()
                        .zip(ca_b)
                        .map(|(opt_a, opt_b)| match (opt_a, opt_b) {
                            (Some(a), Some(b)) => Some(a.len() as i32 + b),
                            _ => None,
                        })
                        .collect();

                    Ok(Some(out.into_series()))
                },
                GetOutput::from_type(DataType::Int32),
            )
            // note: the `'solution_map_elements'` alias is just there to show how you
            // get the same output as in the Python API example.
            .alias("solution_map_elements"),
        (col("keys").str().count_matches(lit("."), true) + col("values"))
            .alias("solution_expr"),
    ])
    .collect()?;
println!("{}", out);

shape: (3, 2)
┌───────────────────────┬───────────────┐
│ solution_map_elements ┆ solution_expr │
│ ---                   ┆ ---           │
│ i64                   ┆ i64           │
╞═══════════════════════╪═══════════════╡
│ 11                    ┆ 11            │
│ 8                     ┆ 8             │
│ 2                     ┆ 2             │
└───────────────────────┴───────────────┘

Structs については次のセクションで詳しく説明します。

戻り値は？

カスタム Python 関数は Polars にとってブラックボックスです。ですので、あなたが何を意図しているのかを推測し、最善を尽くして理解しようとする必要があります。

ユーザーとしては、カスタム関数をよりよく利用するために私たちが何をするかを理解することが役立ちます。

データタイプは自動的に推測されます。私たちは最初の非 null 値を待ち、その値を使用して Series のタイプを決定します。

Python の型から Polars のデータタイプへのマッピングは次の通りです：

int -> Int64
float -> Float64
bool -> Boolean
str -> String
list[tp] -> List[tp] (内部タイプは同じルールで推測)
dict[str, [tp]] -> struct
Any -> object (これは常に避けてください)

Rust の型のマッピングは次の通りです：

i32 または i64 -> Int64
f32 または f64 -> Float64
bool -> Boolean
String または str -> String
Vec<tp> -> List[tp] (内部タイプは同じルールで推測)

ユーザー定義関数（User-defined functions）

map_batches か、それとも map_elements か。

map_elements を使う

map_elements の select 式での使用

カウンターの追加

複数のカラム値の組み合わせ

戻り値は？

`map_batches` か、それとも `map_elements` か。

`map_elements` を使う

`map_elements` の `select` 式での使用