본문 바로가기
AWS

[AWS] python으로 s3에 있는 parquet 파일 읽기

by 초이MS 2022. 2. 3.

목적

  • S3에 저장해놓은 parquet 파일을 AWS ec2 등의 서버가 아닌 개인 컴퓨터에서 쉽게 읽고싶다.
  • aws configure로 계정정보(access key, secret key)를 저장하는 것이 아닌 코드에서 계정정보를 관리하고싶다.
  • 필자는 pyarrow, s3fs가 동작하지 않아 빠르게 이용할 수 있는 다른 방법을 찾고 싶었다.

코드

import boto3
import io
import pandas as pd

...
s3_config = {
    "aws_access_key_id": "{ACCESS KEY}",
    "aws_secret_access_key": "{SECRET KEY}",
    "region_name": "{MY REGION}"
}
...
def pd_read_s3_parquet(key, bucket, s3_client=None, **args):
    if s3_client is None:
        s3_client = boto3.client('s3')
    obj = s3_client.get_object(Bucket=bucket, Key=key)
    return pd.read_parquet(io.BytesIO(obj['Body'].read()), **args)

def pd_read_s3_multiple_parquets(filepath, bucket, s3=None, 
                                 s3_client=None, verbose=False, **args):
    if not filepath.endswith('/'):
        filepath = filepath + '/'  # Add '/' to the end
    if s3_client is None:
        s3_client = boto3.client('s3')
    if s3 is None:
        s3 = boto3.resource('s3')
    s3_keys = [item.key for item in s3.Bucket(bucket).objects.filter(Prefix=filepath)
               if item.key.endswith('.parquet')]
    if not s3_keys:
        print('No parquet found in', bucket, filepath)
    elif verbose:
        print('Load parquets:')
        for p in s3_keys: 
            print(p)
    dfs = [pd_read_s3_parquet(key, bucket=bucket, s3_client=s3_client, **args) 
           for key in s3_keys]
    return pd.concat(dfs, ignore_index=True)

...
session = boto3.session.Session(**s3_config) 
s3_client = session.client('s3')
s3 = session.resource('s3')

# path_to_read="s3://my-bucket/a/b/c"
df = pd_read_s3_multiple_parquets("a/b/c", "my-bucket", s3=s3, s3_client=s3_client)
  • 위의 코드 중 1) s3_config dictionary 변수 부분을 수정 및 2) 최하단의 읽을 파일 경로를 수정하여 a/b/c 하위의 parquet 파일을 pandas dataframe으로 읽을 수 있다.(df 변수)

참고

반응형

댓글