[Python] Polars Tips

Python
2023.02.03

dataframeの作成

データフレームを作成します。

data = {"cols1": ["a", "b"], "cols2": [1, 2]}
df = pl.DataFrame(data)

作成する際に、型も指定できます。

data = {"col1": [1, 2], "col2": [3, 4]}
df = pl.DataFrame(data, schema={"col1": pl.Float32, "col2": pl.Int64})

データの読み込み、書き出し

# csv
# 読み込み
df = pl.read_csv("train.csv")
# 書き出し
df.write_csv("train.csv")

# parquet
# 読み込み
df = pl.read_parquet("train.parquet")
# 書き出し
df.write_parquet("train.parquet")

# よく使う引数
df = pl.read_parquet("train.parquet", n_rows=100, columns=["col1", "col2", "col3"])

filter

データの検索や抜き出すときなどに使用します。

df.filter(pl.col("foo") == 0)

# 複数条件 
df.filter((pl.col("foo") < 3) & (pl.col("ham") == "a"))

新規カラムの作成

新しいカラム名はaliasで指定します。

# カラムaの2倍の値を、カラムcに作成
df = df.with_columns((pl.col("a") ** 2).alias("c"))

リストを新しいカラムとして追加します。サイズが合わないとエラーになるので、注意です

# listの値をnew_colとして追加
df = df.with_columns(pl.Series(name='new_col', values=list))

型変換(cast)

カラムの型を変換します。カラムの型を調べるのはhead()を使うのが楽だと思います。

df = df.with_columns([pl.col("vote_sum").cast(pl.Float32)])
Float32
Float64
Int8
Int16
Int32
Int64
UInt8
UInt16
UInt32
UInt64

データフレームの結合(join)

pandasでいう、mergeとほぼ一緒です。

how: {‘inner’, ‘left’, ‘outer’, ‘semi’, ‘anti’, ‘cross’}
suffix: joinしたときに、重複するカラム名に付加する名前

df = df.join(subs, how='left', on=['session', "aid"])

結合(concat)

df = pl.concat([df1, df2])

# 横方向
df = pl.concat([df1, df2], how="horizontal")

カラム名のリネーム(rename)

カラム名のリネームです。

df = df.rename({'a': 'a_x', 'b': 'b_x'})

ソート(sort)

降順にしたいときは、reverse=Trueにします。

df = df.sort("a", descending=True)

また、カラムごとに昇順降順を指定することができます。

df = df.sort([pl.col("a"), pl.col("b")],descending=[False, True])

groupby.agg

指定したグループごとに、様々な演算を施した値を取得できます。

新しく作るカラム名はsuffixかaliasが使えます。
これ以外にも、多くの関数があるので、公式ドキュメント(https://pola-rs.github.io/polars/py-polars/html/reference/dataframe/groupby.html)を参考にしてください。

data = {"group": ["a", "a", "a", "b", "b", "b"], "value": [1, 2, 3, 4, 5, 6]}
df = pl.DataFrame(data)

df >>>
┌───────┬───────┐
│ group ┆ value │
│ ---   ┆ ---   │
│ str   ┆ i64   │
╞═══════╪═══════╡
│ a     ┆ 1     │
│ a     ┆ 2     │
│ a     ┆ 3     │
│ b     ┆ 4     │
│ b     ┆ 5     │
│ b     ┆ 5     │
└───────┴───────┘
df = df.group_by("group", maintain_order=True).agg(
        [
            pl.sum("value").name.suffix("_sum"),   # groupごとにvalueを合算
            pl.mean("value").name.suffix("_mean"), # groupごとのvalueの平均
            pl.count("value").alias("count"), # groupごとのvalueの数
            pl.n_unique("value").alias("unique_count"), # groupごとのユニークなvalueの数
        ]
)

>>> df
┌───────┬───────────┬────────────┬───────┬──────────────┐
│ group ┆ value_sum ┆ value_mean ┆ count ┆ unique_count │
│ ---   ┆ ---       ┆ ---        ┆ ---   ┆ ---          │
│ str   ┆ i64       ┆ f64        ┆ u32   ┆ u32          │
╞═══════╪═══════════╪════════════╪═══════╪══════════════╡
│ a     ┆ 6         ┆ 2.0        ┆ 3     ┆ 3            │
│ b     ┆ 14        ┆ 4.666667   ┆ 3     ┆ 2            │
└───────┴───────────┴────────────┴───────┴──────────────┘

sortと組み合わせて、グループごとにvalueの高いidを抜き出す、ということもできます。

data = {"group": ["a", "a", "a", "b", "b", "b"], 
        "id": [100, 101, 102, 103, 104, 105], 
        "value": [1, 2, 3, 4, 5, 6]}
df = pl.DataFrame(data)

>>> df
┌───────┬─────┬───────┐
│ group ┆ id  ┆ value │
│ ---   ┆ --- ┆ ---   │
│ str   ┆ i64 ┆ i64   │
╞═══════╪═════╪═══════╡
│ a     ┆ 100 ┆ 1     │
│ a     ┆ 101 ┆ 2     │
│ a     ┆ 102 ┆ 3     │
│ b     ┆ 103 ┆ 4     │
│ b     ┆ 104 ┆ 5     │
│ b     ┆ 105 ┆ 6     │
└───────┴─────┴───────┘
df = (df.sort("value", reverse=True) # valueを降順で並び替える
        .groupby('group').agg([pl.col('id').head(2).alias('labels')]) # groupごとに上からidの値を抜き出し、labelsカラムに格納する
     )

>>> df 
┌───────┬────────────┐
│ group ┆ labels     │
│ ---   ┆ ---        │
│ str   ┆ list[i64]  │
╞═══════╪════════════╡
│ a     ┆ [102, 101] │
│ b     ┆ [105, 104] │
└───────┴────────────┘

列の削除

df = df.drop('col_bool')

欠損値の確認

null_counts = df.select([
    pl.col(name).is_null().sum().alias(name + "_null_count") 
    for name in df.columns
])

# 各カラムのnull数を改行して表示
for column in null_counts.columns:
    print(f"{column}: {null_counts[column][0]}")

欠損値のカラム削除

一つでも欠損値がある行を削除します。

df = df.drop_nulls()

欠損値を埋め(fill_null)

df = df.with_columns(pl.col('missing_cols').fill_null(0))
df = df.with_columns(pl.col('missing_cols').fill_nan(-1))

重複の削除(unique)

df = df.unique(subset=["col1"])

定数の列追加(lit)

df = df.with_columns([
    pl.lit(3.14).alias("pie")
])

apply

data = {"value": ["i ro ha", "ni ho he"]}
df = pl.DataFrame(data)

# valueをスペースで分割しリスト化
df = df.with_columns(pl.col('value').apply(lambda s: s.split()).suffix("_new"))

>>> df
┌──────────┬────────────────────┐
│ value    ┆ value_new          │
│ ---      ┆ ---                │
│ str      ┆ list[str]          │
╞══════════╪════════════════════╡
│ i ro ha  ┆ ["i", "ro", "ha"]  │
│ ni ho he ┆ ["ni", "ho", "he"] │
└──────────┴────────────────────┘

展開(explode)

リストを展開します。applyのところで紹介した、リスト化と組み合わせて使えます。

data = {"group": ["a"], 
        "value": [[1, 2, 3]]}
df = pl.DataFrame(data)

>>> df
┌───────┬───────────┐
│ group ┆ value     │
│ ---   ┆ ---       │
│ str   ┆ list[i64] │
╞═══════╪═══════════╡
│ a     ┆ [1, 2, 3] │
└───────┴───────────┘
df = df.explode("value")

>>> df
┌───────┬───────┐
│ group ┆ value │
│ ---   ┆ ---   │
│ str   ┆ i64   │
╞═══════╪═══════╡
│ a     ┆ 1     │
│ a     ┆ 2     │
│ a     ┆ 3     │
└───────┴───────┘

カラム同士を演算

data = {"col1": [1, 2], "col2": [3, 4]}
df = pl.DataFrame(data)

df = df.with_columns((df.get_column("col1") + df.get_column("col2")).alias("col1+col2"))

>>> df
┌──────┬──────┬───────────┐
│ col1 ┆ col2 ┆ col1+col2 │
│ ---  ┆ ---  ┆ ---       │
│ i64  ┆ i64  ┆ i64       │
╞══════╪══════╪═══════════╡
│ 1    ┆ 3    ┆ 4         │
│ 2    ┆ 4    ┆ 6         │
└──────┴──────┴───────────┘

スペース区切りの文章をリスト化

data = {"cols": ["a b c"]}
df = pl.DataFrame(data)

df = df.with_columns(pl.col('cols').apply(lambda s: s.split()).alias("cols_new")) # スペース区切りの文章をリスト化

>>> df
┌───────┬─────────────────┐
│ cols  ┆ cols_new        │
│ ---   ┆ ---             │
│ str   ┆ list[str]       │
╞═══════╪═════════════════╡
│ a b c ┆ ["a", "b", "c"] │
└───────┴─────────────────┘

値をリストで取り出す

col_list = df["col"].to_list()

グループごとに通し番号を振る

data = {"group": ["a", "a", "a", "b", "b", "b"], 
        "value": [1, 2, 3, 4, 5, 6]}
df = pl.DataFrame(data)

df = (df.with_columns([pl.lit(1).alias("n")]) # 後の、cum_sum関数のために、1で埋めたnカラムを作成
        .with_columns(pl.col("n").cum_sum().over("group").alias("serial")) # groupごとの通し番号をserialカラムとして追加
     )

>>> df
┌───────┬───────┬─────┬────────┐
│ group ┆ value ┆ n   ┆ serial │
│ ---   ┆ ---   ┆ --- ┆ ---    │
│ str   ┆ i64   ┆ i32 ┆ i32    │
╞═══════╪═══════╪═════╪════════╡
│ a     ┆ 1     ┆ 1   ┆ 1      │
│ a     ┆ 2     ┆ 1   ┆ 2      │
│ a     ┆ 3     ┆ 1   ┆ 3      │
│ b     ┆ 4     ┆ 1   ┆ 1      │
│ b     ┆ 5     ┆ 1   ┆ 2      │
│ b     ┆ 6     ┆ 1   ┆ 3      │
└───────┴───────┴─────┴────────┘

ISO 8601 フォーマットの変換

df = df.with_columns(
    (pl.col("timestamp").str.strptime(pl.Datetime, "%Y-%m-%dT%H:%M:%S%Z")),
    (pl.col("timestamp").str.strptime(pl.Datetime, "%Y-%m-%dT%H:%M:%S%Z").dt.year().alias("year")),
    (pl.col("timestamp").str.strptime(pl.Datetime, "%Y-%m-%dT%H:%M:%S%Z").dt.month().alias("month")),
    (pl.col("timestamp").str.strptime(pl.Datetime, "%Y-%m-%dT%H:%M:%S%Z").dt.day().alias("day")),
    (pl.col("timestamp").str.strptime(pl.Datetime, "%Y-%m-%dT%H:%M:%S%Z").dt.hour().alias("hour")),
    (pl.col("timestamp").str.strptime(pl.Datetime, "%Y-%m-%dT%H:%M:%S%Z").dt.minute().alias("minute")),
    (pl.col("timestamp").str.strptime(pl.Datetime, "%Y-%m-%dT%H:%M:%S%Z").dt.second().alias("second"))
)
タイトルとURLをコピーしました