21 lines
784 B
Python
21 lines
784 B
Python
import boto3
|
|
import pandas as pd
|
|
from io import BytesIO
|
|
from typing import Optional
|
|
|
|
|
|
def read_parquet_from_s3(bucket: str, key: str, columns: Optional[list] = None) -> pd.DataFrame:
|
|
"""Read a parquet file from S3 into a DataFrame."""
|
|
s3 = boto3.client("s3", region_name="eu-central-1")
|
|
obj = s3.get_object(Bucket=bucket, Key=key)
|
|
return pd.read_parquet(BytesIO(obj["Body"].read()), columns=columns)
|
|
|
|
|
|
def list_keys(bucket: str, prefix: str) -> list[str]:
|
|
"""List all keys under a prefix."""
|
|
s3 = boto3.client("s3", region_name="eu-central-1")
|
|
paginator = s3.get_paginator("list_objects_v2")
|
|
keys = []
|
|
for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
|
|
keys.extend(obj["Key"] for obj in page.get("Contents", []))
|
|
return keys
|