polarsでstrのcapture groupを参照しつつreplaceする

polars

Python

Rust

Kazuki Moriyama (森山和樹)

polarsのstr.replace

polarsにはstringのカラムに対してreplaceをするための仕組みが備わっている。

df = pl.DataFrame({"id": [1, 2], "text": ["123abc", "abc456"]})
df.with_columns(
    pl.col("text").str.replace(r"abc\b", "ABC")
)  
...
shape: (2, 2)
┌─────┬────────┐
│ id  ┆ text   │
│ --- ┆ ---    │
│ i64 ┆ str    │
╞═════╪════════╡
│ 1   ┆ 123ABC │
│ 2   ┆ abc456 │
└─────┴────────┘

pythonなどいろんな言語でよくやるstringの部分をマッチさせて、その部分を対象の文字列に置き換えたりする処理。
polarsではexpr.strに定義されていて、Exprとして引き回せる。

内部ではrustのregex crateに依存しているので、そのcrateの仕組みが使える。

regexでcapture groupを使ってreplaceする

regexを使ったreplace処理でよくあるのはcapture groupを利用したreplace。
例えばrustとpythonで書くと以下の様。

use regex::Regex;

fn main() {
    // numbered capture group
    let re = Regex::new(r"^.* (\d*)$").unwrap();
    let number = re.replace("Hello 123", "$1");
    println!("{}", number); // => 123

    // named capture group
    let named_re = Regex::new(r"^(?P<name>.*) (?P<number>\d*)$").unwrap();
    let named_number = named_re.replace("Hello 123", "name=$name number=$number");
    println!("{}", named_number); // => name=Hello number=123
}

import re

a = "Hello 123"

# numbered capture
number = re.sub(r"^.* (\d*)$", "\g<1>", a)
print(number) # 123

# named capture
named_number = re.sub(r"^(?P<name>.*) (?P<number>\d*)$", "name=\g<name> number=\g<number>", a)
print(named_number) # name=Hello number=123

上の様に()で囲んだcapture groupに対して、それを置換文字列内で参照できる。
また$1の様に数字だけでなく前もってgroupに名前をつけて置けば$nameの様に名前での参照も可能になる。

当然polarsを使っているときもstringの列に対して上記の様なことをやりたくなる。

polarsでcapture group replace

polarsは前述の通り中身はrustのregex crateなのでrustの書き方と同じ様に書けば普通に動く。
以下はpythonのpolars。

import polars as pl

df = pl.DataFrame(
    {
        "a": ["1 one 2", "2 two 3", "3 three 4"],
    }
)

print(df)

numbered_col = pl.col("a").str.replace(r"(\d) .* (\d)", "head=$1 tail=$2")
named_col = pl.col("a").str.replace(r"(?P<first>\d) .* (?P<second>\d)", "head=$first tail=$second")

# どちらも同じ出力
print(df.select(numbered_col))
print(df.select(named_col))

shape: (3, 1)
┌───────────────┐
│ a             │
│ ---           │
│ str           │
╞═══════════════╡
│ head=1 tail=2 │
│ head=2 tail=3 │
│ head=3 tail=4 │
└───────────────┘