はじめに

Lambdaでは一時領域として/tmpフォルダが提供されていますが、上限512MBのまでの制限があり大容量のCSVファイル生成があった場合に制限がかかっていました。
csvにストリーム出力してS3に保存する方法があるか調べてみました。

変更前

変更前のCSV作成〜s3アップロードは以下になります。
/tmp領域にファイル作成してS3にアップロードしていたので512MBを超えるファイルの作成はできませんでした。

import csv
import boto3
import psycopg2
import psycopg2.extras

def get_conn():
    return psycopg2.connect(
        host='hostname',
        port=5432,
        database='xxxxxx',
        user='xxxxx',
        password='xxxxx',
    )

def get_data():
    with get_conn() as conn:
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            sql = 'SELECT 〜'
            cur.execute(sql)
            rows = cur.fetchall()
            return rows

def upload_s3(src, distname, prefix):
    s3 = boto3.client('s3')
    with open(src, 'rb') as f:
        key = "{}{}".format(prefix, distname)
        s3.put_object(Bucket="BuketName", Key=key, Body=f)
    os.remove(src)

def write_tmp_data():
    with open('/tmp/tmp.csv', 'w', newline='', encoding='utf-8-sig') as f:
        csvfile = csv.writer(f)
        data = get_data()

        csvfile.writerow(['column1', 'column2', 'column3']) # タイトル行出力
        csvfile.writerows(data) # データ書き出し
   upload_s3('/tmp/tmp.csv', 'data.csv', '/tmp')

def handler(event, context):
    write_tmp_data()

変更後

import csv
import boto3
import psycopg2
import psycopg2.extras

def get_conn():
    return psycopg2.connect(
        host='hostname',
        port=5432,
        database='xxxxxx',
        user='xxxxx',
        password='xxxxx',
    )

def get_data():
    with get_conn() as conn:
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            sql = 'SELECT 〜'
            cur.execute(sql)
            rows = cur.fetchall()
            return rows

def upload_s3(buffer, distname, prefix):
    s3 = boto3.client('s3')
    with open(src, 'rb') as f:
        key = "{}{}".format(prefix, distname)
        s3.put_object(Bucket="BuketName", Key=key, Body=buffer)

def write_tmp_data():
    raw = io.BytesIO()
    buffered = io.BufferedWriter(raw)

    with io.TextIOWrapper(buffered, encoding='utf-8-sig', errors='strict', newline='') as f:
        csvfile = csv.writer(f)
        data = get_data()

        csvfile.writerow(['column1', 'column2', 'column3']) # タイトル行出力
        csvfile.writerows(data) # データ書き出し

def handler(event, context):
    write_tmp_data()

io.StringIO（）でもストリームへの書き込みが可能でしたが、エンコードがUTF-8固定となってしまいます。
作成したCSVをExcelで文字化けせずに開けるようにするため、エンコード指定が可能なio. TextIOWrapper()にしました。

最後に

ECSのFargateに移行するかElastic File System(EFS)を導入するか悩みましたが、最小限の修正で大容量ファイルを作成できるようになりました。
あとは、Lambdaのメモリ設定だけを気にすればよいはずです。

参考

第4回　「俺様プチencoding」を実装して理解するPython3.0のioとcodec，encodingの機構：Python 3.0 Hacks｜gihyo.jp … 技術評論社