皆さんこんにちは。 機械学習エンジニアのwasatingです。
今年に入ってから、DS/ML界隈でスタンダードであったpandasから、polarsに乗り換える人が増え始めていると思いますが、皆さんはいかがでしょうか。
私も例に漏れずプライベートでの開発時は完全にpolarsに置き換え、プロダクトでもpolarsへの置き換えを推進しています。
しかしながら、pandasに比べるとpolarsに関する技術ブログ等はなかなか少なく、公式ドキュメントを読んで初めて「そんな機能あったんか…」となることも少なくないと思います。
というわけで今回はpolars.selectorsについて紹介していきたいと思います。
一応公式のリンクもおいておきます。
example
import polars.selectors as cs import polars as pl df = pl.DataFrame( { "w": ["xx", "yyy", "xx", "yyy", "xx"], "x": [1, 2, 1, 4, -2], "y": [3.0, 4.5, 1.0, 2.5, -2.0], "z": ["a", "b", "a", "b", "b"], }, ) print(df.groupby(by=cs.string()).agg(cs.numeric().sum())) ------ shape: (3, 4) ┌─────┬─────┬─────┬──────┐ │ w ┆ z ┆ x ┆ y │ │ --- ┆ --- ┆ --- ┆ --- │ │ str ┆ str ┆ i64 ┆ f64 │ ╞═════╪═════╪═════╪══════╡ │ yy ┆ b ┆ 6 ┆ 7.0 │ │ xx ┆ a ┆ 2 ┆ 4.0 │ │ xx ┆ b ┆ -2 ┆ -2.0 │ └─────┴─────┴─────┴──────┘
ちなみにimport polars.selectors as csとするように公式が推奨しています
上記処理を見ていただくとわかる通り、selectorsはExprと同様に扱うことができます。
さて、そんなselectorsですが、ユースケースとしては
- カラムのdtypeをもとに操作
- カラム名をもとに操作
の2つが主に挙げられると思いますので、以下で紹介していきます。
dtypeをもとに操作
string(), integer(), float(), datetime()等を使用することで、対応するカラムの取得ができます。
また、Exprと同等の扱いができるので、pl.col('hoge') * 2と同様の処理ができます
df = pl.DataFrame(
{
"foo": ["x", "yy"],
"bar": [123, 456],
"baz": [2.0, 5.5],
"zap": [0.0, 1.0],
},
schema_overrides={"baz": pl.Float32, "zap": pl.Float64},
)
df_selector = df.select(
cs.float(),
(cs.float() * 2).suffix('_twice'),
cs.string(),
cs.string().str.n_chars().suffix('_chars'),
)
print(df_selector)
------
shape: (2, 6)
┌─────┬─────┬───────────┬───────────┬─────┬───────────┐
│ baz ┆ zap ┆ baz_twice ┆ zap_twice ┆ foo ┆ foo_chars │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ f32 ┆ f64 ┆ f32 ┆ f64 ┆ str ┆ u32 │
╞═════╪═════╪═══════════╪═══════════╪═════╪═══════════╡
│ 2.0 ┆ 0.0 ┆ 4.0 ┆ 0.0 ┆ x ┆ 1 │
│ 5.5 ┆ 1.0 ┆ 11.0 ┆ 2.0 ┆ yy ┆ 2 │
└─────┴─────┴───────────┴───────────┴─────┴───────────┘
また、演算子として
- or:
| - and:
& - not:
~ - diff:
-
を使えるので、以下のようなこともできます。
df_operation = df.select(cs.float() & ~cs.by_dtype(pl.Float64)) print(df_operation) ------ shape: (2, 1) ┌─────┐ │ baz │ │ --- │ │ f32 │ ╞═════╡ │ 2.0 │ │ 5.5 │ └─────┘
カラム名をもとに操作
contains(), by_name(), ends_with()を使うことで、対応するカラムを取得することができます。
また、こちらも上記した演算子を使うことができます。
df = pl.DataFrame(
{
"foo": ["x", "yy"],
"bar": [123, 456],
"baz": [2.0, 5.5],
"zap": [0.0, 1.0],
},
schema_overrides={"baz": pl.Float32, "zap": pl.Float64},
)
print(df.select(cs.contains('ba')))
------
shape: (2, 2)
┌─────┬─────┐
│ bar ┆ baz │
│ --- ┆ --- │
│ i64 ┆ f32 │
╞═════╪═════╡
│ 123 ┆ 2.0 │
│ 456 ┆ 5.5 │
└─────┴─────┘
print(df.select(cs.ends_with('z', 'p'))) ------ shape: (2, 2) ┌─────┬─────┐ │ baz ┆ zap │ │ --- ┆ --- │ │ f32 ┆ f64 │ ╞═════╪═════╡ │ 2.0 ┆ 0.0 │ │ 5.5 ┆ 1.0 │ └─────┴─────┘
なお、ends_with, by_name等は上述したやり方で複数条件を設定できますが、containsのみ、list, tuple等のiterableなものを渡す必要があり、注意が必要です。
おわりに
いかがでしたか。今回の記事が皆さんのpolarsへの移行の助けになっていれば幸いです。
また、今回は省略しましたが、selectors.matchを使うことで正規表現でカラム指定をすることも可能です
おまけ
実行環境
import polars as pl pl.show_versions() --------Version info--------- Polars: 0.18.15 Index type: UInt32 Platform: Linux-5.10.16.3-microsoft-standard-WSL2-x86_64-with-glibc2.31 Python: 3.9.11 (main, Mar 18 2022, 16:45:24) [GCC 10.2.1 20210110] ----Optional dependencies---- adbc_driver_sqlite: <not installed> cloudpickle: <not installed> connectorx: <not installed> deltalake: <not installed> fsspec: 2022.11.0 matplotlib: 3.6.0 numpy: 1.23.4 pandas: 1.5.1 pyarrow: 12.0.1 pydantic: 1.10.12 sqlalchemy: <not installed> xlsx2csv: <not installed> xlsxwriter: <not installed>
