pudl.io_managers

Dagster IO Managers.

Attributes

logger
MINIMUM_SQLITE_VERSION

Classes

PudlMixedFormatIOManager	Format switching IOManager that supports sqlite and parquet.
SQLiteIOManager	IO Manager that writes and retrieves dataframes from a SQLite database.
PudlParquetIOManager	IOManager that writes pudl tables to pyarrow parquet files.
PudlGeoParquetIOManager	Do some extra work to output valid GeoParquet files when appropriate.
PudlSQLiteIOManager	IO Manager that writes and retrieves dataframes from a SQLite database.
FercSQLiteIOManager	IO Manager for reading tables from FERC databases.
FercDBFSQLiteIOManager	IO Manager for only reading tables from the FERC 1 database.
FercXBRLSQLiteIOManager	IO Manager for only reading tables from the XBRL database.

Functions

get_table_name_from_context(→ str)	Retrieves the table name from the context object.
pudl_mixed_format_io_manager(→ dagster.IOManager)	Create a SQLiteManager dagster resource for the pudl database.
parquet_io_manager(→ dagster.IOManager)	Create a Parquet only IO manager.
geoparquet_io_manager(→ dagster.IOManager)	Create a GeoParquet only IO manager.
ferc1_dbf_sqlite_io_manager(→ FercDBFSQLiteIOManager)	Create a SQLiteManager dagster resource for the ferc1 dbf database.
ferc1_xbrl_sqlite_io_manager(→ FercXBRLSQLiteIOManager)	Create a SQLiteManager dagster resource for the ferc1 xbrl database.
ferc714_xbrl_sqlite_io_manager(→ FercXBRLSQLiteIOManager)	Create a SQLiteManager dagster resource for the ferc714 xbrl database.

Module Contents

pudl.io_managers.logger

pudl.io_managers.MINIMUM_SQLITE_VERSION = '3.32.0'

pudl.io_managers.get_table_name_from_context(context: dagster.OutputContext) → str

Retrieves the table name from the context object.

class pudl.io_managers.PudlMixedFormatIOManager(write_to_parquet: bool = False, read_from_parquet: bool = False)

Bases: dagster.IOManager

Format switching IOManager that supports sqlite and parquet.

This IOManager provides for the use of parquet files along with the standard SQLite database produced by PUDL.

write_to_parquet: bool

If true, data will be written to parquet files.

read_from_parquet: bool

If true, data will be read from parquet files instead of sqlite.

_sqlite_io_manager

_parquet_io_manager

handle_output(context: dagster.OutputContext, obj: pandas.DataFrame | str) → pandas.DataFrame

Passes the output to the appropriate IO manager instance.

load_input(context: dagster.InputContext) → pandas.DataFrame

Reads input from the appropriate IO manager instance.

class pudl.io_managers.SQLiteIOManager(base_dir: str, db_name: str, md: sqlalchemy.MetaData | None = None, timeout: float = 1000.0)

Bases: dagster.IOManager

IO Manager that writes and retrieves dataframes from a SQLite database.

base_dir

db_name

md = None

engine = None

_setup_database(timeout: float = 1000.0) → sqlalchemy.Engine

Create database and metadata if they don’t exist.

Parameters:

timeout – How many seconds the connection should wait before raising an exception, if the database is locked by another connection. If another connection opens a transaction to modify the database, it will be locked until that transaction is committed.

Returns:

SQL Alchemy engine that connects to a database in the base_dir.

Return type:

engine

_get_sqlalchemy_table(table_name: str) → sqlalchemy.Table

Get SQL Alchemy Table object from metadata given a table_name.

Parameters:

table_name – The name of the table to look up.

Returns:

Corresponding SQL Alchemy Table in SQLiteIOManager metadata.

Return type:

table

Raises:

ValueError – if table_name does not exist in the SQLiteIOManager metadata.

_handle_pandas_output(context: dagster.OutputContext, df: pandas.DataFrame)

Write dataframe to the database.

SQLite does not support concurrent writes to the database. Instead, SQLite queues write transactions and executes them one at a time. This allows the assets to be processed in parallel. See the SQLAlchemy docs to learn more about SQLite concurrency.

Parameters:

• context – dagster keyword that provides access to output information like asset name.

• df – dataframe to write to the database.

_handle_str_output(context: dagster.OutputContext, query: str)

Execute a sql query on the database.

This is used for creating output views in the database.

Parameters:

• context – dagster keyword that provides access output information like asset name.

• query – sql query to execute in the database.

handle_output(context: dagster.OutputContext, obj: pandas.DataFrame | str)

Handle an op or asset output.

If the output is a dataframe, write it to the database. If it is a string execute it as a SQL query.

Parameters:

• context – dagster keyword that provides access output information like asset name.

• obj – a sql query or dataframe to add to the database.

Raises:

Exception – if an asset or op returns an unsupported datatype.

load_input(context: dagster.InputContext) → pandas.DataFrame

Load a dataframe from a sqlite database.

Parameters:

context – dagster keyword that provides access output information like asset name.

class pudl.io_managers.PudlParquetIOManager

Bases: dagster.IOManager

IOManager that writes pudl tables to pyarrow parquet files.

handle_output(context: dagster.OutputContext, obj: pandas.DataFrame | polars.LazyFrame) → None

Writes pudl dataframe to parquet file.

load_input(context: dagster.InputContext) → pandas.DataFrame | geopandas.GeoDataFrame | polars.LazyFrame

Loads pudl table from parquet file.

class pudl.io_managers.PudlGeoParquetIOManager

Bases: PudlParquetIOManager

Do some extra work to output valid GeoParquet files when appropriate.

_create_geoparquet_metadata(gdf: geopandas.GeoDataFrame, res: pudl.metadata.classes.Resource) → str

Create GeoParquet metadata JSON string.

handle_output(context: dagster.OutputContext, obj: geopandas.GeoDataFrame) → None

Write a PUDL dataframe to GeoParquet.

class pudl.io_managers.PudlSQLiteIOManager(base_dir: str, db_name: str, package: pudl.metadata.classes.Package | None = None, timeout: float = 1000.0)

Bases: SQLiteIOManager

IO Manager that writes and retrieves dataframes from a SQLite database.

This class extends the SQLiteIOManager class to manage database metadata and dtypes using the pudl.metadata.classes.Package class.

package = None

_handle_str_output(context: dagster.OutputContext, query: str)

Execute a sql query on the database.

This is used for creating output views in the database.

Parameters:

• context – dagster keyword that provides access output information like asset name.

• query – sql query to execute in the database.

_handle_pandas_output(context: dagster.OutputContext, df: pandas.DataFrame)

Enforce PUDL DB schema and write dataframe to SQLite.

load_input(context: dagster.InputContext) → pandas.DataFrame

Load a dataframe from a sqlite database.

Parameters:

context – dagster keyword that provides access output information like asset name.

pudl.io_managers.pudl_mixed_format_io_manager(init_context: dagster.InitResourceContext) → dagster.IOManager

Create a SQLiteManager dagster resource for the pudl database.

pudl.io_managers.parquet_io_manager(init_context: dagster.InitResourceContext) → dagster.IOManager

Create a Parquet only IO manager.

pudl.io_managers.geoparquet_io_manager(init_context: dagster.InitResourceContext) → dagster.IOManager

Create a GeoParquet only IO manager.

class pudl.io_managers.FercSQLiteIOManager(base_dir: str = None, db_name: str = None, md: sqlalchemy.MetaData = None, timeout: float = 1000.0)

Bases: SQLiteIOManager

IO Manager for reading tables from FERC databases.

This class should be subclassed and the load_input and handle_output methods should be implemented.

This IOManager expects the database to already exist.

_setup_database(timeout: float = 1000.0) → sqlalchemy.Engine

Create database engine and read the metadata.

Parameters:

timeout – How many seconds the connection should wait before raising an exception, if the database is locked by another connection. If another connection opens a transaction to modify the database, it will be locked until that transaction is committed.

Returns:

SQL Alchemy engine that connects to a database in the base_dir.

Return type:

engine

abstractmethod handle_output(context: dagster.OutputContext, obj)

Handle an op or asset output.

abstractmethod load_input(context: dagster.InputContext) → pandas.DataFrame

Load a dataframe from a sqlite database.

Parameters:

context – dagster keyword that provides access output information like asset name.

class pudl.io_managers.FercDBFSQLiteIOManager(base_dir: str = None, db_name: str = None, md: sqlalchemy.MetaData = None, timeout: float = 1000.0)

Bases: FercSQLiteIOManager

IO Manager for only reading tables from the FERC 1 database.

This IO Manager is for reading data only. It does not handle outputs because the raw FERC tables are not known prior to running the ETL and are not recorded in our metadata.

abstractmethod handle_output(context: dagster.OutputContext, obj: pandas.DataFrame | str)

Handle an op or asset output.

load_input(context: dagster.InputContext) → pandas.DataFrame

Load a dataframe from a sqlite database.

Parameters:

context – dagster keyword that provides access output information like asset name.

pudl.io_managers.ferc1_dbf_sqlite_io_manager(init_context) → FercDBFSQLiteIOManager

Create a SQLiteManager dagster resource for the ferc1 dbf database.

class pudl.io_managers.FercXBRLSQLiteIOManager(base_dir: str = None, db_name: str = None, md: sqlalchemy.MetaData = None, timeout: float = 1000.0)

Bases: FercSQLiteIOManager

IO Manager for only reading tables from the XBRL database.

This IO Manager is for reading data only. It does not handle outputs because the raw FERC tables are not known prior to running the ETL and are not recorded in our metadata.

static refine_report_year(df: pandas.DataFrame, xbrl_years: list[int]) → pandas.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.

abstractmethod handle_output(context: dagster.OutputContext, obj: pandas.DataFrame | str)

Handle an op or asset output.

load_input(context: dagster.InputContext) → pandas.DataFrame

Load a dataframe from a sqlite database.

Parameters:

context – dagster keyword that provides access output information like asset name.

pudl.io_managers.ferc1_xbrl_sqlite_io_manager(init_context) → FercXBRLSQLiteIOManager

Create a SQLiteManager dagster resource for the ferc1 xbrl database.

pudl.io_managers.ferc714_xbrl_sqlite_io_manager(init_context) → FercXBRLSQLiteIOManager

Create a SQLiteManager dagster resource for the ferc714 xbrl database.