AWS LambdaでDuckDBとAWS Data Wrangler、Polarsの処理性能を比較してみた

#!/usr/bin/env node
import * as cdk from "aws-cdk-lib";
import { LambdaStack } from "../lib/lambda-stack";

const app = new cdk.App();

new LambdaStack(app, `CmKasamaLambdaStack`, {
	description: `CmKasamaLambdaStack`,
	env: {
		account: process.env.CDK_DEFAULT_ACCOUNT,
		region: process.env.CDK_DEFAULT_REGION,
	},
	tags: {
		Repository: `CmKasamaLambdaStack-test-tag`,
	},

	projectName: "cm-kasama",
	envName: "dev",
});

import { Construct } from "constructs";
import * as cdk from "aws-cdk-lib";
import * as lambda from "aws-cdk-lib/aws-lambda";
import * as iam from "aws-cdk-lib/aws-iam";
import * as path from "path";
import { PythonLayerVersion } from "@aws-cdk/aws-lambda-python-alpha";

export interface LambdaStackProps extends cdk.StackProps {
	envName: string;
	projectName: string;
}

export class LambdaStack extends cdk.Stack {
	constructor(scope: Construct, id: string, props: LambdaStackProps) {
		super(scope, id, props);

		const lambdaRole = new iam.Role(this, "LambdaExecutionRole", {
			assumedBy: new iam.ServicePrincipal("lambda.amazonaws.com"),
			managedPolicies: [
				iam.ManagedPolicy.fromAwsManagedPolicyName(
					"service-role/AWSLambdaBasicExecutionRole",
				),
				iam.ManagedPolicy.fromAwsManagedPolicyName("AmazonS3FullAccess"),
			],
		});

		const awswranglerLambdaName = `${props.projectName}-${props.envName}-awswrangler-handler`;
		const duckdbLambdaName = `${props.projectName}-${props.envName}-duckdb-handler`;
		const polarsLambdaName = `${props.projectName}-${props.envName}-polars-handler`;
		const source_bucket = "<your_source_bucket>";
		const filename = "medium_sample";
		const source_key = `src/${filename}.csv`;
		const destination_bucket = "<your_destination_bucket>";
		const destination_key = `target/${filename}/${filename}.parquet`;

		const duckdbLayer = new PythonLayerVersion(this, "DuckdbLayer", {
			entry: path.join(__dirname, "../layers/duckdb"),
			compatibleRuntimes: [lambda.Runtime.PYTHON_3_13],
			compatibleArchitectures: [lambda.Architecture.ARM_64],
			description: "Layer containing DuckDB",
			layerVersionName: `${props.projectName}-${props.envName}-duckdb-layer`,
			bundling: {
				assetHashType: cdk.AssetHashType.SOURCE,
				outputPathSuffix: "python",
			},
		});

		const polarsLayer = new PythonLayerVersion(this, "PolarsLayer", {
			entry: path.join(__dirname, "../layers/polars"),
			compatibleRuntimes: [lambda.Runtime.PYTHON_3_13],
			compatibleArchitectures: [lambda.Architecture.ARM_64],
			description: "Layer containing Polars",
			layerVersionName: `${props.projectName}-${props.envName}-polars-layer`,
			bundling: {
				assetHashType: cdk.AssetHashType.SOURCE,
				outputPathSuffix: "python",
			},
		});

		new lambda.Function(this, "AwswranglerLambda", {
			functionName: awswranglerLambdaName,
			runtime: lambda.Runtime.PYTHON_3_13,
			code: lambda.Code.fromAsset("../resources/lambda"),
			handler: "aws_wrangler_test.lambda_handler",
			memorySize: 10240,
			ephemeralStorageSize: cdk.Size.gibibytes(10),
			timeout: cdk.Duration.seconds(900),
			role: lambdaRole,
			layers: [
				lambda.LayerVersion.fromLayerVersionArn(
					this,
					"AwsSdkPandasLayer",
					"arn:aws:lambda:ap-northeast-1:336392948345:layer:AWSSDKPandas-Python313-Arm64:1",
				),
			],
			environment: {
				source_bucket: source_bucket,
				source_key: source_key,
				destination_bucket: destination_bucket,
				destination_key: destination_key,
			},
			architecture: lambda.Architecture.ARM_64,
		});

		new lambda.Function(this, "DuckdbLambda", {
			functionName: duckdbLambdaName,
			runtime: lambda.Runtime.PYTHON_3_13,
			code: lambda.Code.fromAsset("../resources/lambda"),
			handler: "duckdb_test.lambda_handler",
			memorySize: 10240,
			ephemeralStorageSize: cdk.Size.gibibytes(10),
			timeout: cdk.Duration.seconds(900),
			role: lambdaRole,
			layers: [duckdbLayer],
			environment: {
				source_bucket: source_bucket,
				source_key: source_key,
				destination_bucket: destination_bucket,
				destination_key: destination_key,
			},
			architecture: lambda.Architecture.ARM_64,
		});

		new lambda.Function(this, "polarsLambda", {
			functionName: polarsLambdaName,
			runtime: lambda.Runtime.PYTHON_3_13,
			code: lambda.Code.fromAsset("../resources/lambda"),
			handler: "polars_test.lambda_handler",
			memorySize: 10240,
			ephemeralStorageSize: cdk.Size.gibibytes(10),
			timeout: cdk.Duration.seconds(900),
			role: lambdaRole,
			layers: [polarsLayer],
			environment: {
				source_bucket: source_bucket,
				source_key: source_key,
				destination_bucket: destination_bucket,
				destination_key: destination_key,
			},
			architecture: lambda.Architecture.ARM_64,
		});
	}
}

import os
import time
import json
import awswrangler as wr

def lambda_handler(event, context):
    start_time = time.time()

    # S3バケット情報を取得
    source_bucket = os.environ.get("source_bucket")
    source_key = os.environ.get("source_key")
    destination_bucket = os.environ.get("destination_bucket")
    destination_key = os.environ.get("destination_key")

    source_path = f"s3://{source_bucket}/{source_key}"
    destination_path = f"s3://{destination_bucket}/{destination_key}"

    try:
        # 処理開始をログ出力
        print(f"処理開始: {source_key}")

        # AWS Data Wranglerを使用してCSVを直接読み込み、Parquetとして保存
        df = wr.s3.read_csv(source_path)
        wr.s3.to_parquet(
            df=df,
            path=destination_path,
            compression="snappy",  # 圧縮方式を明示的に指定
        )

        execution_time = time.time() - start_time
        print(f"処理完了: 実行時間 {execution_time:.2f}秒")

        return {
            "statusCode": 200,
            "body": json.dumps(
                {
                    "message": "CSV to Parquet conversion completed successfully with AWS Data Wrangler",
                    "execution_time": execution_time,
                    "source": source_path,
                    "destination": destination_path,
                }
            ),
        }

    except Exception as e:
        print(f"エラー発生: {str(e)}")
        raise e

import os
import time
import json
import duckdb

def lambda_handler(event, context):
    start_time = time.time()

    # S3バケット情報を取得
    source_bucket = os.environ.get("source_bucket")
    source_key = os.environ.get("source_key")
    destination_bucket = os.environ.get("destination_bucket")
    destination_key = os.environ.get("destination_key")

    source_path = f"s3://{source_bucket}/{source_key}"
    destination_path = f"s3://{destination_bucket}/{destination_key}"

    try:
        # 処理開始をログ出力
        print(f"処理開始: {source_key}")

        # DuckDBを使用してCSVをParquetに変換（S3から直接読み込み、S3に直接書き込み）
        con = duckdb.connect(database=":memory:")
        con.execute("SET home_directory='/tmp'")  # Lambdaの一時ディレクトリに設定

        # httpfs拡張をインストールしてロード
        con.execute("INSTALL httpfs;")
        con.execute("LOAD httpfs;")

        # CSVをParquetに変換（圧縮指定あり）
        con.execute(f"""
            COPY (SELECT * FROM read_csv_auto('{source_path}', parallel=True)) 
            TO '{destination_path}' (FORMAT PARQUET, COMPRESSION 'SNAPPY');
        """)
        con.close()

        execution_time = time.time() - start_time
        print(f"処理完了: 実行時間 {execution_time:.2f}秒")

        return {
            "statusCode": 200,
            "body": json.dumps(
                {
                    "message": "CSV to Parquet conversion completed successfully with DuckDB",
                    "execution_time": execution_time,
                    "source": source_path,
                    "destination": destination_path,
                }
            ),
        }

    except Exception as e:
        print(f"エラー発生: {str(e)}")
        raise e

import os
import time
import json
import polars as pl

def lambda_handler(event, context):
    start_time = time.time()
    # 専用のサブディレクトリを作成し、適切な権限を設定
    tmp_dir = "/tmp/polars_cache"

    # 環境変数を設定
    os.environ["POLARS_TEMP_DIR"] = tmp_dir

    # S3バケット情報を取得
    source_bucket = os.environ.get("source_bucket")
    source_key = os.environ.get("source_key")
    destination_bucket = os.environ.get("destination_bucket")
    destination_key = os.environ.get("destination_key")

    source_path = f"s3://{source_bucket}/{source_key}"
    destination_path = f"s3://{destination_bucket}/{destination_key}"

    try:
        # 処理開始をログ出力
        print(f"処理開始: {source_key}")

        # S3から直接CSVをスキャン
        df_lazy = pl.scan_csv(source_path, rechunk=True, low_memory=True)

        # S3に直接Parquetとして保存
        df_lazy.sink_parquet(destination_path, compression="snappy")

        execution_time = time.time() - start_time
        print(f"処理完了: 実行時間 {execution_time:.2f}秒")

        return {
            "statusCode": 200,
            "body": json.dumps(
                {
                    "message": "CSV to Parquet conversion completed successfully with Polars LazyFrame",
                    "execution_time": execution_time,
                    "source": source_path,
                    "destination": destination_path,
                }
            ),
        }

    except Exception as e:
        print(f"エラー発生: {str(e)}")
        return {"statusCode": 500, "body": json.dumps({"error": str(e)})}

import pandas as pd
import random
import os
import sys

def create_sample_csv(filename="sample_data.csv", rows=1000000):
    """
    より一般的な文字列データを含むCSVファイルを作成する関数

    Parameters:
    filename (str): 作成するCSVファイルの名前
    rows (int): 行数
    """
    print(f"{rows}行のデータを生成中...")

    # 会社名のサンプル
    companies = [
        "Acme Corp",
        "Globex",
        "Initech",
        "Umbrella Corp",
        "Stark Industries",
        "Wayne Enterprises",
        "Cyberdyne Systems",
        "Soylent Corp",
        "Massive Dynamic",
        "Weyland-Yutani",
        "Oscorp",
        "Gekko & Co",
        "Wonka Industries",
        "Duff Brewing",
    ]

    # 製品カテゴリ
    categories = [
        "Electronics",
        "Clothing",
        "Food",
        "Home Goods",
        "Office Supplies",
        "Sporting Goods",
        "Toys",
        "Books",
        "Health",
        "Automotive",
    ]

    # 国名
    countries = [
        "USA",
        "Canada",
        "UK",
        "Germany",
        "France",
        "Japan",
        "China",
        "Australia",
        "Brazil",
        "India",
        "Mexico",
        "Spain",
        "Italy",
        "Russia",
    ]

    # 都市名
    cities = [
        "New York",
        "London",
        "Tokyo",
        "Paris",
        "Berlin",
        "Sydney",
        "Toronto",
        "Shanghai",
        "Mumbai",
        "Rio de Janeiro",
        "Mexico City",
        "Madrid",
        "Rome",
        "Moscow",
    ]

    # メールドメイン
    email_domains = [
        "gmail.com",
        "yahoo.com",
        "hotmail.com",
        "outlook.com",
        "company.com",
        "example.org",
        "business.net",
        "mail.co",
        "webmail.info",
        "protonmail.com",
    ]

    # ランダムデータの生成
    data = {
        "customer_id": [f"CUST-{random.randint(10000, 99999)}" for _ in range(rows)],
        "full_name": [
            f"{random.choice(['John', 'Jane', 'Robert', 'Mary', 'David', 'Sarah', 'Michael', 'Lisa', 'James', 'Emily'])} {random.choice(['Smith', 'Johnson', 'Williams', 'Jones', 'Brown', 'Davis', 'Miller', 'Wilson', 'Moore', 'Taylor', 'Anderson', 'Thomas', 'Jackson', 'White'])}"
            for _ in range(rows)
        ],
        "email": [
            f"user{random.randint(100, 999)}@{random.choice(email_domains)}"
            for _ in range(rows)
        ],
        "phone": [
            f"+{random.randint(1, 9)}-{random.randint(100, 999)}-{random.randint(100, 999)}-{random.randint(1000, 9999)}"
            for _ in range(rows)
        ],
        "company": [random.choice(companies) for _ in range(rows)],
        "purchase_amount": [
            f"${random.randint(10, 1000)}.{random.randint(0, 99):02d}"
            for _ in range(rows)
        ],
        "category": [random.choice(categories) for _ in range(rows)],
        "country": [random.choice(countries) for _ in range(rows)],
        "city": [random.choice(cities) for _ in range(rows)],
        "status": [
            random.choice(["Completed", "Pending", "Cancelled", "Refunded", "Shipped"])
            for _ in range(rows)
        ],
    }

    # DataFrameの作成
    df = pd.DataFrame(data)

    # CSVに保存
    df.to_csv(filename, index=False)

    # ファイルサイズの確認
    file_size = os.path.getsize(filename) / (1024 * 1024)  # MBに変換
    print(f"ファイル '{filename}' の作成が完了しました。")
    print(f"行数: {rows}、列数: {len(df.columns)}")
    print(f"ファイルサイズ: {file_size:.2f} MB")

    return filename

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("使用方法: python create_sample_data.py <ファイル名> <行数>")
        sys.exit(1)

    filename = sys.argv[1]
    rows = int(sys.argv[2])

    create_sample_csv(filename, rows)