Query

Data Models

pydantic model dsgrid.query.models.ProjectQueryModel

Represents a user query on a Project.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Fields:

• project (dsgrid.query.models.ProjectQueryParamsModel)

field project: ProjectQueryParamsModel [Required]

Defines the datasets to use and how to transform them.

serialize_cached_content() → dict[str, Any]

Return a JSON-able representation of the model that can be used for caching purposes.

pydantic model dsgrid.query.models.ProjectQueryParamsModel

Defines how to transform a project into a CompositeDataset

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Fields:

• dataset (dsgrid.query.models.DatasetModel)

• excluded_dataset_ids (list[str])

• include_dsgrid_dataset_components (bool)

• mapping_plans (list[dsgrid.query.dataset_mapping_plan.DatasetMappingPlan])

• project_id (str)

• spark_conf_per_dataset (list[dsgrid.query.models.SparkConfByDataset])

• version (str | None)

Validators:

• check_duplicate_dataset_ids » mapping_plans

• check_duplicate_dataset_ids » spark_conf_per_dataset

• check_invalid_dataset_ids » all fields

• check_unsupported_fields » all fields

field dataset: DatasetModel [Required]

Definition of the dataset to create.

Validated by:

• check_invalid_dataset_ids

• check_unsupported_fields

field excluded_dataset_ids: list[str] = []

Datasets to exclude from query

Validated by:

• check_invalid_dataset_ids

• check_unsupported_fields

field include_dsgrid_dataset_components: bool = False

Validated by:

• check_invalid_dataset_ids

• check_unsupported_fields

field mapping_plans: list[DatasetMappingPlan] = []

Defines the order in which to map the dimensions of datasets.

Validated by:

• check_duplicate_dataset_ids

• check_invalid_dataset_ids

• check_unsupported_fields

field project_id: str [Required]

Project ID for query

Validated by:

• check_invalid_dataset_ids

• check_unsupported_fields

field spark_conf_per_dataset: list[SparkConfByDataset] = []

Apply these Spark configuration settings while a dataset is being processed.

Validated by:

• check_duplicate_dataset_ids

• check_invalid_dataset_ids

• check_unsupported_fields

field version: str | None = None

Version of project or dataset on which the query is based. Should not be set by the user

Validated by:

• check_invalid_dataset_ids

• check_unsupported_fields

validator check_duplicate_dataset_ids  »  spark_conf_per_dataset, mapping_plans

validator check_invalid_dataset_ids  »  all fields

validator check_unsupported_fields  »  all fields

get_dataset_mapping_plan(dataset_id: str) → DatasetMappingPlan | None

Return the mapping plan for this dataset_id or None if the user did not specify one.

get_spark_conf(dataset_id: str) → dict[str, Any]

Return the Spark settings to apply while processing dataset_id.

set_dataset_mapper(new_mapper: DatasetMappingPlan) → None

pydantic model dsgrid.query.models.QueryResultParamsModel

Controls post-processing and storage of CompositeDatasets

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Fields:

• aggregate_each_dataset (bool)

• aggregations (list[dsgrid.query.models.AggregationModel])

• column_type (dsgrid.query.models.ColumnType)

• dimension_filters (list[dsgrid.dimension.dimension_filters.DimensionFilterExpressionModel | dsgrid.dimension.dimension_filters.DimensionFilterExpressionRawModel | dsgrid.dimension.dimension_filters.DimensionFilterColumnOperatorModel | dsgrid.dimension.dimension_filters.DimensionFilterBetweenColumnOperatorModel | dsgrid.dimension.dimension_filters.SubsetDimensionFilterModel | dsgrid.dimension.dimension_filters.SupplementalDimensionFilterColumnOperatorModel])

• output_format (str)

• replace_ids_with_names (bool)

• reports (list[dsgrid.query.models.ReportInputModel])

• sort_columns (list[str])

• table_format (dsgrid.dataset.models.PivotedTableFormatModel | dsgrid.dataset.models.UnpivotedTableFormatModel)

• time_zone (str | None)

Validators:

• check_column_type » all fields

• check_format » output_format

• check_pivot_dimension_type » all fields

field aggregate_each_dataset: bool = False

If True, aggregate each dataset before applying the expression to create one overall dataset. This parameter must be set to True for queries that will be adding or subtracting datasets with different dimensionality. Defaults to False, which corresponds to the default behavior of performing one aggregation on the overall dataset. WARNING: For a standard query that performs a union of datasets, setting this value to True could produce rows with duplicate dimension combinations, especially if one or more dimensions are also dropped.

Validated by:

• check_column_type

• check_pivot_dimension_type

field aggregations: list[AggregationModel] = []

Defines how to aggregate dimensions

Validated by:

• check_column_type

• check_pivot_dimension_type

field column_type: ColumnType = ColumnType.DIMENSION_NAMES

Whether to make the result table columns dimension types. Default behavior is to use dimension names. In order to register a result table as a derived dataset, this must be set to dimension_types.

Validated by:

• check_column_type

• check_pivot_dimension_type

field dimension_filters: list[Annotated[DimensionFilterExpressionModel | DimensionFilterExpressionRawModel | DimensionFilterColumnOperatorModel | DimensionFilterBetweenColumnOperatorModel | SubsetDimensionFilterModel | SupplementalDimensionFilterColumnOperatorModel, FieldInfo(annotation=NoneType, required=True, discriminator='filter_type')]] = []

Filters to apply to the result. Must contain columns in the result.

Validated by:

• check_column_type

• check_pivot_dimension_type

field output_format: str = 'parquet'

Output file format: csv or parquet

Validated by:

• check_column_type

• check_format

• check_pivot_dimension_type

field replace_ids_with_names: bool = False

Replace dimension record IDs with their names in result tables.

Validated by:

• check_column_type

• check_pivot_dimension_type

field reports: list[ReportInputModel] = []

Run these pre-defined reports on the result.

Validated by:

• check_column_type

• check_pivot_dimension_type

field sort_columns: list[str] = []

Sort the results by these dimension names.

Validated by:

• check_column_type

• check_pivot_dimension_type

field table_format: Annotated[PivotedTableFormatModel | UnpivotedTableFormatModel, FieldInfo(annotation=NoneType, required=True, title='table_format', description='Defines the format of the value columns of the result table.', discriminator='format_type')] = UnpivotedTableFormatModel(format_type=<TableFormatType.UNPIVOTED: 'unpivoted'>)

Defines the format of the value columns of the result table.

Validated by:

• check_column_type

• check_pivot_dimension_type

field time_zone: str | None = None

Convert the results to this time zone.

Validated by:

• check_column_type

• check_pivot_dimension_type

validator check_column_type  »  all fields

validator check_format  »  output_format

validator check_pivot_dimension_type  »  all fields

pydantic model dsgrid.query.models.DatasetModel

Specifies the datasets to use in a project query.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Fields:

• dataset_id (str)

• expression (str | None)

• params (dsgrid.query.models.ProjectQueryDatasetParamsModel)

• source_datasets (list[dsgrid.query.models.StandaloneDatasetModel | dsgrid.query.models.ProjectionDatasetModel])

Validators:

• handle_expression » expression

field dataset_id: str [Required]

Identifier for the resulting dataset

field expression: str | None = None

Expression to combine datasets. Default is to take a union of all datasets.

Validated by:

• handle_expression

field params: ProjectQueryDatasetParamsModel = ProjectQueryDatasetParamsModel(dimension_filters=[])

Parameters affecting datasets. Used for caching intermediate tables.

field source_datasets: list[Annotated[StandaloneDatasetModel | ProjectionDatasetModel, FieldInfo(annotation=NoneType, required=True, discriminator='dataset_type')]] [Required]

Datasets from which to read. Each must be of type DatasetBaseModel.

validator handle_expression  »  expression

pydantic model dsgrid.query.models.StandaloneDatasetModel

A dataset with energy use data.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Fields:

• dataset_id (str)

• dataset_type (Literal[dsgrid.query.models.DatasetType.STANDALONE])

field dataset_id: str [Required]

Dataset identifier

field dataset_type: Literal[DatasetType.STANDALONE] = DatasetType.STANDALONE

get_dataset_id() → str

Return the primary dataset ID.

Return type:

str

list_source_dataset_ids() → list[str]

Return a list of all source dataset IDs.

pydantic model dsgrid.query.models.ProjectionDatasetModel

A dataset with growth rates that can be applied to a standalone dataset.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Fields:

• base_year (int | None)

• construction_method (dsgrid.query.models.DatasetConstructionMethod)

• dataset_id (str)

• dataset_type (Literal[dsgrid.query.models.DatasetType.PROJECTION])

• growth_rate_dataset_id (str)

• initial_value_dataset_id (str)

field base_year: int | None = None

Base year of the dataset to use in growth rate application. Must be a year defined in the principal dataset’s model year dimension. If None, there must be only one model year in that dimension and it will be used.

field construction_method: DatasetConstructionMethod = DatasetConstructionMethod.EXPONENTIAL_GROWTH

Specifier for the code that applies the growth rate to the principal dataset

field dataset_id: str [Required]

Identifier for the resulting dataset

field dataset_type: Literal[DatasetType.PROJECTION] = DatasetType.PROJECTION

field growth_rate_dataset_id: str [Required]

Growth rate dataset identifier to apply to the principal dataset

field initial_value_dataset_id: str [Required]

Principal dataset identifier

get_dataset_id() → str

Return the primary dataset ID.

Return type:

str

list_source_dataset_ids() → list[str]

Return a list of all source dataset IDs.

pydantic model dsgrid.query.models.AggregationModel

Aggregate on one or more dimensions.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Fields:

• aggregation_function (Any)

• dimensions (dsgrid.query.models.DimensionNamesModel)

Validators:

• check_aggregation_function » aggregation_function

• check_for_metric » dimensions

field aggregation_function: Any = None

Must be a function name in pyspark.sql.functions

Validated by:

• check_aggregation_function

field dimensions: DimensionNamesModel [Required]

Dimensions on which to aggregate

Validated by:

• check_for_metric

validator check_aggregation_function  »  aggregation_function

validator check_for_metric  »  dimensions

iter_dimensions_to_keep() → Generator[tuple[DimensionType, ColumnModel], None, None]

Yield the dimension type and ColumnModel for each dimension to keep.

list_dropped_dimensions() → list[DimensionType]

Return a list of dimension types that will be dropped by the aggregation.

serialize_aggregation_function(function, _)

pydantic model dsgrid.query.models.ColumnModel

Defines one column in a SQL aggregation statement.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Fields:

• alias (str | None)

• dimension_name (str)

• function (Any)

Validators:

• handle_alias » alias

• handle_function » function

field alias: str | None = None

Name of the resulting column.

Validated by:

• handle_alias

field dimension_name: str [Required]

field function: Any = None

Function or name of function in pyspark.sql.functions.

Validated by:

• handle_function

get_column_name()

validator handle_alias  »  alias

validator handle_function  »  function

serialize_function(function, _)

pydantic model dsgrid.query.models.FilteredDatasetModel

Filters to apply to a dataset

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Fields:

• dataset_id (str)

• filters (list[dsgrid.dimension.dimension_filters.DimensionFilterExpressionModel | dsgrid.dimension.dimension_filters.DimensionFilterExpressionRawModel | dsgrid.dimension.dimension_filters.DimensionFilterColumnOperatorModel | dsgrid.dimension.dimension_filters.DimensionFilterBetweenColumnOperatorModel | dsgrid.dimension.dimension_filters.SubsetDimensionFilterModel | dsgrid.dimension.dimension_filters.SupplementalDimensionFilterColumnOperatorModel])

field dataset_id: str [Required]

Dataset ID

field filters: list[Annotated[DimensionFilterExpressionModel | DimensionFilterExpressionRawModel | DimensionFilterColumnOperatorModel | DimensionFilterBetweenColumnOperatorModel | SubsetDimensionFilterModel | SupplementalDimensionFilterColumnOperatorModel, FieldInfo(annotation=NoneType, required=True, discriminator='filter_type')]] [Required]

pydantic model dsgrid.query.models.ReportInputModel

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Fields:

• inputs (Any)

• report_type (dsgrid.query.models.ReportType)

field inputs: Any = None

field report_type: ReportType [Required]

pydantic model dsgrid.query.models.SparkConfByDataset

Defines a custom Spark configuration to use while running a query on a dataset.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Fields:

• conf (dict[str, Any])

• dataset_id (str)

field conf: dict[str, Any] [Required]

field dataset_id: str [Required]

class dsgrid.query.models.ColumnType(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

Defines what the columns of a dataset table represent.

class dsgrid.query.models.DatasetType(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

Defines the type of a dataset in a query.

class dsgrid.query.models.ReportType(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

Pre-defined reports

Submission

class dsgrid.query.query_submitter.ProjectQuerySubmitter(project: Project, *args, **kwargs)

Submits queries for a project.

submit(**kwargs)

Submit a query for execution

Examples

from dsgrid.dimension.base_models import DimensionType
from dsgrid.registry.registry_manager import RegistryManager
from dsgrid.registry.registry_database import DatabaseConnection
from dsgrid.query.models import (
    AggregationModel,
    DatasetModel,
    DimensionQueryNamesModel,
    ProjectQueryParamsModel,
    ProjectQueryModel,
    QueryResultParamsModel,
    StandaloneDatasetModel,
)
from dsgrid.query.query_submitter import ProjectQuerySubmitter

manager = RegistryManager.load(
    DatabaseConnection(
        hostname="dsgrid-registry.hpc.nrel.gov",
        database="standard-scenarios",
    ),
    offline_mode=True
)
project = manager.project_manager.load_project("dsgrid_conus_2022")
query = ProjectQueryModel(
    name="Total Electricity Use By State and Sector",
    project=ProjectQueryParamsModel(
        project_id="dsgrid_conus_2022",
        dataset=DatasetModel(
            dataset_id="electricity_use",
            source_datasets=[
                StandaloneDatasetModel(dataset_id="comstock_conus_2022_projected"),
                StandaloneDatasetModel(dataset_id="resstock_conus_2022_projected"),
                StandaloneDatasetModel(dataset_id="tempo_conus_2022_mapped"),
            ],
        ),
    ),
    result=QueryResultParamsModel(
        aggregations=[
            AggregationModel(
                dimensions=DimensionQueryNamesModel(
                    geography=["state"],
                    metric=["electricity_collapsed"],
                    model_year=[],
                    scenario=[],
                    sector=["sector"],
                    subsector=[],
                    time=[],
                    weather_year=[],
                ),
                aggregation_function="sum",
            ),
        ],
    ),
)