sequential row noising (#509)

Category: feature
JIRA issue: MIC-5982
Add tests for sequential row noising.
Read in data using load_standard_dataset and noise by looping through noise types, all shard-wise except RI with ACS, CPS, and WIC.

Testing
Ran new tests with ACS, CPS, and WIC on USA and RI and SSA on RI.

Author: hussain-jafari

src/pseudopeople/filter.py

@@ -1,9 +1,15 @@
 from __future__ import annotations
 
+from collections.abc import Sequence
 from dataclasses import dataclass
+from typing import cast
 
 import pandas as pd
 
+from pseudopeople.constants.metadata import DatasetNames
+from pseudopeople.schema_entities import DatasetSchema
+from pseudopeople.utilities import get_state_abbreviation
+
 
 @dataclass
 class DataFilter:
@@ -13,3 +19,37 @@ class DataFilter:
 
     def to_tuple(self) -> tuple[str, str, str | int | pd.Timestamp]:
         return self.column_name, self.operator, self.value
+
+
+def get_data_filters(
+    dataset_schema: DatasetSchema, year: int | None = 2020, state: str | None = None
+) -> Sequence[DataFilter]:
+    filters = []
+    if dataset_schema.has_state_filter and state is not None:
+        state_column = cast(str, dataset_schema.state_column_name)
+        filters.append(DataFilter(state_column, "==", get_state_abbreviation(state)))
+
+    if year is not None:
+        try:
+            if dataset_schema.has_year_lower_filter:
+                date_lower_filter = DataFilter(
+                    dataset_schema.date_column_name,
+                    ">=",
+                    pd.Timestamp(year=year, month=1, day=1),
+                )
+                filters.append(date_lower_filter)
+
+            if dataset_schema.has_year_upper_filter:
+                date_upper_filter = DataFilter(
+                    dataset_schema.date_column_name,
+                    "<=",
+                    pd.Timestamp(year=year, month=12, day=31),
+                )
+                filters.append(date_upper_filter)
+        except (pd.errors.OutOfBoundsDatetime, ValueError):
+            raise ValueError(f"Invalid year provided: '{year}'")
+
+        if dataset_schema.has_exact_year_filter:
+            filters.append(DataFilter(dataset_schema.date_column_name, "==", year))
+
+    return filters

src/pseudopeople/interface.py

@@ -14,7 +14,7 @@
 from pseudopeople.constants import paths
 from pseudopeople.dataset import noise_data
 from pseudopeople.exceptions import DataSourceError
-from pseudopeople.filter import DataFilter
+from pseudopeople.filter import DataFilter, get_data_filters
 from pseudopeople.loader import load_standard_dataset
 from pseudopeople.schema_entities import DATASET_SCHEMAS, DatasetSchema
 from pseudopeople.utilities import (
@@ -196,6 +196,9 @@ def _generate_dataset(
 
 
 def validate_source_compatibility(source: Path, dataset_schema: DatasetSchema) -> None:
+    """Validate that a given source is compatible with the provided dataset schema by checking that
+    1) data exist for said schema in the provided source path and that 2) the data is the expected version
+    as specified in its CHANGELOG."""
     # TODO [MIC-4546]: Clean this up w/ metadata and update test_interface.py tests to be generic
     directories = [x.name for x in source.iterdir() if x.is_dir()]
     if dataset_schema.name not in directories:
@@ -341,12 +344,7 @@ def generate_decennial_census(
         The simulated population has no data for this dataset in the
         specified year or state.
     """
-    filters: list[DataFilter] = []
-    if year is not None:
-        filters.append(DataFilter(DATASET_SCHEMAS.census.date_column_name, "==", year))
-    if state is not None:
-        state_column_name = cast(str, DATASET_SCHEMAS.census.state_column_name)
-        filters.append(DataFilter(state_column_name, "==", get_state_abbreviation(state)))
+    filters: Sequence[DataFilter] = get_data_filters(DATASET_SCHEMAS.census, year, state)
     return _generate_dataset(
         DATASET_SCHEMAS.census,
         source,
@@ -472,26 +470,9 @@ def generate_american_community_survey(
         The simulated population has no data for this dataset in the
         specified year or state.
     """
-    filters = []
+    filters: Sequence[DataFilter] = get_data_filters(DATASET_SCHEMAS.acs, year, state)
     if year is not None:
-        try:
-            date_lower_filter = DataFilter(
-                DATASET_SCHEMAS.acs.date_column_name,
-                ">=",
-                pd.Timestamp(year=year, month=1, day=1),
-            )
-            date_upper_filter = DataFilter(
-                DATASET_SCHEMAS.acs.date_column_name,
-                "<=",
-                pd.Timestamp(year=year, month=12, day=31),
-            )
-            filters.extend([date_lower_filter, date_upper_filter])
-        except (pd.errors.OutOfBoundsDatetime, ValueError):
-            raise ValueError(f"Invalid year provided: '{year}'")
         seed = seed * 10_000 + year
-    if state is not None:
-        state_column = cast(str, DATASET_SCHEMAS.acs.state_column_name)
-        filters.append(DataFilter(state_column, "==", get_state_abbreviation(state)))
     return _generate_dataset(
         DATASET_SCHEMAS.acs, source, seed, config, filters, verbose, engine_name=engine
     )
@@ -612,26 +593,9 @@ def generate_current_population_survey(
         The simulated population has no data for this dataset in the
         specified year or state.
     """
-    filters = []
+    filters: Sequence[DataFilter] = get_data_filters(DATASET_SCHEMAS.cps, year, state)
     if year is not None:
-        try:
-            date_lower_filter = DataFilter(
-                DATASET_SCHEMAS.acs.date_column_name,
-                ">=",
-                pd.Timestamp(year=year, month=1, day=1),
-            )
-            date_upper_filter = DataFilter(
-                DATASET_SCHEMAS.acs.date_column_name,
-                "<=",
-                pd.Timestamp(year=year, month=12, day=31),
-            )
-            filters.extend([date_lower_filter, date_upper_filter])
-        except (pd.errors.OutOfBoundsDatetime, ValueError):
-            raise ValueError(f"Invalid year provided: '{year}'")
         seed = seed * 10_000 + year
-    if state is not None:
-        state_column = cast(str, DATASET_SCHEMAS.cps.state_column_name)
-        filters.append(DataFilter(state_column, "==", get_state_abbreviation(state)))
     return _generate_dataset(
         DATASET_SCHEMAS.cps, source, seed, config, filters, verbose, engine_name=engine
     )
@@ -743,13 +707,9 @@ def generate_taxes_w2_and_1099(
         The simulated population has no data for this dataset in the
         specified year or state.
     """
-    filters = []
+    filters: Sequence[DataFilter] = get_data_filters(DATASET_SCHEMAS.tax_w2_1099, year, state)
     if year is not None:
-        filters.append(DataFilter(DATASET_SCHEMAS.tax_w2_1099.date_column_name, "==", year))
         seed = seed * 10_000 + year
-    if state is not None:
-        state_column = cast(str, DATASET_SCHEMAS.tax_w2_1099.state_column_name)
-        filters.append(DataFilter(state_column, "==", get_state_abbreviation(state)))
     return _generate_dataset(
         DATASET_SCHEMAS.tax_w2_1099,
         source,
@@ -878,13 +838,9 @@ def generate_women_infants_and_children(
         The simulated population has no data for this dataset in the
         specified year or state.
     """
-    filters = []
+    filters: Sequence[DataFilter] = get_data_filters(DATASET_SCHEMAS.wic, year, state)
     if year is not None:
-        filters.append(DataFilter(DATASET_SCHEMAS.wic.date_column_name, "==", year))
         seed = seed * 10_000 + year
-    if state is not None:
-        state_column = cast(str, DATASET_SCHEMAS.wic.state_column_name)
-        filters.append(DataFilter(state_column, "==", get_state_abbreviation(state)))
     return _generate_dataset(
         DATASET_SCHEMAS.wic, source, seed, config, filters, verbose, engine_name=engine
     )
@@ -984,18 +940,8 @@ def generate_social_security(
         The simulated population has no data for this dataset in the
         specified year or any prior years.
     """
-    filters = []
+    filters: Sequence[DataFilter] = get_data_filters(DATASET_SCHEMAS.ssa, year)
     if year is not None:
-        try:
-            filters.append(
-                DataFilter(
-                    DATASET_SCHEMAS.ssa.date_column_name,
-                    "<=",
-                    pd.Timestamp(year=year, month=12, day=31),
-                )
-            )
-        except (pd.errors.OutOfBoundsDatetime, ValueError):
-            raise ValueError(f"Invalid year provided: '{year}'")
         seed = seed * 10_000 + year
     return _generate_dataset(
         DATASET_SCHEMAS.ssa, source, seed, config, filters, verbose, engine_name=engine
@@ -1108,13 +1054,9 @@ def generate_taxes_1040(
         The simulated population has no data for this dataset in the
         specified year or state.
     """
-    filters = []
+    filters: Sequence[DataFilter] = get_data_filters(DATASET_SCHEMAS.tax_1040, year, state)
     if year is not None:
-        filters.append(DataFilter(DATASET_SCHEMAS.tax_1040.date_column_name, "==", year))
         seed = seed * 10_000 + year
-    if state is not None:
-        state_column = cast(str, DATASET_SCHEMAS.tax_1040.state_column_name)
-        filters.append(DataFilter(state_column, "==", get_state_abbreviation(state)))
     return _generate_dataset(
         DATASET_SCHEMAS.tax_1040,
         source,

src/pseudopeople/noise_functions.py

@@ -29,6 +29,8 @@
 )
 
 if TYPE_CHECKING:
+    import dask.dataframe as dd
+
     from pseudopeople.configuration.noise_configuration import NoiseConfiguration
     from pseudopeople.dataset import Dataset
 
@@ -95,7 +97,7 @@ def apply_do_not_respond(
 
 # Helper function to format group dataframe and merging with their dependents
 def merge_dependents_and_guardians(
-    dependents_df: pd.DataFrame, full_data: pd.DataFrame
+    dependents_df: pd.DataFrame | dd.DataFrame, full_data: pd.DataFrame | dd.DataFrame
 ) -> pd.DataFrame:
     # Merge dependents with their guardians. We have to merge twice to check
     # if either guardian is living at a separate location from the dependent.

src/pseudopeople/schema_entities.py

@@ -590,6 +590,10 @@ class DatasetSchema:
     date_format: str
     state_column_name: str | None
     row_noise_types: tuple[RowNoiseType, ...]
+    has_state_filter: bool
+    has_year_lower_filter: bool
+    has_year_upper_filter: bool
+    has_exact_year_filter: bool
 
 
 class __DatasetSchemas(NamedTuple):
@@ -626,6 +630,10 @@ class __DatasetSchemas(NamedTuple):
             # NOISE_TYPES.duplication,
         ),
         date_format=DATEFORMATS.MM_DD_YYYY,
+        has_state_filter=True,
+        has_year_lower_filter=False,
+        has_year_upper_filter=False,
+        has_exact_year_filter=True,
     )
     acs: DatasetSchema = DatasetSchema(
         DatasetNames.ACS,
@@ -657,6 +665,10 @@ class __DatasetSchemas(NamedTuple):
             # NOISE_TYPES.duplication,
         ),
         date_format=DATEFORMATS.MM_DD_YYYY,
+        has_state_filter=True,
+        has_year_lower_filter=True,
+        has_year_upper_filter=True,
+        has_exact_year_filter=False,
     )
     cps: DatasetSchema = DatasetSchema(
         DatasetNames.CPS,
@@ -686,6 +698,10 @@ class __DatasetSchemas(NamedTuple):
             # NOISE_TYPES.duplication,
         ),
         date_format=DATEFORMATS.MM_DD_YYYY,
+        has_state_filter=True,
+        has_year_lower_filter=True,
+        has_year_upper_filter=True,
+        has_exact_year_filter=False,
     )
     wic: DatasetSchema = DatasetSchema(
         DatasetNames.WIC,
@@ -713,6 +729,10 @@ class __DatasetSchemas(NamedTuple):
             # NOISE_TYPES.duplication,
         ),
         date_format=DATEFORMATS.MMDDYYYY,
+        has_state_filter=True,
+        has_year_lower_filter=False,
+        has_year_upper_filter=False,
+        has_exact_year_filter=True,
     )
     ssa: DatasetSchema = DatasetSchema(
         DatasetNames.SSA,
@@ -734,6 +754,10 @@ class __DatasetSchemas(NamedTuple):
             # NOISE_TYPES.duplication,
         ),
         date_format=DATEFORMATS.YYYYMMDD,
+        has_state_filter=False,
+        has_year_lower_filter=False,
+        has_year_upper_filter=True,
+        has_exact_year_filter=False,
     )
     tax_w2_1099: DatasetSchema = DatasetSchema(
         DatasetNames.TAXES_W2_1099,
@@ -770,6 +794,10 @@ class __DatasetSchemas(NamedTuple):
             # NOISE_TYPES.duplication,
         ),
         date_format=DATEFORMATS.MM_DD_YYYY,
+        has_state_filter=True,
+        has_year_lower_filter=False,
+        has_year_upper_filter=False,
+        has_exact_year_filter=True,
     )
     tax_1040: DatasetSchema = DatasetSchema(
         DatasetNames.TAXES_1040,
@@ -809,6 +837,10 @@ class __DatasetSchemas(NamedTuple):
         state_column_name=COLUMNS.mailing_state.name,
         row_noise_types=(NOISE_TYPES.omit_row,),
         date_format=DATEFORMATS.MM_DD_YYYY,
+        has_state_filter=True,
+        has_year_lower_filter=False,
+        has_year_upper_filter=False,
+        has_exact_year_filter=True,
     )
 
     ##################