"""Dagster IO managers used by PUDL assets.
This module defines the IO-manager implementations that translate between Dagster asset
execution and PUDL's storage formats, including Parquet (with native GeoParquet support
for assets that return a :class:`geopandas.GeoDataFrame`) and the FERC prerequisite
SQLite databases. Put :class:`dagster.IOManager` and :class:`dagster.ConfigurableIOManager`
classes here, along with configured singleton instances that the default code location
reuses. Keep data-processing logic out of this module; it should focus on persistence,
loading, and storage-compatibility concerns.
For the underlying Dagster concept, see https://docs.dagster.io/guides/build/io-managers
"""
import hashlib
import re
from pathlib import Path
from typing import Any, ClassVar, Literal
import dagster as dg
import geopandas
import pandas as pd
import polars as pl
import pyarrow.parquet as pq
import sqlalchemy as sa
from dagster import DagsterInvariantViolationError, InputContext, OutputContext
from pydantic import PrivateAttr
import pudl.logging_helpers
from pudl.dagster.provenance import (
FercSqliteProvenance,
FercSqliteProvenanceRecord,
ferc_sqlite_provenance_is_compatible,
get_xbrl_extractor_version,
)
from pudl.dagster.resources import (
global_data_config_resource,
pudl_paths_resource,
zenodo_doi_settings_resource,
)
from pudl.helpers import get_parquet_table, get_parquet_table_polars
from pudl.metadata.classes import Resource
from pudl.settings import GlobalDataConfig
from pudl.workspace.datastore import ZenodoDoiSettings
from pudl.workspace.setup import PudlPaths
[docs]
logger = pudl.logging_helpers.get_logger(__name__)
[docs]
def _get_dagster_instance_if_available(
context: InputContext,
) -> dg.DagsterInstance | None:
"""Return the Dagster instance from an input context if one was provided.
Returns ``None`` in two cases where provenance checks should be skipped:
* The context has no attached instance (e.g. ad hoc ``InputContext`` objects built
by notebook or integration-test helpers).
* The instance is ephemeral (created by ``execute_in_process()`` without an explicit
``instance=`` argument). An ephemeral instance has an empty event log, so
provenance checks against it would always raise rather than meaningfully validate.
"""
try:
instance = context.instance
return None if instance.is_ephemeral else instance
except DagsterInvariantViolationError:
return None
[docs]
def get_table_name_from_context(context: InputContext | OutputContext) -> str:
"""Retrieves the table name from the context object."""
if not context.has_asset_key:
raise ValueError(
"These IO managers can only be used with assets, since "
"the table name is taken from the asset key."
)
return context.asset_key.to_python_identifier()
[docs]
class PudlParquetIOManager(dg.ConfigurableIOManager):
"""IOManager that writes pudl tables to pyarrow parquet files."""
[docs]
pudl_paths: dg.ResourceDependency[PudlPaths]
@staticmethod
[docs]
def handle_output(
self,
context: dg.OutputContext,
obj: pd.DataFrame | geopandas.GeoDataFrame | pl.LazyFrame,
) -> None:
"""Writes a pudl dataframe to a Parquet file.
GeoDataFrames are written as GeoParquet using native geopandas output,
which produces spec-compliant CRS metadata readable by DuckDB >= 1.5.
Regular DataFrames and Polars LazyFrames use the PUDL PyArrow schema to
enforce exact column types on disk.
"""
table_name = get_table_name_from_context(context)
res = Resource.from_id(table_name)
parquet_path = self.pudl_paths.parquet_path(table_name)
parquet_path.parent.mkdir(parents=True, exist_ok=True)
if isinstance(obj, geopandas.GeoDataFrame):
gdf = res.enforce_schema(obj)
gdf.to_parquet(
parquet_path,
index=False,
compression=pudl.PARQUET_COMPRESSION,
compression_level=pudl.PARQUET_GEOMETRY_COMPRESSION_LEVEL,
)
elif isinstance(obj, pd.DataFrame):
df = res.enforce_schema(obj)
pa_schema = res.to_pyarrow()
df.to_parquet(
path=parquet_path,
index=False,
schema=pa_schema,
compression=pudl.PARQUET_COMPRESSION,
compression_level=pudl.PARQUET_COMPRESSION_LEVEL,
)
elif isinstance(obj, pl.LazyFrame):
obj.cast(res.to_polars_dtypes()).sink_parquet(
parquet_path,
engine="streaming",
row_group_size=100_000,
compression=pudl.PARQUET_COMPRESSION,
compression_level=pudl.PARQUET_COMPRESSION_LEVEL,
)
else:
raise TypeError(
"PudlParquetIOManager only supports pandas DataFrames, "
f"geopandas GeoDataFrames, and Polars LazyFrames, got {type(obj)}."
)
self._record_parquet_file_metadata(context, parquet_path)
[docs]
parquet_io_manager = PudlParquetIOManager(pudl_paths=pudl_paths_resource)
[docs]
class FercSqliteIOManagerBase(dg.ConfigurableIOManager):
"""Shared lazy-loading behavior for FERC SQLite Dagster IO managers.
Subclasses provide the query details for a particular FERC SQLite backend, while
this base class owns three shared responsibilities:
1. lazily creating and caching a SQLAlchemy engine for the configured database
2. lazily reflecting and caching SQLAlchemy metadata once the database exists
3. checking Dagster provenance metadata before each read
"""
[docs]
global_data_config: dg.ResourceDependency[GlobalDataConfig]
[docs]
pudl_paths: dg.ResourceDependency[PudlPaths]
[docs]
zenodo_dois: dg.ResourceDependency[ZenodoDoiSettings]
[docs]
_engine: sa.Engine | None = PrivateAttr(default=None)
@property
[docs]
def _years_key(self) -> str:
return f"{self.data_format}_years"
@property
[docs]
def db_name(self) -> str:
"""Return the SQLite database name for this dataset and data format."""
return f"{self.dataset}_{self.data_format}"
@property
[docs]
def db_path(self) -> Path:
"""Return the canonical SQLite path for this dataset and data format."""
return self.pudl_paths.sqlite_path(self.db_name)
@property
[docs]
def engine(self) -> sa.Engine:
"""Return a cached SQLAlchemy engine for this FERC SQLite database."""
if self._engine is None:
# Pydantic explicitly permits mutating PrivateAttr fields on frozen
# models; pyrefly doesn't yet model that exception.
self._engine = sa.create_engine(f"sqlite:///{self.db_path}") # type: ignore[read-only]
assert self._engine is not None
return self._engine
[docs]
def teardown_after_execution(self, context: dg.InitResourceContext) -> None:
"""Dispose the cached engine when the resource's lifecycle ends."""
if self._engine is not None:
self._engine.dispose()
# Pydantic explicitly permits mutating PrivateAttr fields on frozen
# models; pyrefly doesn't yet model that exception.
self._engine = None # type: ignore[read-only]
self._metadata = None # type: ignore[read-only]
@property
[docs]
def _get_sqlalchemy_table(self, table_name: str) -> sa.Table:
"""Return reflected SQLAlchemy table metadata for a FERC SQLite table."""
sa_table = self.metadata.tables.get(table_name, None)
if sa_table is None:
raise ValueError(
f"{table_name} not found in {self.db_name} metadata. Either add the "
"table to the metadata or use a different IO Manager. Database is "
f"located at {self.db_path}."
)
return sa_table
[docs]
def _check_provenance(self, context: InputContext) -> None:
"""Check that the existing FERC SQLite database is compatible with this run.
This is intentionally separate from engine and metadata caching because the
compatibility check depends on the Dagster run context rather than on local
process state.
"""
zenodo_doi = self.zenodo_dois.get_doi(self.dataset)
assert self.global_data_config.ferc_to_sqlite is not None, (
"ferc_to_sqlite config must be set to read a FERC SQLite database."
)
provenance = FercSqliteProvenance(
dataset=self.dataset,
data_format=self.data_format,
zenodo_doi=zenodo_doi,
years=self.global_data_config.ferc_to_sqlite.get_dataset_years(
self.dataset, self.data_format
),
ferc_xbrl_extractor_version=get_xbrl_extractor_version(),
)
if ((instance := _get_dagster_instance_if_available(context)) is not None) and (
not ferc_sqlite_provenance_is_compatible(
observed_provenance=FercSqliteProvenanceRecord.from_dagster_instance(
instance=instance,
dataset=self.dataset,
data_format=self.data_format,
),
required_provenance=provenance,
)
):
raise RuntimeError(
f"{self.dataset}_{self.data_format} provenace metadata is not compatible"
" with requirements of current run. Refresh the FERC SQLite assets."
)
[docs]
def handle_output(self, context: dg.OutputContext, obj: pd.DataFrame | str) -> None:
"""Reject writes because these IO managers currently support reads only."""
raise NotImplementedError("Ferc SQLite IO managers can't write outputs yet.")
[docs]
def _query(self, table_name: str, years: list[int]) -> pd.DataFrame:
"""Execute a filtered read against the FERC SQLite database."""
raise NotImplementedError("Subclasses must implement _query.")
[docs]
class FercDbfSqliteIOManager(FercSqliteIOManagerBase):
"""IO manager for reading tables from FERC DBF SQLite databases.
Instantiate with ``dataset`` (``ferc1``, ``ferc714``, etc.)
"""
[docs]
def _query(self, table_name: str, years: list[int]) -> pd.DataFrame:
"""Execute the year-filtered read against the FERC DBF SQLite database."""
_ = self._get_sqlalchemy_table(table_name)
with self.engine.begin() as con:
return pd.read_sql_query(
f"SELECT * FROM {table_name} " # noqa: S608
"WHERE report_year BETWEEN :min_year AND :max_year;",
con=con,
params={
"min_year": min(years),
"max_year": max(years),
},
).assign(sched_table_name=table_name)
[docs]
class FercXbrlSqliteIOManager(FercSqliteIOManagerBase):
"""IO manager for reading tables from a FERC XBRL SQLite database.
Instantiate with ``dataset`` (``ferc1``, ``ferc714``, etc.).
"""
@staticmethod
[docs]
def refine_report_year(df: pd.DataFrame, xbrl_years: list[int]) -> pd.DataFrame:
"""Set a fact's report year by its actual dates.
Sometimes a fact belongs to a context which has no ReportYear associated with
it; other times there are multiple ReportYears associated with a single filing.
In these cases the report year of a specific fact may be associated with the
other years in the filing.
In many cases we can infer the actual report year from the fact's associated
time period - either duration or instant.
"""
is_duration = len({"start_date", "end_date"} - set(df.columns)) == 0
is_instant = "date" in df.columns
def get_year(df: pd.DataFrame, col: str) -> pd.Series:
datetimes = pd.to_datetime(df.loc[:, col], format="%Y-%m-%d", exact=False)
if datetimes.isna().any():
raise ValueError(f"{col} has null values!")
return datetimes.apply(lambda x: x.year)
if is_duration:
start_years = get_year(df, "start_date")
end_years = get_year(df, "end_date")
if not (start_years == end_years).all():
raise ValueError("start_date and end_date are in different years!")
new_report_years = start_years
elif is_instant:
new_report_years = get_year(df, "date")
else:
raise ValueError("Attempted to read a non-instant, non-duration table.")
# we include XBRL data from before our "officially supported" XBRL
# range because we want to use it to set start-of-year values for the
# first XBRL year.
xbrl_years_plus_one_previous = [min(xbrl_years) - 1] + xbrl_years
df = df.assign(report_year=new_report_years)
df = df.loc[df.report_year.isin(xbrl_years_plus_one_previous)]
return df.reset_index(drop=True)
[docs]
def _query(self, table_name: str, years: list[int]) -> pd.DataFrame:
"""Execute the full-table read against the FERC XBRL SQLite database.
Args:
table_name: Name of the table to query (without the ``raw_<db_name>__``
prefix).
years: Years to include in the result set (passed to
:meth:`refine_report_year`).
"""
# TODO (bendnorman): Figure out a better way to handle tables that
# don't have duration and instant variants.
# Not every table contains both instant and duration;
# return an empty dataframe if the table doesn't exist.
if table_name not in self.metadata.tables:
return pd.DataFrame()
sched_table_name = re.sub("_instant|_duration", "", table_name)
with self.engine.begin() as con:
df = pd.read_sql(
f"SELECT {table_name}.* FROM {table_name}", # noqa: S608 - table names not supplied by user
con=con,
).assign(sched_table_name=sched_table_name)
return df.pipe(self.refine_report_year, xbrl_years=years)
[docs]
ferc1_dbf_sqlite_io_manager = FercDbfSqliteIOManager(
global_data_config=global_data_config_resource,
pudl_paths=pudl_paths_resource,
zenodo_dois=zenodo_doi_settings_resource,
dataset="ferc1",
)
[docs]
ferc1_xbrl_sqlite_io_manager = FercXbrlSqliteIOManager(
global_data_config=global_data_config_resource,
pudl_paths=pudl_paths_resource,
zenodo_dois=zenodo_doi_settings_resource,
dataset="ferc1",
)
[docs]
ferc714_xbrl_sqlite_io_manager = FercXbrlSqliteIOManager(
global_data_config=global_data_config_resource,
pudl_paths=pudl_paths_resource,
zenodo_dois=zenodo_doi_settings_resource,
dataset="ferc714",
)
[docs]
default_io_managers: dict[str, Any] = {
"ferc1_dbf_sqlite_io_manager": ferc1_dbf_sqlite_io_manager,
"ferc1_xbrl_sqlite_io_manager": ferc1_xbrl_sqlite_io_manager,
"ferc714_xbrl_sqlite_io_manager": ferc714_xbrl_sqlite_io_manager,
"parquet_io_manager": parquet_io_manager,
}
