📡

Rust版polarsでデータフレームからユニークなカテゴリ列(vec型データ)を取得する方法

2025/01/25に公開

1.はじめに

Rust版Polarsでデータフレームの数値型や文字列型の列から、ユニークなカテゴリ列を Vec 型で取得したい場合があります。
しかし、Polarsが扱うデータ型から直接 Vec 型に変換するメソッドが見当たらなかったため、Vec 型で取得するコードを作成しました。(もし、直接変換するメソッドがあればご教授いただけますと幸いです。)

2. cargo.toml

今回は LazyFrame を使用するため、features に lazy を指定します。

[dependencies]
polars = { version = "0.45.1", features = ["lazy"] }

3. main.rs

数値型の列が指定された場合でも処理できるように、cast メソッドで String 型に変換しています。
LazyFrame型に変換後、DataType::String で文字列にキャストします。

3.1. サンプルデータフレームの作成

    // サンプルデータフレームを作成
    let df = df! [
        "category" => ["A", "B", "A", "C", "B"],
        "values" => [1, 2, 3, 4, 5]
    ]?;
    let column_name = "category";
    let df_casted = df.clone().lazy().select([col(column_name).cast(DataType::String).unique().alias(column_name)]).collect().unwrap();
    let mut unique_vec = Vec::new();

3.2. ユニークなカテゴリ列の取得

変換後のデータフレームから、カテゴリ列を ChunkedArray<String> 型として取得します。

    let unique_chunk = df_casted.column(column_name)?.str()?;

3.3. ChunkedArray<String> 型への変換

into_no_null_iter() メソッドを使用することで、Null 値を含まない要素のみを処理します。

for value in unique_chunk.into_no_null_iter() {
    unique_vec.push(value.to_string());
}

下記に完成版コードを示します。

4. 完成版コード

use polars::prelude::*;
use std::error::Error;

fn main() -> Result<(), Box<dyn Error>> {
    // サンプルデータフレームを作成
    let df = df! [
        "category" => ["A", "B", "A", "C", "B"],
        "values" => [1, 2, 3, 4, 5]
    ]?;
    let column_name = "category";
    let df_casted = df.clone().lazy().select([col(column_name).cast(DataType::String).unique().alias(column_name)]).collect().unwrap();
    let mut unique_vec = Vec::new();
    let unique_chunk = df_casted.column(column_name)?.str()?;
    for value in unique_chunk.into_no_null_iter() {
        unique_vec.push(value.to_string());
    }
    // ユニーク値を出力
    println!("Unique values: {:?}", unique_vec);

    Ok(())
}

Discussion