将Pandas Dataframe写到_csv StringIO而不是文件

时间:2018-06-21 02:41:05

标签: python-3.x pandas dataframe export-to-csv boto3

此代码的目标是将现有的CSV文件从指定的S3存储桶读取到数据帧中,过滤该数据帧以查找所需的列,然后使用StringIO将过滤后的数据帧写入CSV对象,我可以上传到其他S3存储桶。

一切正常, 功能“ prepare_file_for_upload”的代码块除外。下面是完整的代码块:

from io import StringIO
import io #unsued at the moment
import logging
import pandas as pd
import boto3
from botocore.exceptions import ClientError

FORMAT = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
logging.basicConfig(level=logging.INFO, format=FORMAT)
logger = logging.getLogger(__name__)

#S3 parameters
source_bucket = 'REPLACE'
source_folder = 'REPLACE/'
dest_bucket = 'REPLACE'
dest_folder = 'REPLACE'
output_name = 'REPLACE'

def get_file_name():
try:
    s3 = boto3.client("s3")
    logging.info(f'Determining filename from: {source_bucket}/{source_folder}')
    bucket_path = s3.list_objects(Bucket=source_bucket, Prefix=source_folder)
    file_name =[key['Key'] for key in bucket_path['Contents']][1]
    logging.info(file_name)
    return file_name
except ClientError as e:
    logging.info(f'Unable to determine file name from bucket {source_bucket}/{source_folder}')
    logging.info(e)

def get_file_data(file_name):
try:
    s3 = boto3.client("s3")
    logging.info(f'file name from get data: {file_name}')
    obj = s3.get_object(Bucket=source_bucket, Key=file_name)
    body = obj['Body']
    body_string = body.read().decode('utf-8')
    file_data = pd.read_csv(StringIO(body_string))
    #logging.info(file_data)
    return file_data
except ClientError as e:
    logging.info(f'Unable to read {file_name} into datafame')
    logging.info(e)

def filter_file_data(file_data):
try:
    all_columns = list(file_data.columns)
    columns_used = ('col_1', 'col_2', 'col_3')
    desired_columns = [x for x in all_columns if x in columns_used]
    filtered_data = file_data[desired_columns]
    logging.info(type(filtered_data)) #for testing
    return filtered_data
except Exception as e:
    logging.info('Unable to filter file')
    logging.info(e)

我在下面的块中尝试使用StringIO使用“ to_csv”方法而不是创建本地文件来编写传递给函数的现有DF。 to_csv将写入本地文件,但不能与缓冲区一起使用(是的,我尝试将缓冲区游标放置到之后的起始位置,但仍然没有)

def prepare_file_for_upload(filtered_data): #this is the function block where I am stuck
try:
    buffer = StringIO()
    output_name = 'FILE_NAME.csv'
    #code below is writing to file but can not get to write to buffer
    output_file = filtered_data.to_csv(buffer, sep=',')
    df = pd.DataFrame(buffer) #for testing
    logging.info(df) #for testing
    return output_file
except Exception as e:
    logging.info(f'Unable to prepare {output_name} for upload')
    logging.info(e)

def upload_file(adjusted_file):
try:
    #dest_key = f'{dest_folder}/{output_name}'
    dest_key = f'{output_name}'
    s3 = boto3.resource('s3')
    s3.meta.client.upload_file(adjusted_file, dest_bucket, dest_key)
except ClientError as e:
    logging.info(f'Unable to upload {output_name} to {dest_key}')
    logging.info(e)

def execute_program():
file_name = get_file_name()
file_data = get_file_data(file_name)
filtered_data = filter_file_data(file_data)
adjusted_file = prepare_file_for_upload(filtered_data)
upload_file = upload_file(adjusted_file)

if __name__ == '__main__':
execute_program()

1 个答案:

答案 0 :(得分:1)

以下解决方案对我有用:

'csv_buffer = StringIO()' 'output_file = filtered_data.to_csv(csv_buffer)
 s3_resource = boto3.resource('s3')' 's3_resource.Object(dest_bucket,' output_name).
 put(Body=csv_buffer.getvalue())'