📘 Pythonでpandas と Polars の CSV 読み込み速度を徹底比較

import time
import os
import psutil
import pandas as pd
import polars as pl
from datetime import datetime

# ============================================
# 1. 巨大 CSV を生成（20,000,000 行 × 22 列）
# ============================================
# Polars の DataFrame を使って巨大 CSV を作成します。
# ・id 列（0〜n-1）
# ・num_0〜num_19（id に +i した数値列）
# ・category（id % 10）
# ・text_col（固定文字列）
# 合計 22 列のデータを生成します。

def generate_csv(path, n=20_000_000, num_cols=20):
    # id 列を Series で作成（Expr を避けて安定動作）
    df = pl.DataFrame({
        "id": pl.Series(range(n))
    })

    # 数値列を追加（num_0〜num_19）
    for i in range(num_cols):
        df = df.with_columns((pl.col("id") + i).alias(f"num_{i}"))

    # カテゴリ列（0〜9 の文字列）
    df = df.with_columns(((pl.col("id") % 10).cast(pl.Utf8)).alias("category"))

    # 文字列列（全行 "sample_text"）
    df = df.with_columns(pl.Series(["sample_text"] * n).alias("text_col"))

    # CSV 出力
    df.write_csv(path)


# ============================================
# 2. メモリ使用量取得（MB）
# ============================================
# psutil を使って現在プロセスの RSS を取得します。
# 読み込み前後の差分を取ることで、各ライブラリが
# どれだけメモリを消費したかを測定できます。

def get_memory_mb():
    process = psutil.Process(os.getpid())
    return process.memory_info().rss / 1024 / 1024


# ============================================
# 3. ベンチマーク（読み込み時間＋メモリ使用量）
# ============================================
# pandas / Polars（eager / lazy）の 3 パターンで
# CSV 読み込み時間とメモリ使用量を比較します。

def benchmark_read(csv_path):
    results = {}

    # pandas
    mem_before = get_memory_mb()
    start = time.time()
    df_pd = pd.read_csv(csv_path)
    results["pandas"] = {
        "time": time.time() - start,
        "memory": get_memory_mb() - mem_before
    }

    # Polars eager（即時実行）
    mem_before = get_memory_mb()
    start = time.time()
    df_pl = pl.read_csv(csv_path)
    results["polars_eager"] = {
        "time": time.time() - start,
        "memory": get_memory_mb() - mem_before
    }

    # Polars lazy（遅延実行 → collect で実行）
    mem_before = get_memory_mb()
    start = time.time()
    df_lazy = pl.scan_csv(csv_path).collect()
    results["polars_lazy"] = {
        "time": time.time() - start,
        "memory": get_memory_mb() - mem_before
    }

    return results


# ============================================
# 4. Markdown レポート生成
# ============================================
# 実行結果を Markdown 形式でまとめます。
# ブログやドキュメントにそのまま貼れる形式です。

def generate_markdown(results, csv_path, n, num_cols):
    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    md = f"""
# CSV 読み込みベンチマークレポート
実行日時: {now}

## 📊 対象データ
- 行数: {n:,} 行
- 列数: {num_cols + 2} 列（数値 {num_cols} + category + text）
- ファイル: `{csv_path}`

## ⚡ 読み込み速度比較

| ライブラリ | 時間（秒） | メモリ使用量（MB） |
|-----------|------------|---------------------|
| pandas | {results["pandas"]["time"]:.3f} | {results["pandas"]["memory"]:.1f} |
| polars（eager） | {results["polars_eager"]["time"]:.3f} | {results["polars_eager"]["memory"]:.1f} |
| polars（lazy） | {results["polars_lazy"]["time"]:.3f} | {results["polars_lazy"]["memory"]:.1f} |

"""

    return md


# ============================================
# 5. 実行（CSV 生成 → ベンチマーク → レポート出力）
# ============================================

csv_path = "test_big.csv"
n = 20_000_000
num_cols = 20

# 巨大 CSV を生成
generate_csv(csv_path, n=n, num_cols=num_cols)

# ベンチマーク実行
results = benchmark_read(csv_path)

# Markdown レポート生成
markdown_report = generate_markdown(results, csv_path, n, num_cols)

# ファイル出力
with open("benchmark_report.md", "w", encoding="utf-8") as f:
    f.write(markdown_report)

print("Markdownレポートを生成しました: benchmark_report.md")