added storage functionality for raw windows

src/aces/__main__.py

@@ -173,7 +173,8 @@ def main(cfg: DictConfig):
                 logger.warning("Output dataframe is empty; adding an empty patient ID column.")
                 result = result.with_columns(pl.lit(None, dtype=pl.Int64).alias("patient_id"))
                 result = result.head(0)
-
+        if cfg.window_stats_filepath:
+            result.write_parquet(cfg.window_stats_filepath)
         result = get_and_validate_label_schema(result)
         pq.write_table(result, cfg.output_filepath)
     else:

src/aces/configs/aces.yaml

@@ -12,6 +12,8 @@ config_path: ${cohort_dir}/${cohort_name}.yaml
 # Path to store the output file. The `${data._prefix}` addition allows us to add shard specific prefixes in a
 # sharded data mode.
 output_filepath: ${cohort_dir}/${cohort_name}${data._prefix}.parquet
+# Optional path to store the output file with the raw window data.
+window_stats_filepath: null
 
 log_dir: ${cohort_dir}/${cohort_name}/.logs
 

tests/test_meds.py

@@ -368,3 +368,91 @@ def test_meds():
         want_outputs_by_task={TASK_NAME: WANT_SHARDS},
         data_standard="meds",
     )
+
+
+def test_meds_window_storage():
+  import tempfile
+  from pathlib import Path
+  from tests.utils import write_input_files, write_task_configs, run_command, assert_df_equal
+
+  input_files = MEDS_SHARDS
+  task_configs = {TASK_NAME: TASK_CFG}
+  want_outputs_by_task = {TASK_NAME: WANT_SHARDS}
+  data_standard = "meds"
+
+  with tempfile.TemporaryDirectory() as root_dir:
+    root_dir = Path(root_dir)
+    data_dir = root_dir / "sample_data" / "data"
+    cohort_dir = root_dir / "sample_cohort"
+
+    wrote_files = write_input_files(data_dir, input_files)
+    if len(wrote_files) == 0:
+        raise ValueError("No input files were written.")
+    elif len(wrote_files) > 1:
+        sharded = True
+        command = "aces-cli --multirun"
+    else:
+        sharded = False
+        command = "aces-cli"
+
+    wrote_configs = write_task_configs(cohort_dir, task_configs)
+    if len(wrote_configs) == 0:
+        raise ValueError("No task configs were written.")
+
+    for task in task_configs:
+        if sharded:
+            want_outputs = {
+                cohort_dir / task / f"{n}.parquet": df for n, df in want_outputs_by_task[task].items()
+            }
+        else:
+            want_outputs = {cohort_dir / f"{task}.parquet": want_outputs_by_task[task]}
+        window_stats_filepath = cohort_dir / f"{task}_window_stats.parquet"
+        want_window_fp = window_stats_filepath
+
+        extraction_config_kwargs = {
+            "cohort_dir": str(cohort_dir.resolve()),
+            "cohort_name": task,
+            "hydra.verbose": True,
+            "data.standard": data_standard,
+            "window_stats_filepath": str(window_stats_filepath.resolve()),
+        }
+
+        if len(wrote_files) > 1:
+            extraction_config_kwargs["data"] = "sharded"
+            extraction_config_kwargs["data.root"] = str(data_dir.resolve())
+            extraction_config_kwargs['"data.shard'] = f'$(expand_shards {str(data_dir.resolve())})"'
+        else:
+            extraction_config_kwargs["data.path"] = str(list(wrote_files.values())[0].resolve())
+
+        stderr, stdout = run_command(command, extraction_config_kwargs, f"CLI should run for {task}")
+
+        try:
+            if sharded:
+                out_dir = cohort_dir / task
+                all_out_fps = list(out_dir.glob("**/*.parquet"))
+                all_out_fps_str = ", ".join(str(x.relative_to(out_dir)) for x in all_out_fps)
+                if len(all_out_fps) == 0 and len(want_outputs) > 0:
+                    all_directory_contents = ", ".join(
+                        str(x.relative_to(cohort_dir)) for x in cohort_dir.glob("**/*")
+                    )
+
+                    raise AssertionError(
+                        f"No output files found for task '{task}'. Found files: {all_directory_contents}"
+                    )
+
+                assert len(all_out_fps) == len(
+                    want_outputs
+                ), f"Expected {len(want_outputs)} outputs, got {len(all_out_fps)}: {all_out_fps_str}"
+
+            for want_fp, want_df in want_outputs.items():
+                out_shard = want_fp.relative_to(cohort_dir)
+                assert want_fp.is_file(), f"Expected {out_shard} to exist."
+
+                got_df = pl.read_parquet(want_fp)
+                assert_df_equal(
+                    want_df, got_df, f"Data mismatch for shard '{out_shard}':\n{want_df}\n{got_df}"
+                )
+            assert want_window_fp.is_file(), f"Expected window stats file {want_window_fp} to exist."
+        except AssertionError as e:
+            logger.error(f"{stderr}\n{stdout}")
+            raise AssertionError(f"Error running task '{task}': {e}") from e