dsgrid-toolkit 0.2.0__py3-none-any.whl

dsgrid/dataset/dataset_schema_handler_one_table.py

@@ -0,0 +1,196 @@
+import logging
+from typing import Self
+
+from dsgrid.common import VALUE_COLUMN
+from dsgrid.config.dataset_config import DatasetConfig
+from dsgrid.config.project_config import ProjectConfig
+from dsgrid.config.simple_models import DimensionSimpleModel
+from dsgrid.config.time_dimension_base_config import TimeDimensionBaseConfig
+from dsgrid.dataset.models import TableFormatType
+from dsgrid.query.models import DatasetQueryModel
+from dsgrid.registry.data_store_interface import DataStoreInterface
+from dsgrid.spark.types import (
+    DataFrame,
+    StringType,
+)
+from dsgrid.utils.dataset import (
+    convert_types_if_necessary,
+)
+from dsgrid.utils.spark import (
+    check_for_nulls,
+    read_dataframe,
+)
+from dsgrid.utils.timing import timer_stats_collector, track_timing
+from dsgrid.dataset.dataset_schema_handler_base import DatasetSchemaHandlerBase
+from dsgrid.dimension.base_models import DimensionType
+from dsgrid.exceptions import DSGInvalidDataset
+from dsgrid.query.query_context import QueryContext
+
+logger = logging.getLogger(__name__)
+
+
+class OneTableDatasetSchemaHandler(DatasetSchemaHandlerBase):
+    """define interface/required behaviors for ONE_TABLE dataset schema"""
+
+    def __init__(self, load_data_df, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._load_data = load_data_df
+
+    @classmethod
+    def load(
+        cls,
+        config: DatasetConfig,
+        *args,
+        store: DataStoreInterface | None = None,
+        **kwargs,
+    ) -> Self:
+        if store is None:
+            df = read_dataframe(config.load_data_path)
+        else:
+            df = store.read_table(config.model.dataset_id, config.model.version)
+        load_data_df = config.add_trivial_dimensions(df)
+        load_data_df = convert_types_if_necessary(load_data_df)
+        time_dim = config.get_time_dimension()
+        if time_dim is not None:
+            load_data_df = time_dim.convert_time_format(load_data_df)
+        return cls(load_data_df, config, *args, **kwargs)
+
+    @track_timing(timer_stats_collector)
+    def check_consistency(self, missing_dimension_associations: DataFrame | None) -> None:
+        self._check_one_table_data_consistency()
+        self._check_dimension_associations(missing_dimension_associations)
+
+    @track_timing(timer_stats_collector)
+    def check_time_consistency(self):
+        time_dim = self._config.get_time_dimension()
+        if time_dim is not None:
+            if time_dim.supports_chronify():
+                self._check_dataset_time_consistency_with_chronify()
+            else:
+                self._check_dataset_time_consistency(self._load_data)
+
+    @track_timing(timer_stats_collector)
+    def _check_one_table_data_consistency(self):
+        """Dimension check in load_data, excludes time:
+        * check that data matches record for each dimension.
+        * check that all data dimension combinations exist. Time is handled separately.
+        * Check for any NULL values in dimension columns.
+        """
+        logger.info("Check one table dataset consistency.")
+        dimension_types = set()
+        time_dim = self._config.get_time_dimension()
+        time_columns: set[str] = set()
+        if time_dim is not None:
+            time_columns = set(time_dim.get_load_data_time_columns())
+        assert (
+            self._config.get_table_format_type() == TableFormatType.UNPIVOTED
+        ), self._config.get_table_format_type()
+        self._check_load_data_unpivoted_value_column(self._load_data)
+        allowed_columns = DimensionType.get_allowed_dimension_column_names().union(time_columns)
+        allowed_columns.add(VALUE_COLUMN)
+
+        schema = self._load_data.schema
+        for column in self._load_data.columns:
+            if column not in allowed_columns:
+                msg = f"{column=} is not expected in load_data"
+                raise DSGInvalidDataset(msg)
+            if not (column in time_columns or column == VALUE_COLUMN):
+                dim_type = DimensionType.from_column(column)
+                if schema[column].dataType != StringType():
+                    msg = f"dimension column {column} must have data type = StringType"
+                    raise DSGInvalidDataset(msg)
+                dimension_types.add(dim_type)
+        check_for_nulls(self._load_data)
+
+    def _get_load_data_table(self) -> DataFrame:
+        return self._load_data
+
+    @track_timing(timer_stats_collector)
+    def filter_data(self, dimensions: list[DimensionSimpleModel], store: DataStoreInterface):
+        assert (
+            self._config.get_table_format_type() == TableFormatType.UNPIVOTED
+        ), self._config.get_table_format_type()
+        load_df = self._load_data
+        df_columns = set(load_df.columns)
+        stacked_columns = set()
+        for dim in dimensions:
+            column = dim.dimension_type.value
+            if column in df_columns:
+                load_df = load_df.filter(load_df[column].isin(dim.record_ids))
+                stacked_columns.add(column)
+
+        drop_columns = []
+        for dim in self._config.model.trivial_dimensions:
+            col = dim.value
+            count = load_df.select(col).distinct().count()
+            assert count == 1, f"{dim}: {count}"
+            drop_columns.append(col)
+        load_df = load_df.drop(*drop_columns)
+
+        store.replace_table(load_df, self.dataset_id, self._config.model.version)
+        logger.info("Rewrote simplified %s", self._config.model.dataset_id)
+
+    def make_project_dataframe(
+        self, context: QueryContext, project_config: ProjectConfig
+    ) -> DataFrame:
+        plan = context.model.project.get_dataset_mapping_plan(self.dataset_id)
+        if plan is None:
+            plan = self.build_default_dataset_mapping_plan()
+        with context.dataset_mapping_manager(self.dataset_id, plan) as mapping_manager:
+            ld_df = mapping_manager.try_read_checkpointed_table()
+            if ld_df is None:
+                ld_df = self._load_data
+                ld_df = self._prefilter_stacked_dimensions(context, ld_df)
+                ld_df = self._prefilter_time_dimension(context, ld_df)
+
+            ld_df = self._remap_dimension_columns(
+                ld_df,
+                mapping_manager,
+                filtered_records=context.get_record_ids(),
+            )
+            ld_df = self._apply_fraction(ld_df, {VALUE_COLUMN}, mapping_manager)
+            project_metric_records = self._get_project_metric_records(project_config)
+            ld_df = self._convert_units(ld_df, project_metric_records, mapping_manager)
+            input_dataset = project_config.get_dataset(self._config.model.dataset_id)
+            ld_df = self._convert_time_dimension(
+                load_data_df=ld_df,
+                to_time_dim=project_config.get_base_time_dimension(),
+                value_column=VALUE_COLUMN,
+                mapping_manager=mapping_manager,
+                wrap_time_allowed=input_dataset.wrap_time_allowed,
+                time_based_data_adjustment=input_dataset.time_based_data_adjustment,
+                to_geo_dim=project_config.get_base_dimension(DimensionType.GEOGRAPHY),
+            )
+            return self._finalize_table(context, ld_df, project_config)
+
+    def make_mapped_dataframe(
+        self, context: QueryContext, time_dimension: TimeDimensionBaseConfig | None = None
+    ) -> DataFrame:
+        query = context.model
+        assert isinstance(query, DatasetQueryModel)
+        plan = query.mapping_plan
+        if plan is None:
+            plan = self.build_default_dataset_mapping_plan()
+        geography_dimension = self._get_mapping_to_dimension(DimensionType.GEOGRAPHY)
+        metric_dimension = self._get_mapping_to_dimension(DimensionType.METRIC)
+        with context.dataset_mapping_manager(self.dataset_id, plan) as mapping_manager:
+            ld_df = mapping_manager.try_read_checkpointed_table()
+            if ld_df is None:
+                ld_df = self._load_data
+
+            ld_df = self._remap_dimension_columns(ld_df, mapping_manager)
+            ld_df = self._apply_fraction(ld_df, {VALUE_COLUMN}, mapping_manager)
+            if metric_dimension is not None:
+                metric_records = metric_dimension.get_records_dataframe()
+                ld_df = self._convert_units(ld_df, metric_records, mapping_manager)
+            if time_dimension is not None:
+                ld_df = self._convert_time_dimension(
+                    load_data_df=ld_df,
+                    to_time_dim=time_dimension,
+                    value_column=VALUE_COLUMN,
+                    mapping_manager=mapping_manager,
+                    wrap_time_allowed=query.wrap_time_allowed,
+                    time_based_data_adjustment=query.time_based_data_adjustment,
+                    to_geo_dim=geography_dimension,
+                )
+        return ld_df

dsgrid/dataset/dataset_schema_handler_standard.py

@@ -0,0 +1,303 @@
+import logging
+from typing import Self
+
+from dsgrid.common import SCALING_FACTOR_COLUMN, VALUE_COLUMN
+from dsgrid.config.dataset_config import DatasetConfig
+from dsgrid.config.project_config import ProjectConfig
+from dsgrid.config.simple_models import DimensionSimpleModel
+from dsgrid.config.time_dimension_base_config import TimeDimensionBaseConfig
+from dsgrid.dataset.models import TableFormatType
+from dsgrid.dataset.dataset_schema_handler_base import DatasetSchemaHandlerBase
+from dsgrid.dimension.base_models import DimensionType
+from dsgrid.exceptions import DSGInvalidDataset
+from dsgrid.query.models import DatasetQueryModel
+from dsgrid.query.query_context import QueryContext
+from dsgrid.registry.data_store_interface import DataStoreInterface
+from dsgrid.spark.functions import (
+    cache,
+    coalesce,
+    collect_list,
+    except_all,
+    intersect,
+    unpersist,
+)
+from dsgrid.spark.types import (
+    DataFrame,
+    StringType,
+)
+from dsgrid.utils.dataset import (
+    apply_scaling_factor,
+    convert_types_if_necessary,
+)
+from dsgrid.utils.spark import (
+    check_for_nulls,
+    read_dataframe,
+)
+from dsgrid.utils.timing import Timer, timer_stats_collector, track_timing
+
+
+logger = logging.getLogger(__name__)
+
+
+class StandardDatasetSchemaHandler(DatasetSchemaHandlerBase):
+    """define interface/required behaviors for STANDARD dataset schema"""
+
+    def __init__(self, load_data_df, load_data_lookup, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._load_data = load_data_df
+        self._load_data_lookup = load_data_lookup
+
+    @classmethod
+    def load(
+        cls,
+        config: DatasetConfig,
+        *args,
+        store: DataStoreInterface | None = None,
+        **kwargs,
+    ) -> Self:
+        if store is None:
+            load_data_df = read_dataframe(config.load_data_path)
+            load_data_lookup = read_dataframe(config.load_data_lookup_path)
+        else:
+            load_data_df = store.read_table(config.model.dataset_id, config.model.version)
+            load_data_lookup = store.read_lookup_table(
+                config.model.dataset_id, config.model.version
+            )
+
+        load_data_df = convert_types_if_necessary(load_data_df)
+        time_dim = config.get_time_dimension()
+        if time_dim is not None:
+            load_data_df = time_dim.convert_time_format(load_data_df)
+        load_data_lookup = config.add_trivial_dimensions(load_data_lookup)
+        load_data_lookup = convert_types_if_necessary(load_data_lookup)
+        return cls(load_data_df, load_data_lookup, config, *args, **kwargs)
+
+    @track_timing(timer_stats_collector)
+    def check_consistency(self, missing_dimension_associations: DataFrame | None) -> None:
+        self._check_lookup_data_consistency()
+        self._check_dataset_internal_consistency()
+        self._check_dimension_associations(missing_dimension_associations)
+
+    @track_timing(timer_stats_collector)
+    def check_time_consistency(self):
+        time_dim = self._config.get_time_dimension()
+        if time_dim is None:
+            return None
+
+        if time_dim.supports_chronify():
+            self._check_dataset_time_consistency_with_chronify()
+        else:
+            self._check_dataset_time_consistency(self._get_load_data_table())
+
+    def _get_load_data_table(self) -> DataFrame:
+        return self._load_data.join(self._load_data_lookup, on="id")
+
+    def make_project_dataframe(
+        self, context: QueryContext, project_config: ProjectConfig
+    ) -> DataFrame:
+        lk_df = self._load_data_lookup
+        lk_df = self._prefilter_stacked_dimensions(context, lk_df)
+
+        plan = context.model.project.get_dataset_mapping_plan(self.dataset_id)
+        if plan is None:
+            plan = self.build_default_dataset_mapping_plan()
+        with context.dataset_mapping_manager(self.dataset_id, plan) as mapping_manager:
+            ld_df = mapping_manager.try_read_checkpointed_table()
+            if ld_df is None:
+                ld_df = self._load_data
+                ld_df = self._prefilter_stacked_dimensions(context, ld_df)
+                ld_df = self._prefilter_time_dimension(context, ld_df)
+                ld_df = ld_df.join(lk_df, on="id").drop("id")
+
+            ld_df = self._remap_dimension_columns(
+                ld_df,
+                mapping_manager,
+                filtered_records=context.get_record_ids(),
+            )
+            if SCALING_FACTOR_COLUMN in ld_df.columns:
+                ld_df = apply_scaling_factor(ld_df, VALUE_COLUMN, mapping_manager)
+
+            ld_df = self._apply_fraction(ld_df, {VALUE_COLUMN}, mapping_manager)
+            project_metric_records = self._get_project_metric_records(project_config)
+            ld_df = self._convert_units(ld_df, project_metric_records, mapping_manager)
+            input_dataset = project_config.get_dataset(self._config.model.dataset_id)
+            ld_df = self._convert_time_dimension(
+                load_data_df=ld_df,
+                to_time_dim=project_config.get_base_time_dimension(),
+                value_column=VALUE_COLUMN,
+                mapping_manager=mapping_manager,
+                wrap_time_allowed=input_dataset.wrap_time_allowed,
+                time_based_data_adjustment=input_dataset.time_based_data_adjustment,
+                to_geo_dim=project_config.get_base_dimension(DimensionType.GEOGRAPHY),
+            )
+            return self._finalize_table(context, ld_df, project_config)
+
+    def make_mapped_dataframe(
+        self,
+        context: QueryContext,
+        time_dimension: TimeDimensionBaseConfig | None = None,
+    ) -> DataFrame:
+        query = context.model
+        assert isinstance(query, DatasetQueryModel)
+        plan = query.mapping_plan
+        if plan is None:
+            plan = self.build_default_dataset_mapping_plan()
+        geography_dimension = self._get_mapping_to_dimension(DimensionType.GEOGRAPHY)
+        metric_dimension = self._get_mapping_to_dimension(DimensionType.METRIC)
+        with context.dataset_mapping_manager(self.dataset_id, plan) as mapping_manager:
+            ld_df = mapping_manager.try_read_checkpointed_table()
+            if ld_df is None:
+                ld_df = self._load_data
+                lk_df = self._load_data_lookup
+                ld_df = ld_df.join(lk_df, on="id").drop("id")
+
+            ld_df = self._remap_dimension_columns(
+                ld_df,
+                mapping_manager,
+            )
+            if SCALING_FACTOR_COLUMN in ld_df.columns:
+                ld_df = apply_scaling_factor(ld_df, VALUE_COLUMN, mapping_manager)
+
+            ld_df = self._apply_fraction(ld_df, {VALUE_COLUMN}, mapping_manager)
+            if metric_dimension is not None:
+                metric_records = metric_dimension.get_records_dataframe()
+                ld_df = self._convert_units(ld_df, metric_records, mapping_manager)
+            if time_dimension is not None:
+                ld_df = self._convert_time_dimension(
+                    load_data_df=ld_df,
+                    to_time_dim=time_dimension,
+                    value_column=VALUE_COLUMN,
+                    mapping_manager=mapping_manager,
+                    wrap_time_allowed=query.wrap_time_allowed,
+                    time_based_data_adjustment=query.time_based_data_adjustment,
+                    to_geo_dim=geography_dimension,
+                )
+        return ld_df
+
+    @track_timing(timer_stats_collector)
+    def _check_lookup_data_consistency(self):
+        """Dimension check in load_data_lookup, excludes time:
+        * check that data matches record for each dimension.
+        * check that all data dimension combinations exist. Time is handled separately.
+        * Check for any NULL values in dimension columns.
+        """
+        logger.info("Check lookup data consistency.")
+        found_id = False
+        dimension_types = set()
+        for col in self._load_data_lookup.columns:
+            if col == "id":
+                found_id = True
+                continue
+            if col == SCALING_FACTOR_COLUMN:
+                continue
+            if self._load_data_lookup.schema[col].dataType != StringType():
+                msg = f"dimension column {col} must have data type = StringType"
+                raise DSGInvalidDataset(msg)
+            dimension_types.add(DimensionType.from_column(col))
+
+        if not found_id:
+            msg = "load_data_lookup does not include an 'id' column"
+            raise DSGInvalidDataset(msg)
+
+        check_for_nulls(self._load_data_lookup)
+        load_data_dimensions = set(self._list_dimension_types_in_load_data(self._load_data))
+        expected_dimensions = {
+            d
+            for d in DimensionType.get_dimension_types_allowed_as_columns()
+            if d not in load_data_dimensions
+        }
+        missing_dimensions = expected_dimensions.difference(dimension_types)
+        if missing_dimensions:
+            msg = (
+                f"load_data_lookup is missing dimensions: {missing_dimensions}. "
+                "If these are trivial dimensions, make sure to specify them in the Dataset Config."
+            )
+
+    @track_timing(timer_stats_collector)
+    def _check_dataset_internal_consistency(self):
+        """Check load_data dimensions and id series."""
+        logger.info("Check dataset internal consistency.")
+        assert (
+            self._config.get_table_format_type() == TableFormatType.UNPIVOTED
+        ), self._config.get_table_format_type()
+        self._check_load_data_unpivoted_value_column(self._load_data)
+
+        time_dim = self._config.get_time_dimension()
+        time_columns: set[str] = set()
+        if time_dim is not None:
+            time_columns = set(time_dim.get_load_data_time_columns())
+        allowed_columns = (
+            DimensionType.get_allowed_dimension_column_names()
+            .union(time_columns)
+            .union({VALUE_COLUMN, "id", "scaling_factor"})
+        )
+
+        found_id = False
+        for column in self._load_data.columns:
+            if column not in allowed_columns:
+                msg = f"{column=} is not expected in load_data"
+                raise DSGInvalidDataset(msg)
+            if column == "id":
+                found_id = True
+
+        if not found_id:
+            msg = "load_data does not include an 'id' column"
+            raise DSGInvalidDataset(msg)
+
+        check_for_nulls(self._load_data)
+        ld_ids = self._load_data.select("id").distinct()
+        ldl_ids = self._load_data_lookup.select("id").distinct()
+        ldl_id_count = ldl_ids.count()
+        data_id_count = ld_ids.count()
+        joined = ld_ids.join(ldl_ids, on="id")
+        count = joined.count()
+
+        if data_id_count != count or ldl_id_count != count:
+            with Timer(timer_stats_collector, "show load_data and load_data_lookup ID diff"):
+                diff = except_all(ld_ids.unionAll(ldl_ids), intersect(ld_ids, ldl_ids))
+                # Only run the query once (with Spark). Number of rows shouldn't be a problem.
+                cache(diff)
+                diff_count = diff.count()
+                limit = 100
+                diff_list = diff.limit(limit).collect()
+                unpersist(diff)
+                logger.error(
+                    "load_data and load_data_lookup have %s different IDs. Limited to %s: %s",
+                    diff_count,
+                    limit,
+                    diff_list,
+                )
+            msg = f"Data IDs for {self._config.config_id} data/lookup are inconsistent"
+            raise DSGInvalidDataset(msg)
+
+    @track_timing(timer_stats_collector)
+    def filter_data(self, dimensions: list[DimensionSimpleModel], store: DataStoreInterface):
+        lookup = self._load_data_lookup
+        cache(lookup)
+        load_df = self._load_data
+        lookup_columns = set(lookup.columns)
+        for dim in dimensions:
+            column = dim.dimension_type.value
+            if column in lookup_columns:
+                lookup = lookup.filter(lookup[column].isin(dim.record_ids))
+
+        drop_columns = []
+        for dim in self._config.model.trivial_dimensions:
+            col = dim.value
+            count = lookup.select(col).distinct().count()
+            assert count == 1, f"{dim}: count"
+            drop_columns.append(col)
+        lookup = lookup.drop(*drop_columns)
+
+        lookup2 = coalesce(lookup, 1)
+        store.replace_lookup_table(lookup2, self.dataset_id, self._config.model.version)
+        ids = collect_list(lookup2.select("id").distinct(), "id")
+        load_df = self._load_data.filter(self._load_data.id.isin(ids))
+        ld_columns = set(load_df.columns)
+        for dim in dimensions:
+            column = dim.dimension_type.value
+            if column in ld_columns:
+                load_df = load_df.filter(load_df[column].isin(dim.record_ids))
+
+        store.replace_table(load_df, self.dataset_id, self._config.model.version)
+        logger.info("Rewrote simplified %s", self._config.model.dataset_id)

dsgrid/dataset/growth_rates.py

@@ -0,0 +1,162 @@
+import logging
+
+from dsgrid.exceptions import DSGInvalidQuery
+from dsgrid.query.models import ProjectionDatasetModel
+from dsgrid.spark.functions import cross_join, join_multiple_columns, sql_from_df
+from dsgrid.spark.types import DataFrame, F, IntegerType, use_duckdb
+from dsgrid.utils.spark import get_unique_values
+
+
+logger = logging.getLogger(__name__)
+
+
+def apply_exponential_growth_rate(
+    dataset: ProjectionDatasetModel,
+    initial_value_df: DataFrame,
+    growth_rate_df: DataFrame,
+    time_columns,
+    model_year_column,
+    value_columns,
+):
+    """Applies exponential growth rate to the initial_value dataframe as follows:
+    P(t) = P0*(1+r)^(t-t0)
+    where:
+      P(t): quantity at t
+      P0: initial quantity at t0, = P(t0)
+      r: growth rate (per time interval)
+      t-t0: number of time intervals
+
+
+    Parameters
+    ----------
+    dataset : ProjectionDatasetModel
+    initial_value_df : pyspark.sql.DataFrame
+    growth_rate_df : pyspark.sql.DataFrame
+    time_columns : set[str]
+    model_year_column : str
+    value_columns : set[str]
+
+    Returns
+    -------
+    pyspark.sql.DataFrame
+
+    """
+
+    initial_value_df, growth_rate_df = _process_exponential_growth_rate(
+        dataset,
+        initial_value_df,
+        growth_rate_df,
+        model_year_column,
+        value_columns,
+    )
+
+    df = apply_annual_multiplier(
+        initial_value_df,
+        growth_rate_df,
+        time_columns,
+        value_columns,
+    )
+
+    return df
+
+
+def apply_annual_multiplier(
+    initial_value_df: DataFrame,
+    growth_rate_df: DataFrame,
+    time_columns,
+    value_columns,
+):
+    """Applies annual growth rate to the initial_value dataframe as follows:
+    P(t) = P0 * r(t)
+    where:
+      P(t): quantity at year t
+      P0: initial quantity
+      r(t): growth rate per year t (relative to P0)
+
+    Parameters
+    ----------
+    dataset : ProjectionDatasetModel
+    initial_value_df : pyspark.sql.DataFrame
+    growth_rate_df : pyspark.sql.DataFrame
+    time_columns : set[str]
+    value_columns : set[str]
+
+    Returns
+    -------
+    pyspark.sql.DataFrame
+
+    """
+
+    def renamed(col):
+        return col + "_gr"
+
+    orig_columns = initial_value_df.columns
+
+    dim_columns = set(initial_value_df.columns) - value_columns - time_columns
+    df = join_multiple_columns(initial_value_df, growth_rate_df, list(dim_columns))
+    for column in df.columns:
+        if column in value_columns:
+            gr_column = renamed(column)
+            df = df.withColumn(column, df[column] * df[gr_column])
+
+    return df.select(*orig_columns)
+
+
+def _process_exponential_growth_rate(
+    dataset,
+    initial_value_df,
+    growth_rate_df,
+    model_year_column,
+    value_columns,
+):
+    def renamed(col):
+        return col + "_gr"
+
+    initial_value_df, base_year = _check_model_years(
+        dataset, initial_value_df, growth_rate_df, model_year_column
+    )
+
+    gr_df = growth_rate_df
+    for column in value_columns:
+        gr_col = renamed(column)
+        cols = ",".join([x for x in gr_df.columns if x not in (column, gr_col)])
+        if use_duckdb():
+            query = f"""
+                SELECT
+                    {cols}
+                    ,(1 + {column}) ** (CAST({model_year_column} AS INTEGER) - {base_year}) AS {gr_col}
+            """
+            gr_df = sql_from_df(gr_df, query)
+        else:
+            # Spark SQL uses POW instead of **, so keep the DataFrame API method.
+            gr_df = gr_df.withColumn(
+                gr_col,
+                F.pow(
+                    (1 + F.col(column)), F.col(model_year_column).cast(IntegerType()) - base_year
+                ),
+            ).drop(column)
+
+    return initial_value_df, gr_df
+
+
+def _check_model_years(dataset, initial_value_df, growth_rate_df, model_year_column):
+    iv_years = get_unique_values(initial_value_df, model_year_column)
+    iv_years_sorted = sorted((int(x) for x in iv_years))
+
+    if dataset.base_year is None:
+        base_year = iv_years_sorted[0]
+    elif dataset.base_year in iv_years:
+        base_year = dataset.base_year
+    else:
+        msg = f"ProjectionDatasetModel base_year={dataset.base_year} is not in {iv_years_sorted}"
+        raise DSGInvalidQuery(msg)
+
+    if len(iv_years) > 1:
+        # TODO #198: needs test case
+        initial_value_df = initial_value_df.filter(f"{model_year_column} == '{base_year}'")
+
+    initial_value_df = cross_join(
+        initial_value_df.drop(model_year_column),
+        growth_rate_df.select(model_year_column).distinct(),
+    )
+    return initial_value_df, base_year