intermediate changes trying sharded testing

Author: Hussain Jafari

src/pseudopeople/dataset.py

@@ -132,19 +132,17 @@ def _clean_input_data(self) -> None:
 
     def _reformat_dates_for_noising(self) -> None:
         """Formats date columns so they can be noised as strings."""
-        data = self.data.copy()
-
         for date_column in [COLUMNS.dob.name, COLUMNS.ssa_event_date.name]:
             # Format both the actual column, and the shadow version that will be used
             # to copy from a household member
             for column in [date_column, COPY_HOUSEHOLD_MEMBER_COLS.get(date_column)]:
-                if column in data.columns and isinstance(column, str):
+                if column in self.data.columns and isinstance(column, str):
                     # Avoid running strftime on large data, since that will
                     # re-parse the format string for each row
                     # https://github.com/pandas-dev/pandas/issues/44764
                     # Year is already guaranteed to be 4-digit: https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#timeseries-timestamp-limits
-                    is_na = data[column].isna()
-                    data_column = data.loc[~is_na, column]
+                    is_na = self.data[column].isna()
+                    data_column = self.data.loc[~is_na, column]
                     year_string = data_column.dt.year.astype(str)
                     month_string = _zfill_fast(data_column.dt.month.astype(str), 2)
                     day_string = _zfill_fast(data_column.dt.day.astype(str), 2)
@@ -159,10 +157,8 @@ def _reformat_dates_for_noising(self) -> None:
                             f"Invalid date format in {self.dataset_schema.name}."
                         )
 
-                    data[column] = pd.Series(np.nan, dtype=str)
-                    data.loc[~is_na, column] = result
-
-        self.data = data
+                    self.data[column] = pd.Series(np.nan, dtype=str)
+                    self.data.loc[~is_na, column] = result
 
     @staticmethod
     def drop_non_schema_columns(

src/pseudopeople/entity_types.py

@@ -149,7 +149,6 @@ def __call__(
         output_dtype = self.output_dtype_getter(input_dtype)
 
         dataset.data[column_name] = ensure_dtype(dataset.data[column_name], output_dtype)
-
         self.noise_function(
             dataset,
             configuration,

src/pseudopeople/interface.py

@@ -38,32 +38,49 @@ def _generate_dataset(
     filters: Sequence[DataFilter],
     verbose: bool,
     engine_name: Literal["pandas"],
+    concat: Literal[True],
 ) -> pd.DataFrame:
     ...
 
 
 @overload
 def _generate_dataset(
     dataset_schema: DatasetSchema,
-    source: Path | str | None,
+    source: list[pd.DataFrame],
+    seed: int,
+    config: Path | str | dict[str, Any] | None,
+    filters: Sequence[DataFilter],
+    verbose: bool,
+    engine_name: Literal["pandas"],
+    concat_output: Literal[False],
+) -> list[pd.DataFrame]:
+    ...
+
+
+@overload
+def _generate_dataset(
+    dataset_schema: DatasetSchema,
+    source: Path | str | None | dd.DataFrame,
     seed: int,
     config: Path | str | dict[str, Any] | None,
     filters: Sequence[DataFilter],
     verbose: bool,
     engine_name: Literal["dask"],
+    concat_output: bool,
 ) -> dd.DataFrame:
     ...
 
 
 def _generate_dataset(
     dataset_schema: DatasetSchema,
-    source: Path | str | None,
+    source: Path | str | None | list[pd.DataFrame] | dd.DataFrame,
     seed: int,
     config: Path | str | dict[str, Any] | None,
     filters: Sequence[DataFilter],
     verbose: bool = False,
     engine_name: Literal["pandas", "dask"] = "pandas",
-) -> pd.DataFrame | dd.DataFrame:
+    concat_output = True,
+) -> pd.DataFrame | dd.DataFrame | list[pd.DataFrame]:
     """
     Helper for generating noised datasets.
 
@@ -85,18 +102,19 @@ def _generate_dataset(
     :return:
         Noised dataset data in a dataframe
     """
+    breakpoint()
     configure_logging_to_terminal(verbose)
     noise_configuration = get_configuration(config, dataset_schema, filters)
 
     if source is None:
         source = paths.SAMPLE_DATA_ROOT
-    else:
+    elif isinstance(source, Path) or isinstance(source, str):
         source = Path(source)
         validate_source_compatibility(source, dataset_schema)
 
     engine = get_engine_from_string(engine_name)
 
-    noised_dataset: pd.DataFrame | dd.DataFrame
+    noised_dataset: pd.DataFrame | dd.DataFrame | list[pd.DataFrame]
     if engine == PANDAS_ENGINE:
         # We process shards serially
         data_file_paths = get_dataset_filepaths(source, dataset_schema.name)
@@ -137,9 +155,11 @@ def _generate_dataset(
                 "Invalid value provided for 'state' or 'year'. No data found with "
                 f"the user provided 'state' or 'year' filters at {source / dataset_schema.name}."
             )
-        noised_dataset = pd.concat(noised_datasets_list, ignore_index=True)
-
-        noised_dataset = coerce_dtypes(noised_dataset, dataset_schema)
+        if concat_output:
+            noised_dataset = pd.concat(noised_datasets_list, ignore_index=True)
+            noised_dataset = coerce_dtypes(noised_dataset, dataset_schema)
+        else:
+            noised_dataset = [coerce_dtypes(dataset, dataset_schema) for dataset in noised_dataset]
     else:  # dask
         try:
             from distributed.client import default_client

src/pseudopeople/noise_functions.py

@@ -305,7 +305,6 @@ def copy_from_household_member(
     :param to_noise_index: pd.Index of rows to be noised
     :param column_name: String for column that will be noised
     """
-
     copy_values = dataset.data.loc[to_noise_index, COPY_HOUSEHOLD_MEMBER_COLS[column_name]]
     dataset.data.loc[to_noise_index, column_name] = ensure_dtype(
         pd.Series(copy_values, name=column_name, index=to_noise_index),

tests/integration/release/test_release.py

@@ -27,7 +27,7 @@
 from tests.constants import TOKENS_PER_STRING_MAPPER
 from tests.integration.conftest import SEED, IDX_COLS
 from tests.integration.release.conftest import DATASET_ARG_TO_FULL_NAME_MAPPER
-from tests.integration.release.utilities import run_omit_row_tests, run_do_not_respond_tests, run_guardian_duplication_tests
+from tests.integration.release.utilities import load_unnoised_data, run_omit_row_tests, run_do_not_respond_tests, run_guardian_duplication_tests
 from tests.unit.test_configuration import COLUMN_NOISE_TYPES
 from tests.utilities import (
     get_single_noise_type_config,
@@ -38,7 +38,7 @@
 ROW_TEST_FUNCTIONS = {'omit_row': run_omit_row_tests,
 'do_not_respond': run_do_not_respond_tests,
 'duplicate_with_guardian': run_guardian_duplication_tests}
-NEW_PROBABILITY = 0.03
+NEW_PROBABILITY = 0.2
 
 
 def get_high_noise_config(
@@ -49,7 +49,7 @@ def get_high_noise_config(
 
     for noise_type, probabilities in config_dict[dataset_name][Keys.ROW_NOISE].items():
         for probability_name, probability in probabilities.items():
-            config_dict[dataset_name][Keys.ROW_NOISE][noise_type][probability_name] = .03
+            config_dict[dataset_name][Keys.ROW_NOISE][noise_type][probability_name] = NEW_PROBABILITY
 
     for col, noise_types in config_dict[dataset_name][Keys.COLUMN_NOISE].items():
         for noise_type, probabilities in noise_types.items():
@@ -62,7 +62,7 @@ def get_high_noise_config(
                     # NOTE: this will fail default config validations
                     new_probability = {1: 1.0}
                 else:
-                    new_probability = .03
+                    new_probability = NEW_PROBABILITY
                 config_dict[dataset_name][Keys.COLUMN_NOISE][col][noise_type][
                     probability_name
                 ] = new_probability
@@ -98,6 +98,10 @@ def test_release_runs(
     if dataset_func != generate_social_security:
         unnoised_data_kwargs["state"] = state
     unnoised_data = dataset_func(**unnoised_data_kwargs)
+    full_dataset_name = DATASET_ARG_TO_FULL_NAME_MAPPER[dataset_name]
+    dataset_schema = DATASET_SCHEMAS.get_dataset_schema(full_dataset_name)
+    #new_unnoised_data = load_unnoised_data(dataset_schema, source, year, state, engine)
+    #breakpoint()
 
     if source is None:
         population = 'sample'
@@ -106,20 +110,17 @@ def test_release_runs(
     else:
         population = 'usa'
     timestr = time.strftime("%Y%m%d-%H%M%S")
-    filename = f"/ihme/homes/hjafari/ppl_runs_new/{timestr}_{dataset_name}_{population}.o"
+    filename = f"/ihme/homes/hjafari/ppl_runs_new/{timestr}_{dataset_name}_{population}_{NEW_PROBABILITY*100}_percent.o"
 
     # In our standard noising process, i.e. when noising a shard of data, we
     # 1) clean and reformat the data, 2) noise the data, and 3) do some post-processing.
     # We're replicating steps 1 and 2 in this test and skipping 3.
-    full_dataset_name = DATASET_ARG_TO_FULL_NAME_MAPPER[dataset_name]
-    dataset_schema = DATASET_SCHEMAS.get_dataset_schema(full_dataset_name)
     dataset = Dataset(dataset_schema, unnoised_data, SEED)
-    # don't unnecessarily keep in memory
+    # don't unnecessarily keep in memory now that we have the data in dataset
     del unnoised_data
     dataset._clean_input_data()
     # convert datetime columns to datetime types for _reformat_dates_for_noising
-    # because the post-processing that occured in generating the unnoised data
-    # in step 3 mentioned above converts these columns to object dtypes
+    # because we coerce these types into object types when loading unnoised data
     for col in [COLUMNS.dob.name, COLUMNS.ssa_event_date.name]:
         if col in dataset.data:
             dataset.data[col] = pd.to_datetime(
@@ -133,7 +134,9 @@ def test_release_runs(
     dataset._reformat_dates_for_noising()
 
     config = NoiseConfiguration(LayeredConfigTree(get_high_noise_config(full_dataset_name)))
-    
+
+
+
     for noise_type in NOISE_TYPES:
         original_data = dataset.data.copy()
         # if isinstance(noise_type, RowNoiseType):
@@ -145,12 +148,15 @@ def test_release_runs(
             for column in dataset.data.columns:
                 if config.has_noise_type(
                     dataset.dataset_schema.name, noise_type.name, column
-                ):
-                    noise_type(dataset, config, column)
-                    run_column_noising_test(original_data, dataset.data, config, full_dataset_name, noise_type.name, column, fuzzy_checker, filename)
+                ):# and noise_type.name == 'copy_from_household_member' and column == 'age':
+                    if column == COLUMNS.ssa_event_type.name:
+                        pass
+                    else:
+                        noise_type(dataset, config, column)
+                        run_column_noising_test(original_data, dataset.data, config, full_dataset_name, noise_type.name, column, fuzzy_checker, filename)
         with check:
-            # TODO: possible to replace missingness with smaller data structure?
             try:
+            # TODO: possible to replace missingness with smaller data structure?
                 assert dataset.missingness.equals(dataset.is_missing(dataset.data))
             except:
                 breakpoint()
@@ -164,6 +170,72 @@ def test_release_runs(
     # run_final_tests(unnoised_data, dataset.data)
 
 
+def test_release_runs(
+    dataset_params: tuple[
+    str,
+    Callable[..., pd.DataFrame],
+    str | None,
+    int | None,
+    str | None,
+    Literal["pandas", "dask"],
+    ],
+    fuzzy_checker: FuzzyChecker,
+    mocker: MockerFixture,
+) -> None:
+    # keep all columns when generating unnoised data because some of them are used in testing
+    mocker.patch(
+        "pseudopeople.dataset.Dataset.drop_non_schema_columns", side_effect=lambda df, _: df
+    )
+
+    # create unnoised dataset
+    dataset_name, dataset_func, source, year, state, engine = dataset_params
+    unnoised_data_kwargs = {
+        "source": source,
+        "config": NO_NOISE,
+        "year": year,
+        "engine": engine,
+        "concat_output": False,
+    }
+    if dataset_func != generate_social_security:
+        unnoised_data_kwargs["state"] = state
+    unnoised_data = dataset_func(**unnoised_data_kwargs)
+    #full_dataset_name = DATASET_ARG_TO_FULL_NAME_MAPPER[dataset_name]
+    #dataset_schema = DATASET_SCHEMAS.get_dataset_schema(full_dataset_name)
+    #new_unnoised_data = load_unnoised_data(dataset_schema, source, year, state, engine)
+
+    # In our standard noising process, i.e. when noising a shard of data, we
+    # 1) clean and reformat the data, 2) noise the data, and 3) do some post-processing.
+    # We're replicating steps 1 and 2 in this test and skipping 3.
+    full_dataset_name = DATASET_ARG_TO_FULL_NAME_MAPPER[dataset_name]
+    dataset_schema = DATASET_SCHEMAS.get_dataset_schema(full_dataset_name)
+    dataset = Dataset(dataset_schema, unnoised_data, SEED)
+    # don't unnecessarily keep in memory now that we have the data in dataset
+    del unnoised_data
+    dataset._clean_input_data()
+    # convert datetime columns to datetime types for _reformat_dates_for_noising
+    # because we coerce these types into object types when loading unnoised data
+    for col in [COLUMNS.dob.name, COLUMNS.ssa_event_date.name]:
+        if col in dataset.data:
+            dataset.data[col] = pd.to_datetime(
+                dataset.data[col], format=dataset_schema.date_format
+            )
+            copy_col = "copy_" + col
+            if copy_col in dataset.data:
+                dataset.data[copy_col] = pd.to_datetime(
+                    dataset.data[copy_col], format=dataset_schema.date_format
+                )
+    # TODO: mock this to do nothing
+    dataset._reformat_dates_for_noising()
+
+    # TODO: generate dictionary of noise configs with key as noise type and value is config
+
+    # for noise_type, config in config_dict:
+    #     # TODO: case pandas vs dask 
+    #     pre_noised = [x.copy() for x in unnoised_data]
+    #     noised_data = dataset_func(source=prenoised_data, config)
+    #     check_noise(pre_noised, noised_data)
+
+
 def _get_common_datasets(
     unnoised_dataset: Dataset, noised_dataset: pd.DataFrame
 ) -> tuple[pd.DataFrame, pd.DataFrame, pd.Index[int]]:
@@ -218,12 +290,9 @@ def run_column_noising_test(
 ) -> None:
     dataset_schema = DATASET_SCHEMAS.get_dataset_schema(dataset_name)
     original_dataset = Dataset(dataset_schema, original_data, SEED)
-    try:
-        check_noised, check_original, shared_idx = _get_common_datasets(
-            original_dataset, noised_data
-        )
-    except:
-        breakpoint()
+    check_noised, check_original, shared_idx = _get_common_datasets(
+        original_dataset, noised_data
+    )
     # TODO: remove population param which was just used in testing the tests
     check_column_noising(dataset_name, config, fuzzy_checker, check_noised, check_original, shared_idx, noise_type, column, filename)
 
@@ -305,10 +374,7 @@ def check_column_noising(
 
         # This is accumulating not_noised over all noise types
         expected_noise = avg_probability_any_token_noised * expected_noise
-    #         if no_differences:
-    #             with open(filename, "a") as f:
-    #                 info = f"no differences for NOISE_TYPE_{noise_type}_COL_{col.name} and expected noise level of {expected_noise} or {expected_noise * len(check_original.loc[to_compare_idx, col.name])} simulants\n"
-    #                 f.write(info)
+    open(filename, 'a').close()
     try:
         fuzzy_checker.fuzzy_assert_proportion(
             name=noise_type,
@@ -332,7 +398,7 @@ def check_column_noising(
                 f.write(info)
         else:
             with open(filename, "a") as f:
-                info = f"NOISE_TYPE_{noise_type}_COL_{col.name} issue fuzzy checking: expected {expected_noise} but got {noise_level / len(check_original.loc[to_compare_idx, col.name])}\n"
+                info = f"NOISE_TYPE_{noise_type}_COL_{col.name} issue fuzzy checking: expected {expected_noise} but got {noise_level / len(check_original.loc[to_compare_idx, col.name])} from {noise_level} / {len(check_original.loc[to_compare_idx, col.name])}\n"
                 f.write(info)