[小ネタ] PythonでS3 Selectを使用するときは、結果を連結するように書かかないと正しいデータを取得できない

サーバーレス開発部の夏目です。

案件の中でS3のデータからS3 Selectを使って一部抜き出して、JSON化して読み込むということをPythonでしていたのですが、エラーが出たり出なかったりしていたので調査してみました。

問題になったコード

def execute_s3_select(option):
  resp = s3.select_object_content(**option)
  text = None
  for event in resp['Payload']:
      if 'Records' in event:
          raw = event['Records']['Payload'].decode('utf-8')
          text = raw
  return text

S3 Selectのサンプルコードなどと同じように、返り値のPayloadをループで回してRecordsの内容を代入している。

このコードで得た結果をファイルに書き出したら、以下のようになった。
(実際の結果はもっと大きいものだったけど、わかりやすいように小さめのデータで表現している)

ey003":"3"},{"key001":"ttt","key002":"hey","key003":"4"}

何故かデータが途中からしか取得できていない。 それ故にJSONの形式が崩れてしまっている。

Payloadの中身ってどうなってるんだろう?

def default_proc(obj):
    if isinstance(obj, bytes):
        return obj.decode('utf-8')
    return obj

def write_payload(option, path):
    resp = s3.select_object_content(**option)
    json.dump(list(resp['Payload']), open(path, mode='w'), indent=2, default=default_proc)

S3 Selectの結果のPayloadをループで回してデータを取得できているならリスト化できるんじゃないの?って考えた。
途中からデータしか取得できなかった際のoptionを使って、Payload全体を保存してみた。
bytes型のデータが含まれてて変換が必要になる、ということもあったけど無事リスト化して、出力することができた。

以下がその結果(実際の結果はもっとサイズが大きいものだったけど、見やすくするために小さいデータにした)。

[
  {
    "Records": {
      "Payload": "{\"key001\":\"aaa\",\"key002\":\"say\",\"key003\":\"2\"},{\"key001\":\"bbb\",\"key002\":\"hey\",\"k"
    }
  },
  {
    "Records": {
      "Payload": "ey003\":\"3\"},{\"key001\":\"ttt\",\"key002\":\"hey\",\"key003\":\"4\"},"
    }
  },
  {
    "Stats": {
      "Details": {
        "BytesScanned": 197599,
        "BytesProcessed": 197599,
        "BytesReturned": 115275
      }
    }
  },
  {
    "End": {}
  }
]

どうも、結果のサイズが一定以上だと複数のRecordsに分割されるみたい。
毎回代入してるから、最後のRecordsしか存在しないことになっていた。
となると、S3 Selectからのデータを正しく扱うためにはRecordsの内容を連結する必要がある。

コードを修正してみた

def execute_s3_select(option):
  resp = s3.select_object_content(**option)
  text = ''
  for event in resp['Payload']:
      if 'Records' in event:
          raw = event['Records']['Payload'].decode('utf-8')
          text += raw
  return text

データを連結するように書いてみた。 これでS3 Selectで問題なくデータを取得できるようになった。

まとめ

S3 Selectからデータを取得してJSONに加工して読み込む際、エラーが出たり出なかったりしていたので困っていました。
が、データ量が大きいときでも問題なく読み込めるようになったので安心できました。