from airflow.exceptions import AirflowFailException
from config import Config
from time import sleep
from .s3 import S3Util
[docs]
class DataLakeUtil:
"""
Athena and glue utilities
"""
[docs]
@staticmethod
def apply_column_mapping(
df, column_mapping: dict[str, str | list[str]] | None
) -> None:
"""
Apply column mapping to a DataFrame in-place.
For each target column, finds the first matching alias in the DataFrame
and renames it to the target name. Validates during application that no
alias appears in multiple target lists.
Args:
df: pandas DataFrame (optional, pass None to validate only)
column_mapping: dict where keys are target column names and values are
alias(es) to look for. Example:
{"md5_pkey_hash": ["md5", "md5_pkey_hash", "md5hash"]}
Raises:
ValueError: if the same alias appears in multiple target lists
"""
if not column_mapping:
return
seen_aliases: dict[str, str] = {}
renames = {}
df_columns_lower = (
{col.lower(): col for col in df.columns} if df is not None else {}
)
for target, aliases in column_mapping.items():
if isinstance(aliases, str):
aliases = [aliases]
for alias in aliases:
alias_lower = alias.lower()
if alias_lower in seen_aliases:
raise ValueError(
f"Column mapping error: alias '{alias}' is mapped to both "
f"'{seen_aliases[alias_lower]}' and '{target}'. Each alias can only appear once."
)
seen_aliases[alias_lower] = target
if df is not None and alias_lower in df_columns_lower:
actual_col = df_columns_lower[alias_lower]
if actual_col != target:
renames[actual_col] = target
print(f"Column mapping: '{actual_col}' -> '{target}'")
break
else:
break
if df is not None and renames:
df.rename(columns=renames, inplace=True)
[docs]
@staticmethod
def get_athena_table_size(table_name: str, db_name: str = Config.DATA_LAKE_DB):
import awswrangler as wr
print(f"Finding size of athena table: {db_name}.{table_name}")
if not wr.catalog.does_table_exist(
table=table_name,
database=db_name,
):
print(f"{db_name=}.{table_name=} doesnt exist")
return -1
return wr.athena.read_sql_query(
f"SELECT count(*) AS cnt FROM {table_name}",
database=db_name,
).iloc[0, 0]
[docs]
@staticmethod
def get_table_path_info(tgt_table, tgt_db=Config.DATA_LAKE_DB):
"""Find the full S3 path, bucket, and prefix of a table in the Data Lake.
If table does not exists, return None"""
import awswrangler as wr
if not wr.catalog.does_table_exist(
database=tgt_db,
table=tgt_table,
):
print(f"Path/metadata info: Table {tgt_db}.{tgt_table} does not exist!")
return None
# Example output: 's3://463002504035-shoreline/curated/lsd/c_lsd'
full_s3_path = wr.catalog.get_table_location(database=tgt_db, table=tgt_table)
infer_path = full_s3_path.split("s3://", 1)[-1]
infer_bucket = infer_path.split("/", 1)[0]
infer_prefix = infer_path.split("/", 1)[-1]
return full_s3_path, infer_bucket, infer_prefix
[docs]
@staticmethod
def drop_athena_partition(partition_key, partition_val, parent_tbname):
"""Drop partition (single partitioning scheme) if EXISTS - both metadata and data"""
import awswrangler as wr
print(
f"Try dropping partition {partition_key}='{partition_val}' in parent table {parent_tbname}..."
)
path_info = DataLakeUtil.get_table_path_info(parent_tbname)
if path_info:
_, s3_bucket, s3_prefix_parent = path_info
s3_prefix_partition = f"{s3_prefix_parent}/{partition_key}={partition_val}/"
print(f"Try removing logical data for {partition_key}='{partition_val}'")
S3Util.delete_s3_folder(bucket=s3_bucket, prefix=s3_prefix_partition)
wr.catalog.delete_partitions(
parent_tbname,
database=Config.DATA_LAKE_DB,
partitions_values=[[partition_val]],
)
print(
f"Successfully dropped partition metadata for {partition_key}='{partition_val}'"
)
[docs]
@staticmethod
def get_athena_columns_defs(
table_name: str,
lake_db: str = Config.DATA_LAKE_DB,
):
"""
Returns columns_types (dictionary of non-partition cols and dtypes), partitions_types (separate partition cols and dtypes)
:rtype tuple[dict[str, str], [dict[str, str]]]: a tuple of size 2: each item is a dictionary of columns and their athena dtypes e.g. {'state': 'varchar(2)', 'volume': 'int', 'fa_ver': 'int', 'stage': 'int'}
This is a working replacement that rfor unreliable read_parquet_metadata https://aws-sdk-pandas.readthedocs.io/en/stable/stubs/awswrangler.s3.read_parquet_metadata.html.
Cons of read_parquet_metadata are:
Unwieldy params: ignore_null, dtype, sampling, ...
It failed to recognize partition cols (partitions_types=None!). Examples:
table c_lsd: sampling=1.0 tries to read entire gigantic table: SLOW;
table r_cai_sbo: due to NULL values in partition cols? "Field expected_size has incompatible types: int32 vs int64"
"""
import awswrangler as wr
from simple_ddl_parser import DDLParser
if table_name and not wr.catalog.does_table_exist(
database=lake_db,
table=table_name,
):
raise Exception(f"Table {lake_db}.{table_name} does not exist!")
ddl_str = wr.athena.show_create_table(database=lake_db, table=table_name)
print(f"Full raw DDL for table {lake_db}.{table_name}: {ddl_str}")
parsed = DDLParser(ddl_str).run(output_mode="athena")
if parsed:
partitions_list = parsed[0].get("partitioned_by", [])
else:
import boto3
glue_client = boto3.client("glue", region_name="us-east-1")
glue_table = glue_client.get_table(DatabaseName=lake_db, Name=table_name)[
"Table"
]
partitions_list = [
{"name": pk["Name"]} for pk in glue_table.get("PartitionKeys", [])
]
all_cols = wr.catalog.get_table_types(database=lake_db, table=table_name)
partitions_types = {}
for dict_item in partitions_list:
colname = dict_item["name"].replace("`", "")
# For getting athena native dtypes, better to use wr.get_table_types (avoid manually parsing dtypes from simple_ddl_parser.DDLParser)!
coltype = all_cols[colname]
partitions_types[colname] = coltype
# Modify columns_types: MUST EXCLUDE partition columns else ERROR; Review: https://youtu.be/2EXMKl9zKTQ?t=100
columns_types = {k: v for k, v in all_cols.items() if k not in partitions_types}
print(
f"""
Columns definitions for {lake_db}.{table_name} are:
{columns_types=}
{partitions_types=}
"""
)
return columns_types, partitions_types
[docs]
@staticmethod
def msck_repair_table(
table_name: str,
database: str = Config.DATA_LAKE_DB,
):
"""
Helper function to run MSCK REPAIR TABLE on an Athena table
:param table_name: the table name
:type table_name: str
:param lake_db: the database name
:type lake_db: str
:raises Exception: if the table does not exist
"""
import awswrangler as wr
if not wr.catalog.does_table_exist(
database=database,
table=table_name,
):
raise AirflowFailException(f"Table {database}.{table_name} does not exist!")
wr.athena.start_query_execution(
f"MSCK REPAIR TABLE `{table_name}`;",
database=database,
)
[docs]
@staticmethod
def replicate_athena_table(
old_name: str,
new_name: str,
lake_db: str = Config.DATA_LAKE_DB,
new_description: str = None,
new_columns_comments: dict = None,
new_s3_path: str = "",
):
"""Moves an athena table to new name, with new descriptions;
Perform COPY DDL if new_s3_path is specified: then a new table will be created with empty data at new location; old table is not affected!
Perform MOVE TABLE if new_s3_path is not specified: new table will have exactly same data as old table (and old table will be dropped from Glue catalog!)
It is essentially a wrapper around awswrangler function wr create_parquet_table.
param: old_name: name of old existing table
param: new_name: name of new table to create
param: lake_db: aws database
param: new_description: must include if want to include these comments in the new table
param: new_columns_comments: must include if want to include these comments in the new table
param: new_s3_path: EXTERNAL LOCATION for the NEW and EMPTY table;
Source inspiration: https://stackoverflow.com/a/74394685/12572612
"""
import awswrangler as wr
assert (
old_name != new_name
), f"Cannot replicate Athena tables with same name: {old_name=} == {new_name=}!"
if not wr.catalog.does_table_exist(
database=lake_db,
table=old_name,
):
raise Exception(f"Table {lake_db}.{old_name} does not exist!")
if wr.catalog.does_table_exist(
database=lake_db,
table=new_name,
):
raise Exception(
f"""
Table {lake_db}.{new_name} already exists...
We dont want the underlying wr create_parquet_table (default `mode='overwrite'`)
to accidentally overwrite the existing table {new_name}!
"""
)
old_s3_path = wr.catalog.get_table_location(database=lake_db, table=old_name)
# REQUIRED PARAMS
args = {
"database": lake_db,
"table": new_name,
"compression": "snappy",
}
# OPTIONAL PARAMS
if new_description is not None:
args["description"] = new_description
if new_columns_comments is not None:
args["columns_comments"] = new_columns_comments
# IMPORTANT PARAMS: partitions_types, columns_types!
old_regular_dict, old_partitions_dict = DataLakeUtil.get_athena_columns_defs(
old_name
)
if len(old_partitions_dict):
args["partitions_types"] = old_partitions_dict
args["columns_types"] = old_regular_dict
print(
f"""MAIN: CREATING new EMPTY table with following unpacked dictionary:
{args}
"""
)
if new_s3_path:
assert (
new_s3_path != old_s3_path
), f"Cannot use same S3 path {new_s3_path} for both: old table {old_name} and new table {new_name}!!!"
print(
f"""
Creating EMPTY new table {lake_db}.{new_name}::{new_s3_path}
with same DDL as old table {old_name}::{old_s3_path} (paths are different)
"""
)
args["path"] = new_s3_path
else:
print(
f"""
Moving old table {old_name} to new table {lake_db}.{new_name}...
Preserve DDL, and data at same S3 location: {old_s3_path}
"""
)
args["path"] = old_s3_path
# DEV NOTE: UNSURE about ordering of partition_types and possible effect when original table is partitioned by multiple cols! Might need OrderDict!
# This function works because multiple external glue tables could be based off same S3 location!
wr.catalog.create_parquet_table(**args)
if len(old_partitions_dict) > 0 and not new_s3_path:
print(
f"""
Running MSCK REPAIR TABLE to detect and load partitions for new table {new_name}, from data in old s3: {old_s3_path}....
NOTE: this is required after create_parquet_table from a PARTITIONED TABLE (e.g. c_cai_sbo)
Normal AWS Glue partitions are un-projected (similar to Partitioning Using Inheritance in Postgres).
Sources: https://stackoverflow.com/a/77833677/12572612;
https://docs.aws.amazon.com/athena/latest/ug/partition-projection.html;
https://aws-sdk-pandas.readthedocs.io/en/stable/tutorials/017%20-%20Partition%20Projection.html
"""
)
DataLakeUtil.msck_repair_table(
table_name=new_name,
database=lake_db,
)
# VALIDATE new against old: UNCHANGED VALUES
new_regular_dict, new_partitions_dict = DataLakeUtil.get_athena_columns_defs(
new_name
)
same_regular_cols = old_regular_dict == new_regular_dict
same_partitions = old_partitions_dict == new_partitions_dict
# VALIDATE new against old: CHANGED VALUES
if same_regular_cols & same_partitions:
new_s3_location = wr.catalog.get_table_location(
database=lake_db, table=new_name
)
new_count = wr.athena.read_sql_query(
database=lake_db, sql=f"SELECT COUNT(*) AS new_count FROM {new_name}"
).iloc[0, 0]
if new_s3_path == new_s3_location and new_count == 0:
print(
f"""
SUCCESSFULLY COPIED DDL (metadata only)
FROM: {old_s3_path}: Table {old_name}
TO: {new_s3_path}: Table {new_name}
THE END!
"""
)
elif new_s3_location == old_s3_path:
old_count = wr.athena.read_sql_query(
database=lake_db,
sql=f"SELECT COUNT(*) AS old_count FROM {old_name}",
).iloc[0, 0]
if new_count == old_count:
print(
f"""
SUCCESSFULLY RENAMED {old_name} to {new_name}, with unchanged S3 data and location..
>> Lastly, dropping old table {old_name} (metadata ONLY) from Glue catalog!"""
)
wr.catalog.delete_table_if_exists(database=lake_db, table=old_name)
else:
raise Exception(
f"""FAILED TO REPLICATE athena table: CONDITIONS::
{new_s3_path=}/{new_s3_location=}
{old_s3_path=}
{new_count=}/{old_count=}:
>> NEXT STEP: Inspect new table {lake_db}.{new_name} (not dropped) (comparing to old table {old_name})!
"""
)
else:
raise Exception(
f"""FAILED TO REPLICATE athena table: CONDITIONS::
{same_regular_cols=}/{same_partitions=}:
>> NEXT STEP: Inspect new table {lake_db}.{new_name} (not dropped) (comparing to old table {old_name})!
"""
)
[docs]
@staticmethod
def remove_datalake_table(
tgt_table: str,
tgt_prefix: str = "",
tgt_db: str = Config.DATA_LAKE_DB,
tgt_bucket: str = Config.DATA_LAKE_BUCKET,
):
"""
[This method is deprecated and will be removed in a future release.] Should use new DeleteAthenaTableOperator!
Drop Athena table from Glue catalog (metadata) and also delete all S3 data, if table exists
:param tgt_table: the target table
:param tgt_prefix: the target prefix: OPTIONAL, if provided, will delete the specified prefix instead of inferring from table catalog (and even when the catalog table is missing — clears orphans from failed CTAS runs)
:param tgt_db: the target db
:param tgt_bucket: the target bucket: used only when the catalog table is missing and tgt_prefix is provided; otherwise inferred from the catalog
"""
import awswrangler as wr
path_info = DataLakeUtil.get_table_path_info(tgt_db=tgt_db, tgt_table=tgt_table)
if not path_info:
# A failed CTAS can leave orphaned S3 files with NO catalog table — the
# next clean+CTAS then hits HIVE_PATH_ALREADY_EXISTS forever (ID stage3,
# DATASME-2028, 2026-08-31). When the caller named the prefix explicitly,
# purge it regardless of catalog state; only the catalog-inferred path
# requires the table to exist.
if tgt_prefix:
print(
f"Table {tgt_db}.{tgt_table} does not exist; removing possible "
f"orphaned S3 data at bucket {tgt_bucket} prefix {tgt_prefix}"
)
S3Util.delete_s3_folder(bucket=tgt_bucket, prefix=tgt_prefix)
return
print(
f"Table {tgt_db}.{tgt_table} does not exist, nothing to remove from Data Lake"
)
return
full_s3_path, infer_bucket, infer_prefix = path_info
wr.catalog.delete_table_if_exists(database=tgt_db, table=tgt_table)
# Refresher on Python falsiness: if False or None or "" or () or {} or set() or [] or range(0) or 0 or 0.0 or 0j: print("Falsey")
if tgt_prefix:
print(
f"Removing S3 data from {full_s3_path} i.e. bucket {infer_bucket} and prefix {tgt_prefix}"
)
S3Util.delete_s3_folder(bucket=infer_bucket, prefix=tgt_prefix)
return
print(
f"Removing S3 data from {full_s3_path} i.e. bucket {infer_bucket} and prefix {infer_prefix}"
)
S3Util.delete_s3_folder(bucket=infer_bucket, prefix=infer_prefix)
[docs]
@staticmethod
def ctas(
prefix: str,
query: str,
tgt_table: str,
partition_by: list[str] = [],
bucket: str = Config.DATA_LAKE_BUCKET,
src_db: str = Config.DATA_LAKE_DB,
tgt_db: str = Config.DATA_LAKE_DB,
workgroup: str = Config.DEFAULT_WORKGROUP,
clean: bool = True,
data_source: str = "AwsDataCatalog",
bucketing_info: tuple[list[str], int] = None,
):
"""
[Deprecated] This method is deprecated and will be removed in a future release. Should use new AthenaCTASOperator instead!
Create a table as select. Uses `awswrangler.athena.create_ctas_table` under the hood.
See [docs](https://aws-sdk-pandas.readthedocs.io/en/stable/stubs/awswrangler.athena.create_ctas_table.html#awswrangler.athena.create_ctas_table)
for a more in depth explanation of the parameters.
Uses sane defaults for data lake
:param prefix: the prefix
:param query: the query
:param tgt_table: the target table
:param partition_by: the partitions
:param bucket: the bucket
:param tgt_db: the target db
:param workgroup: the workgroup
:param clean: whether or not to clean
"""
import awswrangler as wr
if clean:
DataLakeUtil.remove_datalake_table(
tgt_table=tgt_table,
tgt_db=tgt_db,
tgt_prefix=prefix,
tgt_bucket=bucket,
)
wr.athena.create_ctas_table(
sql=query,
database=src_db,
ctas_database=tgt_db,
data_source=data_source,
ctas_table=tgt_table,
storage_format="PARQUET",
write_compression="snappy",
s3_output=f"s3://{bucket}/{prefix}",
wait=True,
workgroup=workgroup,
partitioning_info=partition_by,
bucketing_info=bucketing_info,
)
[docs]
@staticmethod
def create_and_start_crawler(
crawler_name: str,
tgt_path: str,
table_prefix: str = "",
tgt_db: str = Config.DATA_LAKE_DB,
tgt_bucket: str = Config.DATA_LAKE_BUCKET,
wait: bool = True,
max_iterations: int = Config.MAX_ITERATIONS,
sleep_time: int = Config.CRAWLER_SLEEP_TIME,
append_dev: bool = True,
):
"""
Create and start a glue crawler
:param crawler_name: the crawler name
:param tgt_path: the target path
:param table_prefix: the table prefix for boto client.create_crawler to catalog tables that are created i.e. all external tables stored at tgt_path that have this names starting with table_prefix will be crawled (scanned)
:param tgt_db: the target db
:param tgt_bucket: the target bucket
:param wait: whether or not to wait
:param max_iterations: the max iterations
:param sleep_time: the sleep time
:param append_dev: whether or not to append dev to the crawler name
"""
import boto3
if append_dev:
crawler_name += f"{'' if Config.PRODUCTION else '_dev'}"
client = boto3.client("glue")
crawler_response = client.list_crawlers()
crawler_names = crawler_response["CrawlerNames"]
while crawler_response.get("NextToken", False):
crawler_response = client.list_crawlers(
NextToken=crawler_response["NextToken"]
)
crawler_names.extend(crawler_response.get("CrawlerNames", []))
if crawler_name in crawler_names:
client.delete_crawler(Name=crawler_name)
client.create_crawler(
Name=crawler_name,
Role=Config.CRAWLER_ROLE,
DatabaseName=tgt_db,
Targets={
"S3Targets": [
{
"Path": f"s3://{tgt_bucket}/{tgt_path}",
},
],
},
TablePrefix=table_prefix,
)
client.start_crawler(Name=crawler_name)
crawler = client.get_crawler(Name=crawler_name)
count = 0
while (
wait and crawler["Crawler"]["State"] != "READY" and count < max_iterations
):
crawler = client.get_crawler(Name=crawler_name)
sleep(sleep_time)
count += 1
if count == max_iterations:
print("Crawler timed out")
[docs]
@staticmethod
def get_valid_fabric_table_from_version(
version_number: int,
tgt_db: str = Config.DATA_LAKE_DB,
) -> str | None:
"""Checks if the `version_number` is a valid number. It converts that to the table name and validates it
version_number: int a number. 0 means the latest version of fabric (v3).
:param version_number: int a number. 0 means the latest version of fabric (v3).
:param tgt_db: the target db
:raises AirflowFailException: if table name is invalid/does not exist
:return: fabric fix availabilty table name
"""
import awswrangler as wr
fabric_fix_availability_table = f't_fcc_fabric_fixed_availability{"_v{}".format(version_number) if version_number else ""}'
is_fabric_version_valid = wr.catalog.does_table_exist(
database=tgt_db, table=fabric_fix_availability_table
)
if is_fabric_version_valid == False:
raise AirflowFailException(
f"Wrong fabric fix availability version selected {version_number}"
)
return fabric_fix_availability_table
[docs]
@staticmethod
def list_states_with_cpc(tgt_database: str = Config.DATA_LAKE_DB) -> list[str]:
"""Check the datalake and returns the list of the states that their `c_cpc_merge_tile_` is generated. This function ignores states which have sub city lie mi_detroit
:param tgt_database: the datalake table, defaults to Config.DATA_LAKE_DB
:type tgt_database: str, optional
:return: an array of state_abbr
:rtype: list[str]
"""
import awswrangler as wr
from config.fips import FipsCodes
df_tables = wr.catalog.tables(
name_prefix="c_cpc_merge_tile_", database=tgt_database
)
if not len(df_tables["Table"]):
return []
all_cpc_table_suffixes = df_tables["Table"].str.replace("c_cpc_merge_tile_", "")
return [
state_abbr.upper()
for state_abbr in all_cpc_table_suffixes.to_list()
if state_abbr.upper() in FipsCodes.ALL.keys()
]
[docs]
@staticmethod
def s3_file_to_parquet(
s3_file_path: str,
raw_table_prefix: str,
raw_table_name: str,
chunk_size: int = 1_000_000,
tgt_bucket: str = Config.DATA_LAKE_BUCKET,
tgt_db: str = Config.DATA_LAKE_DB,
partition_cols: list[str] = [],
column_mapping: dict[str, str | list[str]] | None = None,
drop_empty_rows: bool = False,
dtypes: dict[str, str] | None = None,
) -> None:
"""Load the raw data from a s3 path and store it in the datalake as parquet file.
Use this function when you do not know what is expected from the client shared data.
All fields are forced to be string by default.
:param s3_file_path: the s3 path of the file
:type s3_file_path: str
:param raw_table_prefix: the prefix of the table. Example is "raw/uploads/Ohio/unzipped/data_revision_requests/{raw_table_name}"
:type raw_table_prefix: str
:param raw_table_name: the name of the table
:type raw_table_name: str
:param tgt_bucket: the target bucket, defaults to Config.DATA_LAKE_BUCKET
:type tgt_bucket: str, optional
:param tgt_db: the target db, defaults to Config.DATA_LAKE_DB
:type tgt_db: str, optional
:param column_mapping: dict where keys are target column names and values are
alias(es) to look for. Example: {"md5_pkey_hash": ["md5", "md5hash"]}
:type column_mapping: dict, optional
:param dtypes: opt-in map of column name -> Athena type (e.g. "double",
"bigint", "date", "timestamp") for columns that should be
typed instead of left as string. Keys must use the
normalized (lowercased) column names. Defaults to None,
which keeps every column a string (the historical
behavior — all other callers are unaffected).
:type dtypes: dict, optional
:param mode: the mode of the write, defaults to "overwrite", options are "overwrite", "append", "upsert"
"""
import awswrangler as wr
import geopandas as gpd
import pandas as pd
from tempfile import NamedTemporaryFile, TemporaryDirectory
import os
from zipfile import ZipFile
import fiona
# https://github.com/geopandas/geopandas/issues/2481
fiona.drvsupport.supported_drivers["KML"] = "rw"
def process_chunk(
df, raw_table_name, raw_table_prefix, sheet_name="", partition_cols=[]
):
df.columns = [
str(c)
.lower()
.replace("*", "")
.replace("#", "")
.strip()
.replace(" ", "_")
.replace("__", "")
.replace(" ", "")
.strip()
for c in df.columns
]
# remove the '_' at the end of the column name
df.columns = [c[:-1] if c.endswith("_") else c for c in df.columns]
# apply column mapping after normalization
DataLakeUtil.apply_column_mapping(df, column_mapping)
# remove duplicated columns
df = df.loc[:, ~df.columns.duplicated()].reset_index(drop=True)
# Drop rows where all values are empty/nan/null
if drop_empty_rows:
empty_values = {"", "nan", "null", "na", "none"}
mask = df.apply(
lambda row: not all(
str(v).strip().lower() in empty_values or pd.isna(v)
for v in row
),
axis=1,
)
original_count = len(df)
df = df[mask].reset_index(drop=True)
dropped = original_count - len(df)
if dropped > 0:
print(f"Dropped {dropped} empty rows")
data_columns = df.columns.to_list()
# Default: every column is registered as string. An optional `dtypes`
# map (column -> Athena type) opts specific columns into real types,
# coerced here so the parquet is physically written to match — you
# can't ALTER a string parquet column to numeric after the fact.
# Unlisted columns stay string. Keys use the normalized names above.
typed = {c: t.lower() for c, t in (dtypes or {}).items() if c in df.columns}
for col, atype in typed.items():
if atype in ("double", "float", "real", "decimal"):
df[col] = pd.to_numeric(df[col], errors="coerce")
elif atype in ("bigint", "int", "integer", "smallint", "tinyint"):
df[col] = pd.to_numeric(df[col], errors="coerce").astype("Int64")
elif atype == "date":
df[col] = pd.to_datetime(df[col], errors="coerce").dt.date
elif atype == "timestamp":
df[col] = pd.to_datetime(df[col], errors="coerce")
elif atype in ("boolean", "bool"):
df[col] = df[col].astype("boolean")
# string-ify only the columns we are NOT typing
str_cols = [c for c in data_columns if c not in typed]
if str_cols:
df[str_cols] = df[str_cols].astype(str)
data_columns_dict = {c: typed.get(c, "string") for c in data_columns}
wr.s3.to_parquet(
df=df,
path="s3://{tgt_bucket}/{raw_table_prefix}{sheet_name}".format(
tgt_bucket=tgt_bucket,
raw_table_prefix=raw_table_prefix,
sheet_name=sheet_name,
),
dataset=True,
mode="append",
database=tgt_db,
table=raw_table_name + sheet_name,
index=False,
dtype=data_columns_dict,
partition_cols=partition_cols,
)
print(f"File loaded at {raw_table_name}{sheet_name}")
print("*" * 60)
GEO_CHUNK_SIZE = 200_000
def _process_single_geo_chunk(df, partition_cols=[]) -> None:
"""Process a single GeoDataFrame chunk: reproject, convert geometry to WKT, write."""
if df.crs is None:
df.crs = "EPSG:4326"
else:
df.to_crs(epsg=4326, inplace=True)
df["wkt"] = df.geometry.to_wkt()
df.drop(columns="geometry", inplace=True)
process_chunk(
df, raw_table_name, raw_table_prefix, partition_cols=partition_cols
)
def process_geopanda_df(
df_or_path, partition_cols=[], layer=None, geo_chunk_size=GEO_CHUNK_SIZE
) -> None:
"""
Handles geopandas dataframes and loads them into the datalake.
Accepts either a GeoDataFrame or a file path string for chunked reading.
"""
if isinstance(df_or_path, str):
# Chunked path: read from file in slices to limit memory
total = len(fiona.open(df_or_path, layer=layer))
print(
f"Chunked geo read: {total} features, chunk_size={geo_chunk_size}"
)
for start in range(0, total, geo_chunk_size):
end = min(start + geo_chunk_size, total)
chunk = gpd.read_file(
df_or_path, layer=layer, rows=slice(start, end)
)
_process_single_geo_chunk(chunk, partition_cols=partition_cols)
print(f" Geo chunk {start:,}-{end:,} of {total:,} written")
else:
# Legacy: already a GeoDataFrame in memory
_process_single_geo_chunk(df_or_path, partition_cols=partition_cols)
otherfiles = [".geojson", ".shp", ".gpkg", ".kml", "kmz"]
s3_file_lower = s3_file_path.lower()
if s3_file_lower.endswith(".csv"):
for encoding in ["utf-8", "cp1252", "latin1", "iso-8859-1"]:
try:
for i, chunk in enumerate(
wr.s3.read_csv(
s3_file_path, chunksize=chunk_size, encoding=encoding
)
):
process_chunk(
chunk,
raw_table_name,
raw_table_prefix,
partition_cols=partition_cols,
)
print(f"chunk {i} loaded")
print(f"Successfully read file at {s3_file_path} with {encoding=}")
break
except UnicodeDecodeError as ude:
print(
f"Failed to read with {encoding=}, Error: {ude}; trying next..."
)
print("All chunks loaded")
elif s3_file_lower.endswith(".xlsx") or s3_file_lower.endswith(".xls"):
print("processing the excel file")
# iterate through all sheets in this excel and ingest them as separate tables
sheet_names = wr.s3.read_excel(s3_file_path, sheet_name=None).keys()
if len(sheet_names) == 1:
df = wr.s3.read_excel(s3_file_path)
process_chunk(
df, raw_table_name, raw_table_prefix, partition_cols=partition_cols
)
else:
import re
for sheet_name in sheet_names:
df = wr.s3.read_excel(s3_file_path, sheet_name=sheet_name)
sheet_name_underscore = [" ", "-", "."]
sheet_name_clean = sheet_name.lower().strip()
for char in sheet_name_underscore:
sheet_name_clean = sheet_name_clean.replace(char, "_")
# Erase any character that isn't alphanumeric or an underscore
sheet_name_clean = re.sub(r"[^a-zA-Z0-9_]", "", sheet_name_clean)
process_chunk(
df,
raw_table_name,
raw_table_prefix,
sheet_name=sheet_name_clean,
partition_cols=partition_cols,
)
elif s3_file_lower.endswith(".gdb"):
# The GDB is actually a directory, so we need to download it and read it with geopandas
with TemporaryDirectory(suffix=".gdb") as temp_dir:
S3Util.download_folder(s3_file_path, temp_dir, verbose=True)
layers = fiona.listlayers(temp_dir)
print("Layers found in the GDB file:")
for layer in layers:
print(layer)
for layer in layers:
process_geopanda_df(
temp_dir, partition_cols=partition_cols, layer=layer
)
elif s3_file_lower.endswith(".gdb.zip") or (
# sometimes the gdb files are zipped with multiple files inside
".gdb" in s3_file_lower
and ".zip" in s3_file_lower
):
print("processing the zipped gdb file")
print("*" * 60)
# The GDB is actually a directory, so we need to download it and read it with geopandas
with NamedTemporaryFile() as temp_file:
# NamedTemporaryFile only removes temp_file.name (an empty
# placeholder); the real download lives at a sibling path and
# must be removed explicitly or it is orphaned in /tmp.
gdb_path = temp_file.name + ".gdb.zip"
try:
wr.s3.download(s3_file_path, gdb_path)
print(temp_file.name)
layers = fiona.listlayers(gdb_path)
print("Layers found in the GDB file:")
for layer in layers:
print(layer)
for layer in layers:
process_geopanda_df(
gdb_path, partition_cols=partition_cols, layer=layer
)
finally:
if os.path.exists(gdb_path):
os.remove(gdb_path)
elif any([s3_file_lower.endswith(x) for x in otherfiles]):
# download the data to a temporary file and read it with geopandas
file_format = s3_file_path.split(".")[-1]
with NamedTemporaryFile() as temp_file:
# NamedTemporaryFile only removes temp_file.name (an empty
# placeholder); the real download(s) live at sibling paths and
# must be removed explicitly or they are orphaned in /tmp.
local_path = temp_file.name + "." + file_format
downloaded = [local_path]
try:
wr.s3.download(s3_file_path, local_path)
# fiona.errors.DriverError: Unable to open .shx
if file_format == "shp":
s3_file_path_shx = s3_file_path.replace(".shp", ".shx")
local_path_shx = temp_file.name + ".shx"
downloaded.append(local_path_shx)
try:
wr.s3.download(s3_file_path_shx, local_path_shx)
print(f"Downloaded {s3_file_path_shx} as well")
except Exception as e:
print(f"No {s3_file_path_shx} file")
pass
# If the .dbf file exists, we want to download it as well
s3_file_path_dbf = s3_file_path.replace(".shp", ".dbf")
local_path_dbf = temp_file.name + ".dbf"
downloaded.append(local_path_dbf)
try:
wr.s3.download(s3_file_path_dbf, local_path_dbf)
print(f"Downloaded {s3_file_path_dbf} as well")
except Exception as e:
print(f"No {s3_file_path_dbf} file, but its not required")
try:
if file_format in ("kml", "kmz"):
# KML/KMZ don't support rows= slicing — read fully
if file_format == "kml":
df = gpd.read_file(local_path, driver="KML")
else:
df = gpd.read_file(
"/vsizip"
+ temp_file.name
+ ".kmz/"
+ Config.KMZ_KML_DEFAULT,
driver="KML",
)
_process_single_geo_chunk(df, partition_cols=partition_cols)
else:
# gpkg, geojson, shp — use chunked reading
process_geopanda_df(local_path, partition_cols=partition_cols)
except Exception as e:
raise ValueError(f"Error reading the {s3_file_path}. Error: {e}")
finally:
for path in downloaded:
if os.path.exists(path):
os.remove(path)
elif s3_file_lower.endswith(".zip"):
# unzip the file and detect content type
print("Unzipping the file")
with NamedTemporaryFile(suffix=".zip") as temp_file:
wr.s3.download(s3_file_path, temp_file.name)
zip_dir = "./tmp/zipfiled_s3"
os.makedirs(zip_dir, exist_ok=True)
file_size = os.path.getsize(temp_file.name)
print(f"Downloaded file size: {file_size} bytes")
with ZipFile(temp_file.name, "r") as zip_ref:
zip_ref.extractall(zip_dir)
all_files = []
for root_dir, dirs, files_list in os.walk(zip_dir):
for f in files_list:
all_files.append(os.path.join(root_dir, f))
extensions = {os.path.splitext(f)[1].lower() for f in all_files}
print(f"Extracted {len(all_files)} files, extensions: {extensions}")
if ".shp" in extensions:
# Zipped shapefiles — chunked reading
print("Detected shapefiles in zip")
shp_files = [f for f in all_files if f.lower().endswith(".shp")]
for shp_file in shp_files:
print(f" Reading shapefile: {shp_file}")
process_geopanda_df(shp_file, partition_cols=partition_cols)
elif ".gpkg" in extensions:
# Zipped GeoPackage — chunked reading per layer
print("Detected GeoPackage in zip")
gpkg_files = [f for f in all_files if f.lower().endswith(".gpkg")]
for gpkg_file in gpkg_files:
print(f" Reading GeoPackage: {gpkg_file}")
layers = fiona.listlayers(gpkg_file)
for layer in layers:
print(f" Layer: {layer}")
process_geopanda_df(
gpkg_file, partition_cols=partition_cols, layer=layer
)
elif ".geojson" in extensions:
# Zipped GeoJSON — chunked reading
print("Detected GeoJSON in zip")
geojson_files = [
f for f in all_files if f.lower().endswith(".geojson")
]
for gj_file in geojson_files:
print(f" Reading GeoJSON: {gj_file}")
process_geopanda_df(gj_file, partition_cols=partition_cols)
elif ".csv" in extensions or ".CSV" in {
os.path.splitext(f)[1] for f in all_files
}:
# Zipped CSV
print("Detected CSV in zip")
for i, chunk in enumerate(
wr.s3.read_csv(s3_file_path, chunksize=chunk_size)
):
process_chunk(
chunk,
raw_table_name,
raw_table_prefix,
partition_cols=partition_cols,
)
print(f"chunk {i} loaded")
print("All chunks loaded")
else:
# Check for GDB directories inside the extracted zip
gdb_dirs = [
os.path.join(r, d)
for r, dirs, _ in os.walk(zip_dir)
for d in dirs
if d.lower().endswith(".gdb")
]
if not gdb_dirs:
raise ValueError(
f"Unsupported content in zip. Found extensions: {extensions}. "
"Supporting shapefiles, GeoPackage, GeoJSON, GDB, or CSV inside zip."
)
print("Detected GDB directory in zip")
for gdb_dir in gdb_dirs:
layers = fiona.listlayers(gdb_dir)
print(f" GDB: {gdb_dir}, layers: {layers}")
for layer in layers:
process_geopanda_df(
gdb_dir, partition_cols=partition_cols, layer=layer
)
# remove the tmp folder after data successful loaded
os.system(f"yes | rm -rf {zip_dir}")
print("Temporary files removed")
elif s3_file_lower.endswith(".parquet"):
df = gpd.read_parquet(s3_file_path)
process_chunk(
df, raw_table_name, raw_table_prefix, partition_cols=partition_cols
)
else:
raise ValueError(
"File format not supported. Supporting zipped csv, plain csv, xlsx, geojson, or geopackage."
)
return None
[docs]
@staticmethod
def read_multiple_queries(
queries: list[str], database=Config.DATA_LAKE_DB, max_concurrent: int = None
) -> list:
"""
Executes multiple queries simultaneously and returns the results as a list of dataframes
:param queries: the queries to run
:param database: the database to run the queries on
:param max_concurrent: the maximum number of queries to run simultaneously.
Defaults to 50 if all queries are only DML and 10 otherwise
:return: the results as a list of dataframes
"""
import awswrangler as wr
dml_keywords = [
"SELECT",
"INSERT",
"UPDATE",
"DELETE",
"MERGE",
"CALL",
"EXPLAIN",
"LOCK",
]
dml_only = all([q.strip().split()[0].upper() in dml_keywords for q in queries])
max_con = (50 if dml_only else 10) if max_concurrent is None else max_concurrent
query_pages = [
queries[i : i + max_con] for i in range(0, len(queries), max_con)
]
# TODO: The efficiency can be improved by running these as a pool instead of paged
def run_queries(queries):
execution_ids = [
wr.athena.start_query_execution(query, database) for query in queries
]
return [
wr.athena.get_query_results(query_execution_id=e_id)
for e_id in execution_ids
]
return [result for page in query_pages for result in run_queries(page)]