从 S3 流式传输 / 分块 csv 到 Python

Question

我打算使用 Python 对存储在 S3 中的一个非常大的 csv 文件执行一些内存密集型操作，目的是将脚本移动到 AWS Lambda。我知道我可以读取整个 csv nto 内存，但我肯定会遇到 Lambda 的内存和存储限制，使用如此大的文件是否有任何方法可以使用 boto3 一次将 csv 的块读入或读入 Python /botocore，最好通过指定要读入的行号？

以下是我已经尝试过的一些事情：

1) 使用S3.get_object 中的range 参数指定要读入的字节范围。不幸的是，这意味着最后一行在中间被截断，因为无法指定要读入的行数. 有一些混乱的解决方法，例如扫描最后一个换行符，记录索引，然后将其用作下一个字节范围的起点，但如果可能的话，我想避免这种笨拙的解决方案。

2) 使用 S3 select 编写 sql 查询以选择性地从 S3 存储桶中检索数据。不幸的是，不支持 row_numbers SQL 函数，而且看起来没有办法读取行子集。

Answer 1

假设您的文件未压缩，这应该涉及从流中读取并在换行符处拆分。读取一块数据，找到该块中换行符的最后一个实例，拆分并处理。

s3 = boto3.client('s3')
body = s3.get_object(Bucket=bucket, Key=key)['Body']

# number of bytes to read per chunk
chunk_size = 1000000

# the character that we'll split the data with (bytes, not string)
newline = '\n'.encode()   
partial_chunk = b''

while (True):
    chunk = partial_chunk + body.read(chunk_size)

    # If nothing was read there is nothing to process
    if chunk == b'':
        break

    last_newline = chunk.rfind(newline)

    # write to a smaller file, or work against some piece of data
    result = chunk[0:last_newline+1].decode('utf-8')

    # keep the partial line you've read here
    partial_chunk = chunk[last_newline+1:]

如果你有 gzip 文件，那么你需要在循环中使用 BytesIO 和 GzipFile 类；这是一个更难的问题，因为您需要保留 Gzip 压缩细节。

Answer 2

我开发了一个类似于@Kirk Broadhurst 的代码，但是如果每个块的处理时间超过 5 分钟（大约），就会发生连接超时。以下代码通过为每个块打开一个新连接来工作。

import boto3
import pandas as pd
import numpy as np

# The following credentials should not be hard coded, it's best to get these from cli.
region_name = 'region'
aws_access_key_id = 'aws_access_key_id'
aws_secret_access_key = 'aws_secret_access_key'

s3 =boto3.client('s3',region_name=region_name,aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key)

obj = s3.get_object(Bucket='bucket', Key='key')

total_bytes = obj['ContentLength']
chunk_bytes = 1024*1024*5 # 5 MB as an example.
floor = int(total_bytes//chunk_bytes)
whole = total_bytes/chunk_bytes
total_chunks = [1+floor if floor<whole else floor][0]

chunk_size_list = [(i*chunk_bytes, (i+1)*chunk_bytes-1) for i in range(total_chunks)]
a,b = chunk_size_list[-1]
b = total_bytes
chunk_size_list[-1] = (a,b)
chunk_size_list = [f'bytes={a}-{b}' for a,b in chunk_size_list]

prev_str = ''

for i,chunk in enumerate(chunk_size_list):
    s3 = boto3.client('s3', region_name=region_name, aws_access_key_id=aws_access_key_id, 
                      aws_secret_access_key=aws_secret_access_key)
    byte_obj = s3.get_object(Bucket='bucket', Key='key', Range=chunk_size_list[i])
    byte_obj = byte_obj['Body'].read()
    str_obj = byte_obj.decode('utf-8')
    del byte_obj
    list_obj = str_obj.split('\n')
    # You can use another delimiter instead of ',' below.
    if len(prev_str.split(',')) < len(list_obj[1].split(',')) or len(list_obj[0].split(',')) < len(list_obj[1].split(',')):
        list_obj[0] = prev_str+list_obj[0]
    else:
        list_obj = [prev_str]+list_obj
    prev_str = list_obj[-1]
    del str_obj, list_obj[-1] 
    list_of_elements = [st.split(',') for st in list_obj]
    del list_obj
    df = pd.DataFrame(list_of_elements)
    del list_of_elements
    gc.collect()
    # You can process your pandas dataframe here, but you need to cast it to correct datatypes.
    # casting na values to numpy nan type.
    na_values = ['', '#N/A', '#N/A N/A', '#NA', '-1.#IND', '-1.#QNAN', '-NaN', '-nan', '1.#IND', '1.#QNAN', 'N/A', 'NA', 'NULL', 'NaN', 'n/a', 'nan', 'null']
    df = df.replace(na_values, np.nan)
    dtypes = {col1: 'float32', col2:'category'}
    df = df.astype(dtype=dtypes, copy=False)