【DuckDB】DuckDBを用いてS3のデータを操作するLambdaを作ってみた

はじめに

以前、DuckDBを用いて、S3バケット上に保管されているALBアクセスログをクエリする取り組みをしました。

しかし、この記事ではCloud9上で実施しており、Lambda上で実施したいと考えていました。本記事では、Lambdaで実現する際に注意が必要な点について説明できればと思います。

今回やること

S3バケットにあるcsvデータを取得し、クエリを実行をした結果をDataFrameで出力するところまでを実施してみたいと思います。

使用するcsvデータはKaggleの有名なデータセットであるTitanicです。

Lambda Layerについて

今回、DuckDBライブラリを利用するため、ライブラリをレイヤーに追加する必要があります。しかし、本記事ではLambdaの実装に関する説明をしたいので割愛させていただきます。

レイヤー追加に関しては以下記事で手順などを説明していますので、興味のある方はこちらをご覧ください。

作成したコード

以下のようなコードを作成しました。

長いため折り畳み

import json
import duckdb
import os
import pandas as pd

def setup_duckdb_connection():
    """DuckDBの接続と設定を行う"""
    # Lambdaの一時ディレクトリを使用する
    temp_dir = '/tmp'
    os.environ['HOME'] = temp_dir
    
    # DuckDBコネクションを作成
    conn = duckdb.connect(database=':memory:')
    
    # 設定を適用
    conn.execute(f"SET home_directory='{temp_dir}';")
    conn.execute("INSTALL httpfs;")
    conn.execute("LOAD httpfs;")
    conn.execute("SET s3_region='ap-northeast-1';")
    
    return conn

def execute_s3_query(conn, bucket, file_path, limit):
    """S3のデータに対してクエリを実行する"""
    query = f"""
    SELECT PassengerID, Survived, Fare 
    FROM read_csv_auto('s3://{bucket}/{file_path}') 
    WHERE Fare > 30 
    LIMIT {limit};
    """
    
    return conn.execute(query).fetchdf()

def lambda_handler(event, context):
    try:
        # リクエストからパラメータ取得
        bucket = event.get('bucket', 'duckdb-demo-qiita')
        file_path = event.get('file_path', 'train.csv')
        limit = event.get('limit', 10)
        
        # DuckDB接続のセットアップと実行
        conn = setup_duckdb_connection()
        df = execute_s3_query(conn, bucket, file_path, limit)
        
        print(df)
        
        return {
            'statusCode': 200,
            'body': json.dumps({'message': 'Query executed successfully'})
        }
    
    except Exception as e:
        print(f"エラーが発生しました: {str(e)}")
        return {
            'statusCode': 500,
            'body': json.dumps({'error': str(e)})
        }

実行結果

きちんとクエリ結果を取得できていそうです。

PassengerId  Survived      Fare
0            2         1   71.2833
1            4         1   53.1000
2            7         0   51.8625
3           10         1   30.0708
4           14         0   31.2750
5           24         1   35.5000
6           26         1   31.3875
7           28         0  263.0000
8           32         1  146.5208
9           35         0   82.1708

注意が必要な点

DuckDBからS3バケットにアクセスするためにhttpfs拡張機能をインストール・ロードする必要があります。

その際、ホームディレクトリの設定が必要であることに注意する必要があります。

これは以下で説明するLambdaのローカルファイルシステムの制約があるためです。

Lambdaのローカルファイルシステムの制約

Lambda関数内では、/tmpディレクトリのみ書き込み可能となっています。これは、Lambda関数が実行される間だけ存在する一時的なストレージで、エフェメラルストレージと呼ばれます。
（デフォルトでは、512MBのサイズとなっており、最大10GBまで設定することができます。ただし、デフォルトから追加したエフェメラルストレージに関しては料金が発生します。）

DuckDBは、デフォルトではHOMEディレクトリ以下にファイル書き込みしようとするため、HOMEディレクトリを/tmpに設定する必要があります。

まとめ

• DuckDBからS3バケットにアクセスするためにhttpfs拡張機能をインストール・ロードする必要がある
• DuckDBは、デフォルトではHOMEディレクトリ以下にファイル書き込みしようとする
• Lambda上では、/tmpディレクトリのみ書き込み可能である
• 以上より、HOMEディレクトリを/tmpに設定する必要がある

おわりに

本記事では、DuckDBをLambdaで扱うことに取り組みました。少し設定が必要なものの、無事S3バケットのデータを取り扱うことができました。
今後、他のライブラリを利用するときにも今回の知識が活かせるのではないかと思います。
ありがとうございました。