added storage functionality for raw windows

src/aces/__main__.py

@@ -173,7 +173,9 @@ 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_dir:
+            Path(cfg.window_stats_filename).parent.mkdir(exist_ok=True, parents=True)
+            result.write_parquet(cfg.window_stats_filename)
         result = get_and_validate_label_schema(result)
         pq.write_table(result, cfg.output_filepath)
     else:

src/aces/configs/aces.yaml

@@ -12,6 +12,9 @@ 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_dir: null
+window_stats_filename: ${window_stats_dir}/${cohort_name}${data._prefix}.parquet
 
 log_dir: ${cohort_dir}/${cohort_name}/.logs
 

tests/test_meds.py

@@ -5,15 +5,23 @@
 
 root = rootutils.setup_root(__file__, dotenv=True, pythonpath=True, cwd=True)
 
+import tempfile
 from io import StringIO
+from pathlib import Path
 
 import polars as pl
 import pyarrow as pa
 from loguru import logger
 from meds import label_schema
 from yaml import load as load_yaml
 
-from .utils import cli_test
+from .utils import (
+    assert_df_equal,
+    cli_test,
+    run_command,
+    write_input_files,
+    write_task_configs,
+)
 
 try:
     from yaml import CLoader as Loader
@@ -368,3 +376,91 @@ def test_meds():
         want_outputs_by_task={TASK_NAME: WANT_SHARDS},
         data_standard="meds",
     )
+
+
+def test_meds_window_storage():
+    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)
+        assert len(wrote_files) > 1, "No input files were written."
+        sharded = True
+        command = "aces-cli --multirun"
+
+        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:
+            want_outputs = {
+                cohort_dir / task / f"{n}.parquet": df for n, df in want_outputs_by_task[task].items()
+            }
+            window_dir = Path(cohort_dir / "window_stats")
+            want_window_output_files = [
+                window_dir / task / f"{n}.parquet" for n in want_outputs_by_task[task]
+            ]
+
+            extraction_config_kwargs = {
+                "cohort_dir": str(cohort_dir.resolve()),
+                "cohort_name": task,
+                "hydra.verbose": True,
+                "data.standard": data_standard,
+                "window_stats_dir": str(window_dir.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 window_dir.exists(), f"Expected window stats directory {window_dir} to exist."
+                out_fps = list(window_dir.glob("**/*.parquet"))
+                assert len(out_fps) == len(
+                    want_window_output_files
+                ), f"Expected {len(want_window_output_files)} window output files, got {len(out_fps)}"
+
+                for want_fp in want_window_output_files:
+                    out_shard = want_fp.relative_to(window_dir)
+                    assert want_fp.is_file(), f"Expected {out_shard} to exist."
+                    got_df = pl.read_parquet(want_fp)
+                    assert "patient_id" in got_df.columns, f"Expected 'patient_id' column in {out_shard}."
+            except AssertionError as e:
+                logger.error(f"{stderr}\n{stdout}")
+                raise AssertionError(f"Error running task '{task}': {e}") from e