Rust LambdaからS3 Selectを実行するPython Lambdaを呼び出す

こんにちは。サービスグループの武田です。

Rust + AWS CLIとPythonとJavaScriptからS3 Selectを実行したパフォーマンス計測をしてみました。

最後に「S3 Selectを実行するPython Lambda関数をRustから呼び出した方が早そう」なんて書いてしまったので、じゃあやったろやないかい。ってことでやってみました。

やってみた

先にPythonのLambda関数を作成しておきます。コードは先述したエントリのものを少し改変したものです。また関数名はs3-select-sample-pythonとしました（s3:GetObject権限が必要です）。

import boto3
import json

s3 = boto3.client("s3")


def lambda_handler(event, context):
    params = {
        "Bucket": "testdata-xxxx",
        "Key": "test_data.json",
        "InputSerialization": {
            "JSON": {
                "Type": "LINES",
            }
        },
        "OutputSerialization": {
            "JSON": {
                "RecordDelimiter": "\n",
            }
        },
        "ExpressionType": "SQL",
        "Expression": "SELECT * FROM s3object s",
    }

    res = s3.select_object_content(**params)

    text = ""
    for event in res["Payload"]:
        if "Records" in event:
            raw = event["Records"]["Payload"].decode("UTF-8")
            text += raw

    return text

次にRustプロジェクトを作成していきます。適当なディレクトリに移動してコマンドを実行しましょう。

$ cargo new rust-lambda-call-python-lambda

プロジェクトが作成できたら、必要な依存関係を書いておきます（dependencies以外は省略）。

[dependencies]
lamedh_runtime = "0.3"
tokio = { version = "1", features = ["full"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
rusoto_core = "0.46"
rusoto_lambda = "0.46"
log = "0.4"
env_logger = "0.8"

コードは先日のものをベースにLambda関数を呼び出すように書き換えます。なおLambdaをinvokeした戻り値は"でくくられた上エスケープされているため、文字列操作で取り除く必要があります。

use lamedh_runtime::{Context, Error};
use log::{debug, info};
use rusoto_lambda::{InvocationRequest, Lambda, LambdaClient};
use serde::Deserialize;
use serde_json::Value;
use std::env;

#[derive(Debug, Deserialize)]
struct TestData {
    name: String,
    code: u32,
    tags: Option<String>,
    lang: Option<String>,
}

#[tokio::main]
async fn main() -> Result<(), Error> {
    env::set_var("RUST_LOG", "rust_lambda_call_python_lambda=debug");
    env_logger::init();

    lamedh_runtime::run(lamedh_runtime::handler_fn(handler)).await?;

    Ok(())
}

async fn handler(_: Value, _: Context) -> Result<(), Error> {
    debug!("handler start");

    let lambda_client = LambdaClient::new(Default::default());

    let res = lambda_client
        .invoke(InvocationRequest {
            function_name: "s3-select-sample-python".into(),
            ..Default::default()
        })
        .await?;

    let text = unwrap_string_quote(String::from_utf8(res.payload.unwrap().to_vec())?);

    debug!("{:?}", text);

    for line in text.lines() {
        let d: TestData = serde_json::from_str(line)?;
        info!("{:?}", d);
    }

    Ok(())
}

fn unwrap_string_quote(s: String) -> String {
    let mut s: String = s.replace("\\\"", "\"").replace("\\n", "\n");
    s.remove(0);
    s.remove(s.len() - 1);
    s
}

コードが準備できたらsoftprops/lambda-rustを利用してLambda用にコンパイルします。

$ docker run --rm \
  -v $PWD:/code \
  -v $HOME/.cargo/registry:/root/.cargo/registory \
  -v $HOME/.cargo/git:/root/.cargo/git \
  softprops/lambda-rust

コンパイルできたら、Lambda関数を作成してzipファイルをアップロードします。なお、ロールにはlambda:InvokeFunction権限が必要ですので、忘れずに付与しておきましょう。

注目の実行結果はこちら！

START RequestId: a3fee59d-1711-4eb4-8b8e-d6b8159f833c Version: $LATEST
[2021-02-26T09:20:59Z DEBUG rust_lambda_call_python_lambda] handler start
[2021-02-26T09:20:59Z DEBUG rust_lambda_call_python_lambda] "{\"name\":\"Test\",\"code\":1939,\"tags\":\"Dev\",\"lang\":\"ja\"}\n{\"name\":\"IT Division\",\"code\":1,\"tags\":\"Prod\",\"lang\":\"ja\"}\n{\"name\":\"Sample\",\"code\":31,\"lang\":\"en\"}\n{\"name\":\"Classmethod\",\"code\":2,\"lang\":\"en\"}\n{\"name\":\"Classmethod2\",\"code\":19}\n"
[2021-02-26T09:20:59Z INFO  rust_lambda_call_python_lambda] TestData { name: "Test", code: 1939, tags: Some("Dev"), lang: Some("ja") }
[2021-02-26T09:20:59Z INFO  rust_lambda_call_python_lambda] TestData { name: "IT Division", code: 1, tags: Some("Prod"), lang: Some("ja") }
[2021-02-26T09:20:59Z INFO  rust_lambda_call_python_lambda] TestData { name: "Sample", code: 31, tags: None, lang: Some("en") }
[2021-02-26T09:20:59Z INFO  rust_lambda_call_python_lambda] TestData { name: "Classmethod", code: 2, tags: None, lang: Some("en") }
[2021-02-26T09:20:59Z INFO  rust_lambda_call_python_lambda] TestData { name: "Classmethod2", code: 19, tags: None, lang: None }
END RequestId: a3fee59d-1711-4eb4-8b8e-d6b8159f833c
REPORT RequestId: a3fee59d-1711-4eb4-8b8e-d6b8159f833c	Duration: 173.65 ms	Billed Duration: 174 ms	Memory Size: 128 MB	Max Memory Used: 39 MB

メモリ128MBでも 173.65ms ！Python Lambdaの実行時間を加味しても（合計実行時間）300msはいかなそうです。これなら悪くないですね。

まとめ

RustでなんとかS3 Selectしたいということから始め、結局Pythonに頼るという結果になりました。いやいや、実現方法は二の次ですよ！ローカル環境であればAWS CLIを外部コマンドとして呼び出す方式もよさそうですが、Lambdaの場合は素直にPythonなどで実装したものを呼び出した方がよさそうです。