ibis-frameworkでPolarsとSQLをつかってみた

Posted at 2024-07-21

最近ibis-frameworkというデータ処理のライブラリを見つけました。
1つのAPIで20以上のバックエンドに対応し、接続先のバックエンドと同じくらい優れたパフォーマンスと、一貫したユーザーエクスペリエンスを実現するよう開発されているそうです。
公式サイトの画像を見ると、pandasやBigqueryなど馴染みのあるものから、最近よく聞くようになったPolarsやSnowflake、その他初めて聞くものもあり幅広く対応しているようです。

https://ibis-project.org/

やりたいこと

Ibisをつかったデータ加工を実際にやってみる
pandas / Polarsそのままでの処理と、IbisのPolarsバックエンドでの処理の実行時間を比較する
SQLも利用できるっぽいので試してみる

Googleトレンドでpandas Polars ibis-frameworkを比較したスクリーンショット↓
Polarsは最近聞くようになったとはいえ、まだまだpandasの方が人気のようです。
Ibisはこれから伸びていくのでしょうか。
お盆、年末年始、GWにpandasの検索が下がっているのはおもしろいですね。

前提

筆者の利用経験は以下の通り

SQL:3年以上
pandas:2年以上
Polars:なし
Ibis:なし

使用データ

架空のスーパーマーケットのデータを2種類

購買履歴データ
ポイントカード会員データ

Polarsの速さを体験したいので大きめのデータを生成し、テーブルの結合や欠損補完も試したいので以下のような処理を加え、CSVでエクスポート、Googleドライブに格納しておく

架空のデータを作成

import pandas as pd
import numpy as np

# データのサイズ設定
num_transactions = 40_000_000
num_members = 100_000

# 購買履歴データの作成
transaction_data = {
    "transaction_id": np.arange(1, num_transactions + 1),
    "customer_id": np.random.randint(1, num_members + 1, size=num_transactions),
    "product_id": np.random.randint(1, 10001, size=num_transactions),
    "quantity": np.random.randint(1, 20, size=num_transactions),
    "price": np.random.randint(10, 1000, size=num_transactions),
    "purchase_date": pd.date_range(start='1/1/2023', periods=num_transactions, freq='S').strftime('%Y%m%d'),
    "purchase_time": pd.date_range(start='1/1/2023', periods=num_transactions, freq='S').strftime('%H:%M:%S')
}

df_transactions = pd.DataFrame(transaction_data)

# 購買履歴データの20%を会員コードが欠損値（会員でない）に設定
non_member_indices = np.random.choice(df_transactions.index, size=int(num_transactions * 0.2), replace=False)
df_transactions.loc[non_member_indices, 'customer_id'] = 0

print('transactions', df_transactions.shape)
display(df_transactions.head())

df_transactions.to_csv('/content/drive/MyDrive/transactions.csv', index=False)

# ポイントカード会員データの作成
member_data = {
    "customer_id": np.arange(1, num_members + 1),
    "sex": np.random.choice(['Male', 'Female'], size=num_members),
    "age": np.random.randint(18, 80, size=num_members),
    "address_code": np.random.randint(100, 999, size=num_members)
}

df_members = pd.DataFrame(member_data)

print('members', df_members.shape)
display(df_members.head())

df_members.to_csv('/content/drive/MyDrive/members.csv', index=False)

それぞれのファイルサイズは以下の通り

データの読み込み・加工・集計の実施

以下の処理をいくつかのやり方で試してみる

作成したCSVデータを読み込み
購買履歴データとポイントカード会員データを結合
欠損を補完
数量と単価から売上金額を算出
指定した条件での売上金額の上位3件の商品IDと金額を抽出
（住所コードが200番台の30代男性でフィルタリング）

pandasそのまま

まずは使い慣れたpandasで実行してみます。

pandas

%%time
import pandas as pd

# CSVファイルを読み込む
transactions = pd.read_csv('/content/drive/MyDrive/transactions.csv')
members = pd.read_csv('/content/drive/MyDrive/members.csv')

# データを結合
df = transactions.merge(members, on='customer_id', how='left')

# ポイントカード会員でない場合の補完
df['sex'].fillna('Unknown', inplace=True)
df['age'].fillna(999, inplace=True)
df['address_code'].fillna(9999, inplace=True)

# 売上金額のカラムを追加
df['sales'] = df['quantity'] * df['price']

# フィルタリングと集計
df_filtered = df[(df['sex'] == 'Male') & (df['age'].between(30, 39)) & (df['address_code'].between(200, 299))]
result = df_filtered.groupby('product_id', as_index=False).agg(total_sales=('sales', 'sum')).sort_values(by='total_sales', ascending=False).reset_index(drop=True).head(3)

# 集計結果を表示
display(result)

1分20秒ほどかかりました。体感としてはちょっと長いかなという感じです。

Polarsそのまま

続いてPolarsを試してみます。Ibisを使わずにpolarsで読み込んで処理します。

polars

%%time
import polars as pl

# CSVファイルを読み込む
transactions = pl.read_csv('/content/drive/MyDrive/transactions.csv')
members = pl.read_csv('/content/drive/MyDrive/members.csv')

# データを結合
df = transactions.join(members, on='customer_id', how='left')

# ポイントカード会員でない場合の補完
df = df.with_columns([
    pl.col('sex').fill_null('Unknown'),
    pl.col('age').fill_null(999),
    pl.col('address_code').fill_null(9999)
])

# 売上金額のカラムを追加
df = df.with_columns((pl.col('quantity') * pl.col('price')).alias('sales'))

# フィルタリングと集計
df_filtered = df.filter(
    (pl.col('sex') == 'Male') &
    (pl.col('age').is_between(30, 39)) &
    (pl.col('address_code').is_between(200, 299))
)

result = df_filtered.group_by('product_id').agg(pl.sum('sales').alias('total_sales')).sort('total_sales', descending=True).limit(3)

# 集計結果を表示
print(result)

30秒ほどで完了しました。pandasの半分以下でやはりPolarsは速いですね。

Ibisを使ってみる（エンジンはPolars, Ibis記法）

次にIbisのPolarsエンジンで実行してみます。記法はpandasのメソッドチェーンやRのdplyrのようなSQLっぽいような感じですね。

ibis[polars]

%%time
import ibis
import polars as pl

# polarsバックエンドに接続
con = ibis.polars.connect()

# CSVファイルを読み込む
transactions = con.read_csv('/content/drive/MyDrive/transactions.csv')
members = con.read_csv('/content/drive/MyDrive/members.csv')

# データの補完と売上金額のカラム追加
transactions_with_members = (
    transactions
    .left_join(members, transactions.customer_id == members.customer_id)
    .mutate(
        sex=ibis.case().when(members.sex.notnull(), members.sex).else_('Unknown').end(),
        age=ibis.case().when(members.age.notnull(), members.age).else_(999).end(),
        address_code=ibis.case().when(members.address_code.notnull(), members.address_code).else_(9999).end(),
        sales=transactions.quantity * transactions.price
    )
)

# フィルタリングと集計
result = (
    transactions_with_members
    .filter(
        (transactions_with_members.sex == 'Male') &
        (transactions_with_members.age.between(30, 39)) &
        (transactions_with_members.address_code.between(200, 299))
    )
    .group_by('product_id')
    .aggregate(total_sales=transactions_with_members.sales.sum())
    .order_by(ibis.desc('total_sales'))
    .limit(3)
    .execute()
)

# 集計結果を表示
display(result)

Polarsそのままで実行したのと同じ30秒ほどで完了しました。

Ibisを使ってみる（エンジンはPolars, SQL）

最後にIbisのPolarsバックエンドをSQLで書いて同じように処理できるか試してみます。
read_csvのオプションでtable_name=を設定する必要があるようです。

ibis[polars] SQLで書く場合


%%time
import ibis
import polars as pl

# polarsバックエンドに接続
con = ibis.polars.connect()

# CSVファイルを読み込む
transactions = con.read_csv('/content/drive/MyDrive/transactions.csv', table_name='transactions')
members = con.read_csv('/content/drive/MyDrive/members.csv', table_name='members')

# SQL
sql = """
with
transactions_with_members as (
    select
        transactions.*
        , case when sex is not null then sex else 'Unknown' end as sex
        , case when age is not null then age else 999 end as age
        , case when address_code is not null then address_code else 9999 end as address_code
        , quantity * price as sales
    from
        transactions as t
        left join
            members as m
        on
            t.customer_id = m.customer_id
),

result as (
select
    product_id
    , sum(sales) as total_sales
from
    transactions_with_members
where
    sex = 'Male'
    and age between 30 and 39
    and address_code between 200 and 299
group by
    product_id
order by
    total_sales desc
limit
    3
)

select * from result
"""

# 結果を実行して取得
result = con.sql(sql).execute()

# 集計結果を表示
display(result)

こちらも30秒ほどとPolars、Ibis[polars]と同様の実行時間で完了しました。
SQLでも同様に処理できるというのはすごいですね。

まとめ

実行コード	実行時間
pandas	79 sec
polars	32 sec
ibis[polars]	31 sec
ibis[polars(SQL)]	30 sec

ということでIbisをつかってpolarsと同様の結果を得られました。
今後もPolarsのようにいろいろなデータ処理のライブラリが出てきた際にはIbisがカバーしてくれるのであれば、Ibisさえ身につけておけばある程度のことは対応できるというふうになっていくのかもしれません。

Ibisならデータ分析プロジェクトが進捗するにつれて、データが大規模になったり、pandasやSQLしかつかったことがないメンバーが参画したりした際にもスムーズに対応できそうですね。

参考

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up