Skip to content

集計(Aggregation)

Polars は、lazy API だけでなく、eager API でも強力な構文を実装しています。それがどういう意味かを見ていきましょう。

米国議会データセット(US congress dataset)から始めましょう。

DataFrame · Categorical

url = "https://theunitedstates.io/congress-legislators/legislators-historical.csv"

dtypes = {
    "first_name": pl.Categorical,
    "gender": pl.Categorical,
    "type": pl.Categorical,
    "state": pl.Categorical,
    "party": pl.Categorical,
}

dataset = pl.read_csv(url, dtypes=dtypes).with_columns(
    pl.col("birthday").str.to_date(strict=False)
)

DataFrame · Categorical · Available on feature dtype-categorical

use std::io::Cursor;

use reqwest::blocking::Client;

let url = "https://theunitedstates.io/congress-legislators/legislators-historical.csv";

let mut schema = Schema::new();
schema.with_column(
    "first_name".into(),
    DataType::Categorical(None, Default::default()),
);
schema.with_column(
    "gender".into(),
    DataType::Categorical(None, Default::default()),
);
schema.with_column(
    "type".into(),
    DataType::Categorical(None, Default::default()),
);
schema.with_column(
    "state".into(),
    DataType::Categorical(None, Default::default()),
);
schema.with_column(
    "party".into(),
    DataType::Categorical(None, Default::default()),
);
schema.with_column("birthday".into(), DataType::Date);

let data: Vec<u8> = Client::new().get(url).send()?.text()?.bytes().collect();

let dataset = CsvReader::new(Cursor::new(data))
    .has_header(true)
    .with_dtypes(Some(Arc::new(schema)))
    .with_try_parse_dates(true)
    .finish()?;

println!("{}", &dataset);

基本的な集計

list に複数の式を追加することで、簡単に異なる集計を組み合わせられます。 集計の数に上限はなく、好きな組み合わせを作成できます。 以下のスニペットでは、次のような集計を行っています:

"first_name" グループごとに

  • "party" 列の行数をカウントします:
    • 短縮形: pl.count("party")
    • 完全形: pl.col("party").count()
  • "gender" 値をグループ化して集計します:
    • 完全形: pl.col("gender")
  • グループ内の "last_name" 列の最初の値を取得します:
    • 短縮形: pl.first("last_name")(Rustでは使えません)
    • 完全形: pl.col("last_name").first()

集計の後、結果をすぐにソートし、上位 5 件に制限して、 わかりやすい概要を得ています。

group_by

q = (
    dataset.lazy()
    .group_by("first_name")
    .agg(
        pl.len(),
        pl.col("gender"),
        pl.first("last_name"),
    )
    .sort("len", descending=True)
    .limit(5)
)

df = q.collect()
print(df)

group_by

let df = dataset
    .clone()
    .lazy()
    .group_by(["first_name"])
    .agg([len(), col("gender"), col("last_name").first()])
    .sort(
        ["len"],
        SortMultipleOptions::default()
            .with_order_descending(true)
            .with_nulls_last(true),
    )
    .limit(5)
    .collect()?;

println!("{}", df);

shape: (5, 4)
┌────────────┬──────┬───────────────────┬───────────┐
│ first_name ┆ len  ┆ gender            ┆ last_name │
│ ---        ┆ ---  ┆ ---               ┆ ---       │
│ cat        ┆ u32  ┆ list[cat]         ┆ str       │
╞════════════╪══════╪═══════════════════╪═══════════╡
│ John       ┆ 1256 ┆ ["M", "M", … "M"] ┆ Walker    │
│ William    ┆ 1022 ┆ ["M", "M", … "M"] ┆ Few       │
│ James      ┆ 714  ┆ ["M", "M", … "M"] ┆ Armstrong │
│ Thomas     ┆ 453  ┆ ["M", "M", … "M"] ┆ Tucker    │
│ Charles    ┆ 439  ┆ ["M", "M", … "M"] ┆ Carroll   │
└────────────┴──────┴───────────────────┴───────────┘

条件式

簡単ですね!さらに進めましょう。 "state" の代表者が "Pro" または "Anti" 政権かどうかを知りたいとします。 lambdaDataFrame の整理に頼ることなく、集計の中で直接クエリを使えます。

group_by

q = (
    dataset.lazy()
    .group_by("state")
    .agg(
        (pl.col("party") == "Anti-Administration").sum().alias("anti"),
        (pl.col("party") == "Pro-Administration").sum().alias("pro"),
    )
    .sort("pro", descending=True)
    .limit(5)
)

df = q.collect()
print(df)

group_by

let df = dataset
    .clone()
    .lazy()
    .group_by(["state"])
    .agg([
        (col("party").eq(lit("Anti-Administration")))
            .sum()
            .alias("anti"),
        (col("party").eq(lit("Pro-Administration")))
            .sum()
            .alias("pro"),
    ])
    .sort(
        ["pro"],
        SortMultipleOptions::default().with_order_descending(true),
    )
    .limit(5)
    .collect()?;

println!("{}", df);

shape: (5, 3)
┌───────┬──────┬─────┐
│ state ┆ anti ┆ pro │
│ ---   ┆ ---  ┆ --- │
│ cat   ┆ u32  ┆ u32 │
╞═══════╪══════╪═════╡
│ CT    ┆ 0    ┆ 3   │
│ NJ    ┆ 0    ┆ 3   │
│ NC    ┆ 1    ┆ 2   │
│ SC    ┆ 0    ┆ 1   │
│ PA    ┆ 1    ┆ 1   │
└───────┴──────┴─────┘

同様のことは、ネストされた GROUP BY でも行えますが、これらの素晴らしい機能を示すのに役立ちません。 😉

group_by

q = (
    dataset.lazy()
    .group_by("state", "party")
    .agg(pl.count("party").alias("count"))
    .filter(
        (pl.col("party") == "Anti-Administration")
        | (pl.col("party") == "Pro-Administration")
    )
    .sort("count", descending=True)
    .limit(5)
)

df = q.collect()
print(df)

group_by

let df = dataset
    .clone()
    .lazy()
    .group_by(["state", "party"])
    .agg([col("party").count().alias("count")])
    .filter(
        col("party")
            .eq(lit("Anti-Administration"))
            .or(col("party").eq(lit("Pro-Administration"))),
    )
    .sort(
        ["count"],
        SortMultipleOptions::default()
            .with_order_descending(true)
            .with_nulls_last(true),
    )
    .limit(5)
    .collect()?;

println!("{}", df);

shape: (5, 3)
┌───────┬─────────────────────┬───────┐
│ state ┆ party               ┆ count │
│ ---   ┆ ---                 ┆ ---   │
│ cat   ┆ cat                 ┆ u32   │
╞═══════╪═════════════════════╪═══════╡
│ NJ    ┆ Pro-Administration  ┆ 3     │
│ VA    ┆ Anti-Administration ┆ 3     │
│ CT    ┆ Pro-Administration  ┆ 3     │
│ NC    ┆ Pro-Administration  ┆ 2     │
│ PA    ┆ Anti-Administration ┆ 1     │
└───────┴─────────────────────┴───────┘

フィルタリング

グループをフィルタリングすることもできます。 グループごとの平均を計算したいが、そのグループのすべての値を含めたくない、 また DataFrame の行をフィルタリングしたくない(別の集計に必要なため)場合などです。

以下の例では、これがどのように行えるかを示しています。

Note

Python 関数を明確にするためのメモ。これらの関数にはコストがかかりません。なぜなら、Polars エクスプレッションのみを作成し、クエリの実行時にカスタム関数を Series 上で適用しないためです。もちろん、Rust でもエクスプレッションを返す関数を作ることができます。

group_by

from datetime import date


def compute_age():
    return date.today().year - pl.col("birthday").dt.year()


def avg_birthday(gender: str) -> pl.Expr:
    return (
        compute_age()
        .filter(pl.col("gender") == gender)
        .mean()
        .alias(f"avg {gender} birthday")
    )


q = (
    dataset.lazy()
    .group_by("state")
    .agg(
        avg_birthday("M"),
        avg_birthday("F"),
        (pl.col("gender") == "M").sum().alias("# male"),
        (pl.col("gender") == "F").sum().alias("# female"),
    )
    .limit(5)
)

df = q.collect()
print(df)

group_by

fn compute_age() -> Expr {
    lit(2022) - col("birthday").dt().year()
}

fn avg_birthday(gender: &str) -> Expr {
    compute_age()
        .filter(col("gender").eq(lit(gender)))
        .mean()
        .alias(&format!("avg {} birthday", gender))
}

let df = dataset
    .clone()
    .lazy()
    .group_by(["state"])
    .agg([
        avg_birthday("M"),
        avg_birthday("F"),
        (col("gender").eq(lit("M"))).sum().alias("# male"),
        (col("gender").eq(lit("F"))).sum().alias("# female"),
    ])
    .limit(5)
    .collect()?;

println!("{}", df);

shape: (5, 5)
┌───────┬────────────────┬────────────────┬────────┬──────────┐
│ state ┆ avg M birthday ┆ avg F birthday ┆ # male ┆ # female │
│ ---   ┆ ---            ┆ ---            ┆ ---    ┆ ---      │
│ cat   ┆ f64            ┆ f64            ┆ u32    ┆ u32      │
╞═══════╪════════════════╪════════════════╪════════╪══════════╡
│ WY    ┆ 140.717949     ┆ 68.0           ┆ 39     ┆ 2        │
│ AK    ┆ 123.411765     ┆ null           ┆ 17     ┆ 0        │
│ AS    ┆ 84.0           ┆ null           ┆ 2      ┆ 0        │
│ PI    ┆ 148.0          ┆ null           ┆ 13     ┆ 0        │
│ DK    ┆ 194.333333     ┆ null           ┆ 9      ┆ 0        │
└───────┴────────────────┴────────────────┴────────┴──────────┘

ソート

GROUP BY 操作の順序を管理するために、DataFrame をソートすることは一般的です。州ごとの最年長および最年少の政治家の名前を取得したいとします。その際は、SORT と GROUP BY を使うことができます。

group_by

def get_person() -> pl.Expr:
    return pl.col("first_name") + pl.lit(" ") + pl.col("last_name")


q = (
    dataset.lazy()
    .sort("birthday", descending=True)
    .group_by("state")
    .agg(
        get_person().first().alias("youngest"),
        get_person().last().alias("oldest"),
    )
    .limit(5)
)

df = q.collect()
print(df)

group_by

fn get_person() -> Expr {
    col("first_name") + lit(" ") + col("last_name")
}

let df = dataset
    .clone()
    .lazy()
    .sort(
        ["birthday"],
        SortMultipleOptions::default()
            .with_order_descending(true)
            .with_nulls_last(true),
    )
    .group_by(["state"])
    .agg([
        get_person().first().alias("youngest"),
        get_person().last().alias("oldest"),
    ])
    .limit(5)
    .collect()?;

println!("{}", df);

shape: (5, 3)
┌───────┬───────────────────────┬─────────────────┐
│ state ┆ youngest              ┆ oldest          │
│ ---   ┆ ---                   ┆ ---             │
│ cat   ┆ str                   ┆ str             │
╞═══════╪═══════════════════════╪═════════════════╡
│ MN    ┆ Erik Paulsen          ┆ Cyrus Aldrich   │
│ KY    ┆ John Edwards          ┆ Matthew Lyon    │
│ CA    ┆ Edward Gilbert        ┆ William Gwin    │
│ NY    ┆ Cornelius Schoonmaker ┆ Philip Schuyler │
│ ID    ┆ Raúl Labrador         ┆ William Wallace │
└───────┴───────────────────────┴─────────────────┘

ただし、もし 名前をアルファベット順にソートしたい場合、これは機能しません。幸いにも、group_by 式で DataFrame とは別にソートができます。

group_by

def get_person() -> pl.Expr:
    return pl.col("first_name") + pl.lit(" ") + pl.col("last_name")


q = (
    dataset.lazy()
    .sort("birthday", descending=True)
    .group_by("state")
    .agg(
        get_person().first().alias("youngest"),
        get_person().last().alias("oldest"),
        get_person().sort().first().alias("alphabetical_first"),
    )
    .limit(5)
)

df = q.collect()
print(df)

group_by

let df = dataset
    .clone()
    .lazy()
    .sort(
        ["birthday"],
        SortMultipleOptions::default()
            .with_order_descending(true)
            .with_nulls_last(true),
    )
    .group_by(["state"])
    .agg([
        get_person().first().alias("youngest"),
        get_person().last().alias("oldest"),
        get_person()
            .sort(Default::default())
            .first()
            .alias("alphabetical_first"),
    ])
    .limit(5)
    .collect()?;

println!("{}", df);

shape: (5, 4)
┌───────┬──────────────────┬──────────────────┬─────────────────────────┐
│ state ┆ youngest         ┆ oldest           ┆ alphabetical_first      │
│ ---   ┆ ---              ┆ ---              ┆ ---                     │
│ cat   ┆ str              ┆ str              ┆ str                     │
╞═══════╪══════════════════╪══════════════════╪═════════════════════════╡
│ WY    ┆ Liz Cheney       ┆ Stephen Nuckolls ┆ Alan Simpson            │
│ AK    ┆ Mark Begich      ┆ Thomas Cale      ┆ Anthony Dimond          │
│ DK    ┆ George Mathews   ┆ John Todd        ┆ George Mathews          │
│ AS    ┆ Eni Faleomavaega ┆ Fofó Sunia       ┆ Eni Faleomavaega        │
│ PI    ┆ Carlos Romulo    ┆ Pablo Ocampo     ┆ Benito Legarda Y Tuason │
└───────┴──────────────────┴──────────────────┴─────────────────────────┘

group_by 式の中で別の列を基準にソートすることもできます。アルファベット順にソートされた名前が男性か女性かを知りたい場合は:pl.col("gender").sort_by("first_name").first().alias("gender") と記述できます。

group_by

def get_person() -> pl.Expr:
    return pl.col("first_name") + pl.lit(" ") + pl.col("last_name")


q = (
    dataset.lazy()
    .sort("birthday", descending=True)
    .group_by("state")
    .agg(
        get_person().first().alias("youngest"),
        get_person().last().alias("oldest"),
        get_person().sort().first().alias("alphabetical_first"),
        pl.col("gender")
        .sort_by(pl.col("first_name").cast(pl.Categorical("lexical")))
        .first(),
    )
    .sort("state")
    .limit(5)
)

df = q.collect()
print(df)

group_by

let df = dataset
    .clone()
    .lazy()
    .sort(
        ["birthday"],
        SortMultipleOptions::default()
            .with_order_descending(true)
            .with_nulls_last(true),
    )
    .group_by(["state"])
    .agg([
        get_person().first().alias("youngest"),
        get_person().last().alias("oldest"),
        get_person()
            .sort(Default::default())
            .first()
            .alias("alphabetical_first"),
        col("gender")
            .sort_by(["first_name"], SortMultipleOptions::default())
            .first()
            .alias("gender"),
    ])
    .sort(["state"], SortMultipleOptions::default())
    .limit(5)
    .collect()?;

println!("{}", df);

shape: (5, 5)
┌───────┬───────────────────────┬──────────────────────┬────────────────────┬────────┐
│ state ┆ youngest              ┆ oldest               ┆ alphabetical_first ┆ gender │
│ ---   ┆ ---                   ┆ ---                  ┆ ---                ┆ ---    │
│ cat   ┆ str                   ┆ str                  ┆ str                ┆ cat    │
╞═══════╪═══════════════════════╪══════════════════════╪════════════════════╪════════╡
│ NY    ┆ Cornelius Schoonmaker ┆ Philip Schuyler      ┆ A. Foster          ┆ M      │
│ NC    ┆ John Ashe             ┆ Samuel Johnston      ┆ Abraham Rencher    ┆ M      │
│ NE    ┆ Samuel Daily          ┆ Experience Estabrook ┆ Albert Jefferis    ┆ M      │
│ MS    ┆ Narsworthy Hunter     ┆ Thomas Greene        ┆ Aaron Ford         ┆ M      │
│ IL    ┆ Benjamin Stephenson   ┆ Shadrack Bond        ┆ Aaron Schock       ┆ M      │
└───────┴───────────────────────┴──────────────────────┴────────────────────┴────────┘

並列処理を阻害しない

Python ユーザーのみ

以下のセクションは Python に固有のものであり、Rust には適用されません。Rust では、ブロックとクロージャ(ラムダ)を並行して実行できるためです。

Python は遅くて "スケールしない" というのは、誰もが耳にしたことがあるでしょう。 "遅い" バイトコードを実行するオーバーヘッドに加えて、Python は Global Interpreter Lock(GIL)の制約の中にいなければなりません。 つまり、並列化フェーズで lambda やカスタム Python 関数を使用する場合、 Polars の速度は Python コードの実行によって制限され、 複数のスレッドが関数を実行することを妨げます。

これはとてもうっとうしい制限に感じられますが、特に .group_by() ステップでは lambda 関数が必要になることが多いです。 このアプローチは Polars でまだサポートされていますが、バイトコード GIL のコストを支払う必要があることを念頭に置いてください。エクスプレッションの構文を使ってクエリを解決することをお勧めします。 lambda の使用については、ユーザー定義関数セクション を参照してください。

まとめ

上記の例では、エクスプレッションを組み合わせることで多くのことができることを見てきました。そうすることで、(Python と GIL の遅い性質によって)クエリを遅くする Python のカスタム関数の使用を遅らせられます。

もしエクスプレッションのタイプが見つからない場合は、feature requestを開いてお知らせください!