DuckDB SQLクエリをテストする

• 2024-12-28
• DuckDB

はじめに

最近では、比べて数百万行規模の大規模データを高速に処理できるため、データ分析にDuckDB SQLクエリを用いるようになりました。ところで、クエリが複雑になるほど、その動作や結果を検証するのは難しくなります。ここでは、クエリの正確性を確認し、バグを防ぐためにツールを組み合わせて DuckDB SQLクエリをテストするアプローチを紹介します。

サンプルコード

パラメータを埋め込む形でSQLクエリテンプレートをJinjaで定義します。

-- queries/filtering.sql

SELECT 
  * 
FROM {{ table_name }} 
WHERE value > {{ threshold }};

テンプレートエンジンで生成したクエリをDuckDBで実行し、その結果をpandasのDataFrameとして検証します。

以下pytestのコード例では、minijinjaを使っています。

# tests/test_query.py
import duckdb
import pandas as pd
from minijinja import Environment

def test_filtering_query():
    # 入力データをpd.DataFrameで準備
    df = pd.DataFrame(
        data={
            "id": [1, 2, 3],
            "value": [10, 20, 30],
        }
    )

    # SQLテンプレートの読み込みとクエリ生成
    with open('queries/filtering.sql', 'r') as f:
        env = Environment(templates={"filtering_template": f.read()})
    filtering_query = env.render_template(
        "filtering_template",
        table_name="df",
        threshold=15
    )

    # DuckDBでクエリ実行
    actual = duckdb.query(filtering_query).fetchdf()

    # 期待値の定義
    desired = pd.DataFrame({'id': [2, 3], 'value': [20, 30]})

    # 結果を検証
    pd.testing.assert_frame_equal(actual, desired)

まとめ

SQLクエリの検証を効率化するために DuckDB、pandas、pytest、minijinja (Jinja2)を組み合わせる方法を紹介しました。手作業を削減し、テストの高速化を実現するシンプルな手法ですが、SQLロジックの信頼性をさらに向上させてください。