sumeds 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
sumeds-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2025 Robin van de Water
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
sumeds-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,100 @@
1
+ Metadata-Version: 2.4
2
+ Name: sumeds
3
+ Version: 0.1.0
4
+ Summary: Lazy, privacy-aware code occurrence summaries for MEDS datasets
5
+ Keywords: MEDS,healthcare,Polars,privacy
6
+ Author: Robin van de Water
7
+ Author-email: Robin van de Water <r.p.vandewater@gmail.com>
8
+ License-Expression: MIT
9
+ License-File: LICENSE
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Programming Language :: Python :: 3
12
+ Classifier: Programming Language :: Python :: 3 :: Only
13
+ Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
14
+ Requires-Dist: meds>=0.4.1,<0.5
15
+ Requires-Dist: polars>=1.30
16
+ Requires-Dist: pyarrow>=15
17
+ Requires-Dist: pyyaml>=6
18
+ Requires-Python: >=3.10, <3.15
19
+ Description-Content-Type: text/markdown
20
+
21
+ # suMEDS
22
+
23
+ `suMEDS` creates a compact code catalog for defining tasks on very large
24
+ [Medical Event Data Standard (MEDS)](https://github.com/Medical-Event-Data-Standard/meds)
25
+ datasets. It joins canonical descriptions and parent codes to event and unique-subject
26
+ counts, then masks or removes rare codes before writing Parquet, CSV, or JSON.
27
+
28
+ Polars scans, aggregations, joins, and writes stay lazy. Patient-level rows are never
29
+ collected into Python memory.
30
+
31
+ ## Install and run
32
+
33
+ ```bash
34
+ uv sync
35
+ uv run suMEDS MIMICIV_DEMO/MEDS_cohort \
36
+ --output code-summary.parquet \
37
+ --min-subjects 20
38
+ ```
39
+
40
+ Or use YAML:
41
+
42
+ ```yaml
43
+ summary:
44
+ per_split: false
45
+ split_columns: false
46
+ privacy:
47
+ min_subjects: 20
48
+ min_split_subjects: 1
49
+ rare_code_action: bucket
50
+ rare_code_label: __RARE__
51
+ round_counts_to: 5
52
+ ```
53
+
54
+ ```bash
55
+ uv run suMEDS /path/to/MEDS -o summary.parquet -c examples/summary.yaml
56
+ # The suffix selects CSV or JSON instead:
57
+ uv run suMEDS /path/to/MEDS -o summary.json -c examples/summary.yaml
58
+ ```
59
+
60
+ CLI flags override YAML. The defaults bucket codes seen in fewer than 20 unique
61
+ subjects and retain exact counts.
62
+
63
+ ## Python API
64
+
65
+ ```python
66
+ from sumeds import SummaryConfig, summarize
67
+
68
+ path = summarize(
69
+ "/path/to/MEDS",
70
+ "code-summary.parquet",
71
+ SummaryConfig(min_subjects=20, rare_code_action="bucket"),
72
+ )
73
+ ```
74
+
75
+ The output format is inferred from `.parquet`, `.csv`, or `.json` (with
76
+ `.jsonl`/`.ndjson` also supported). It preserves all code-metadata extension columns and adds
77
+ `event_count`, `subject_count`, and `is_masked`. `per_split=True` emits one row
78
+ per split. `split_columns=True` keeps the total row and adds columns such as
79
+ `event_count_train` and `subject_count_train` from `subject_splits.parquet`.
80
+ Set `min_split_subjects` above 1 to suppress rare split-level cells.
81
+
82
+ ## Documentation
83
+
84
+ ```bash
85
+ uv run mkdocs serve
86
+ uv run mkdocs build --strict
87
+ ```
88
+
89
+ See the [documentation site](docs/index.md) for the data flow, complete CLI and
90
+ YAML reference, output schema, Python API, and privacy limitations.
91
+
92
+ ## Development
93
+
94
+ ```bash
95
+ uv run pre-commit run --all-files
96
+ uv run pytest
97
+ uv run mkdocs build --strict
98
+ ```
99
+
100
+ The included integration test uses `MIMICIV_DEMO/MEDS_cohort` when present.
sumeds-0.1.0/README.md ADDED
@@ -0,0 +1,80 @@
1
+ # suMEDS
2
+
3
+ `suMEDS` creates a compact code catalog for defining tasks on very large
4
+ [Medical Event Data Standard (MEDS)](https://github.com/Medical-Event-Data-Standard/meds)
5
+ datasets. It joins canonical descriptions and parent codes to event and unique-subject
6
+ counts, then masks or removes rare codes before writing Parquet, CSV, or JSON.
7
+
8
+ Polars scans, aggregations, joins, and writes stay lazy. Patient-level rows are never
9
+ collected into Python memory.
10
+
11
+ ## Install and run
12
+
13
+ ```bash
14
+ uv sync
15
+ uv run suMEDS MIMICIV_DEMO/MEDS_cohort \
16
+ --output code-summary.parquet \
17
+ --min-subjects 20
18
+ ```
19
+
20
+ Or use YAML:
21
+
22
+ ```yaml
23
+ summary:
24
+ per_split: false
25
+ split_columns: false
26
+ privacy:
27
+ min_subjects: 20
28
+ min_split_subjects: 1
29
+ rare_code_action: bucket
30
+ rare_code_label: __RARE__
31
+ round_counts_to: 5
32
+ ```
33
+
34
+ ```bash
35
+ uv run suMEDS /path/to/MEDS -o summary.parquet -c examples/summary.yaml
36
+ # The suffix selects CSV or JSON instead:
37
+ uv run suMEDS /path/to/MEDS -o summary.json -c examples/summary.yaml
38
+ ```
39
+
40
+ CLI flags override YAML. The defaults bucket codes seen in fewer than 20 unique
41
+ subjects and retain exact counts.
42
+
43
+ ## Python API
44
+
45
+ ```python
46
+ from sumeds import SummaryConfig, summarize
47
+
48
+ path = summarize(
49
+ "/path/to/MEDS",
50
+ "code-summary.parquet",
51
+ SummaryConfig(min_subjects=20, rare_code_action="bucket"),
52
+ )
53
+ ```
54
+
55
+ The output format is inferred from `.parquet`, `.csv`, or `.json` (with
56
+ `.jsonl`/`.ndjson` also supported). It preserves all code-metadata extension columns and adds
57
+ `event_count`, `subject_count`, and `is_masked`. `per_split=True` emits one row
58
+ per split. `split_columns=True` keeps the total row and adds columns such as
59
+ `event_count_train` and `subject_count_train` from `subject_splits.parquet`.
60
+ Set `min_split_subjects` above 1 to suppress rare split-level cells.
61
+
62
+ ## Documentation
63
+
64
+ ```bash
65
+ uv run mkdocs serve
66
+ uv run mkdocs build --strict
67
+ ```
68
+
69
+ See the [documentation site](docs/index.md) for the data flow, complete CLI and
70
+ YAML reference, output schema, Python API, and privacy limitations.
71
+
72
+ ## Development
73
+
74
+ ```bash
75
+ uv run pre-commit run --all-files
76
+ uv run pytest
77
+ uv run mkdocs build --strict
78
+ ```
79
+
80
+ The included integration test uses `MIMICIV_DEMO/MEDS_cohort` when present.
@@ -0,0 +1,39 @@
1
+ [project]
2
+ name = "sumeds"
3
+ version = "0.1.0"
4
+ description = "Lazy, privacy-aware code occurrence summaries for MEDS datasets"
5
+ readme = "README.md"
6
+ license = "MIT"
7
+ license-files = ["LICENSE"]
8
+ keywords = ["MEDS", "healthcare", "Polars", "privacy"]
9
+ classifiers = [
10
+ "Development Status :: 3 - Alpha",
11
+ "Programming Language :: Python :: 3",
12
+ "Programming Language :: Python :: 3 :: Only",
13
+ "Topic :: Scientific/Engineering :: Medical Science Apps.",
14
+ ]
15
+ authors = [
16
+ { name = "Robin van de Water", email = "r.p.vandewater@gmail.com" }
17
+ ]
18
+ requires-python = ">=3.10, <3.15"
19
+ dependencies = [
20
+ "meds>=0.4.1,<0.5",
21
+ "polars>=1.30",
22
+ "pyarrow>=15",
23
+ "pyyaml>=6",
24
+ ]
25
+
26
+ [project.scripts]
27
+ suMEDS = "sumeds.cli:main"
28
+
29
+ [build-system]
30
+ requires = ["uv_build>=0.11.15,<0.12.0"]
31
+ build-backend = "uv_build"
32
+
33
+ [dependency-groups]
34
+ dev = [
35
+ "mkdocs>=1.6",
36
+ "pre-commit>=4.0",
37
+ "pytest>=8",
38
+ "ruff>=0.15.22",
39
+ ]
@@ -0,0 +1,26 @@
1
+ """Compact, lazy summaries for Medical Event Data Standard datasets."""
2
+
3
+ from .config import SummaryConfig
4
+ from .scan import (
5
+ code_modifier_columns,
6
+ dataset_root,
7
+ event_files,
8
+ read_dataset_metadata,
9
+ scan_code_metadata,
10
+ scan_events,
11
+ scan_subject_splits,
12
+ )
13
+ from .summary import code_occurrences, summarize
14
+
15
+ __all__ = [
16
+ "SummaryConfig",
17
+ "code_modifier_columns",
18
+ "code_occurrences",
19
+ "dataset_root",
20
+ "event_files",
21
+ "read_dataset_metadata",
22
+ "scan_code_metadata",
23
+ "scan_events",
24
+ "scan_subject_splits",
25
+ "summarize",
26
+ ]
@@ -0,0 +1,87 @@
1
+ """Command-line interface for :mod:`sumeds`."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import argparse
6
+ from importlib.metadata import version
7
+ from pathlib import Path
8
+ from typing import Sequence
9
+
10
+ from .config import SummaryConfig
11
+ from .summary import summarize
12
+
13
+
14
+ def build_parser() -> argparse.ArgumentParser:
15
+ """Build the public CLI parser."""
16
+
17
+ parser = argparse.ArgumentParser(
18
+ prog="suMEDS",
19
+ description="Create a lazy, privacy-aware MEDS code occurrence catalog.",
20
+ )
21
+ parser.add_argument("dataset_root", type=Path, help="MEDS dataset root")
22
+ parser.add_argument(
23
+ "-o",
24
+ "--output",
25
+ type=Path,
26
+ required=True,
27
+ help="output .parquet, .csv, or .json file",
28
+ )
29
+ parser.add_argument("-c", "--config", type=Path, help="YAML configuration file")
30
+ parser.add_argument(
31
+ "--per-split", action=argparse.BooleanOptionalAction, default=None
32
+ )
33
+ parser.add_argument(
34
+ "--split-columns",
35
+ action=argparse.BooleanOptionalAction,
36
+ default=None,
37
+ help="add per-split count columns alongside totals",
38
+ )
39
+ parser.add_argument(
40
+ "--min-subjects", type=int, help="minimum unique subjects per released code"
41
+ )
42
+ parser.add_argument(
43
+ "--min-split-subjects",
44
+ type=int,
45
+ help="minimum unique subjects per visible wide split cell",
46
+ )
47
+ parser.add_argument("--rare-code-action", choices=("bucket", "drop"))
48
+ parser.add_argument("--rare-code-label", help="label used by bucket mode")
49
+ parser.add_argument(
50
+ "--round-counts-to", type=int, help="round released counts to this multiple"
51
+ )
52
+ parser.add_argument("--version", action="version", version=version("suMEDS"))
53
+ return parser
54
+
55
+
56
+ def main(argv: Sequence[str] | None = None) -> int:
57
+ """Run the CLI and return a process exit code."""
58
+
59
+ parser = build_parser()
60
+ args = parser.parse_args(argv)
61
+ try:
62
+ config = (
63
+ SummaryConfig.from_yaml(args.config) if args.config else SummaryConfig()
64
+ )
65
+ config = config.with_overrides(
66
+ per_split=args.per_split,
67
+ split_columns=args.split_columns,
68
+ min_subjects=args.min_subjects,
69
+ min_split_subjects=args.min_split_subjects,
70
+ rare_code_action=args.rare_code_action,
71
+ rare_code_label=args.rare_code_label,
72
+ round_counts_to=args.round_counts_to,
73
+ )
74
+ except (OSError, ValueError) as exc:
75
+ parser.error(str(exc))
76
+
77
+ try:
78
+ output = summarize(args.dataset_root, args.output, config)
79
+ except Exception as exc: # CLI boundary: library callers retain typed exceptions.
80
+ parser.exit(1, f"suMEDS: error: {exc}\n")
81
+
82
+ print(f"Wrote {output}")
83
+ return 0
84
+
85
+
86
+ if __name__ == "__main__":
87
+ raise SystemExit(main())
@@ -0,0 +1,90 @@
1
+ """Configuration for MEDS summaries."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from dataclasses import dataclass, fields, replace
6
+ from pathlib import Path
7
+ from typing import Any
8
+
9
+ import yaml
10
+
11
+
12
+ @dataclass(frozen=True)
13
+ class SummaryConfig:
14
+ """Options controlling grouping and privacy-aware release.
15
+
16
+ ``min_subjects`` is evaluated on unique subjects, not event rows. A value
17
+ of one disables useful rarity protection and should only be used on public
18
+ or synthetic data.
19
+ """
20
+
21
+ per_split: bool = False
22
+ split_columns: bool = False
23
+ min_subjects: int = 20
24
+ min_split_subjects: int = 1
25
+ rare_code_action: str = "bucket"
26
+ rare_code_label: str = "__RARE__"
27
+ round_counts_to: int | None = None
28
+
29
+ def __post_init__(self) -> None:
30
+ if not isinstance(self.per_split, bool) or not isinstance(
31
+ self.split_columns, bool
32
+ ):
33
+ raise ValueError("per_split and split_columns must be true or false")
34
+ if self.per_split and self.split_columns:
35
+ raise ValueError("per_split and split_columns are mutually exclusive")
36
+ for name in ("min_subjects", "min_split_subjects"):
37
+ value = getattr(self, name)
38
+ if not isinstance(value, int) or isinstance(value, bool) or value < 1:
39
+ raise ValueError(f"{name} must be a positive integer")
40
+ if self.rare_code_action not in {"bucket", "drop"}:
41
+ raise ValueError("rare_code_action must be 'bucket' or 'drop'")
42
+ if not self.rare_code_label:
43
+ raise ValueError("rare_code_label must not be empty")
44
+ if self.round_counts_to is not None and (
45
+ not isinstance(self.round_counts_to, int)
46
+ or isinstance(self.round_counts_to, bool)
47
+ or self.round_counts_to < 1
48
+ ):
49
+ raise ValueError("round_counts_to must be null or a positive integer")
50
+
51
+ @classmethod
52
+ def from_yaml(cls, path: str | Path) -> SummaryConfig:
53
+ """Load a strict ``summary``/``privacy`` YAML configuration file."""
54
+
55
+ source = Path(path)
56
+ raw = yaml.safe_load(source.read_text()) or {}
57
+ if not isinstance(raw, dict):
58
+ raise ValueError(f"YAML root must be a mapping: {source}")
59
+ unknown_sections = set(raw) - {"summary", "privacy"}
60
+ if unknown_sections:
61
+ raise ValueError(f"Unknown YAML sections: {sorted(unknown_sections)}")
62
+
63
+ summary = _mapping(raw.get("summary"), "summary")
64
+ privacy = _mapping(raw.get("privacy"), "privacy")
65
+ allowed = {field.name for field in fields(cls)}
66
+ summary_options = {"per_split", "split_columns"}
67
+ unknown = (set(summary) - summary_options) | (
68
+ set(privacy) - (allowed - summary_options)
69
+ )
70
+ if unknown:
71
+ raise ValueError(
72
+ f"Unknown configuration options or wrong sections: {sorted(unknown)}"
73
+ )
74
+ values: dict[str, Any] = {**summary, **privacy}
75
+ return cls(**values)
76
+
77
+ def with_overrides(self, **values: Any) -> SummaryConfig:
78
+ """Return a copy with non-``None`` CLI overrides applied."""
79
+
80
+ return replace(
81
+ self, **{key: value for key, value in values.items() if value is not None}
82
+ )
83
+
84
+
85
+ def _mapping(value: Any, name: str) -> dict[str, Any]:
86
+ if value is None:
87
+ return {}
88
+ if not isinstance(value, dict):
89
+ raise ValueError(f"YAML section '{name}' must be a mapping")
90
+ return value
File without changes
@@ -0,0 +1,110 @@
1
+ """Validated, projection-friendly MEDS scans."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import json
6
+ from pathlib import Path
7
+ from typing import Any, Sequence
8
+
9
+ import meds
10
+ import polars as pl
11
+ import pyarrow as pa
12
+ import pyarrow.parquet as pq
13
+
14
+
15
+ def dataset_root(path: str | Path) -> Path:
16
+ """Return a validated MEDS root containing ``data`` and ``metadata``."""
17
+
18
+ root = Path(path).expanduser().resolve()
19
+ for child in (meds.data_subdirectory, Path(meds.dataset_metadata_filepath).parent):
20
+ if not (root / child).is_dir():
21
+ raise FileNotFoundError(f"MEDS directory not found: {root / child}")
22
+ return root
23
+
24
+
25
+ def read_dataset_metadata(root: str | Path) -> dict[str, Any]:
26
+ """Read and validate MEDS ``metadata/dataset.json``."""
27
+
28
+ path = dataset_root(root) / meds.dataset_metadata_filepath
29
+ if not path.is_file():
30
+ raise FileNotFoundError(f"MEDS dataset metadata not found: {path}")
31
+ value = json.loads(path.read_text())
32
+ meds.DatasetMetadataSchema.validate(value)
33
+ return value
34
+
35
+
36
+ def code_modifier_columns(metadata: dict[str, Any]) -> list[str]:
37
+ """Return declared code modifiers, rejecting duplicates and invalid names."""
38
+
39
+ modifiers = metadata.get("code_modifier_columns", [])
40
+ if not isinstance(modifiers, list) or not all(
41
+ isinstance(item, str) and item for item in modifiers
42
+ ):
43
+ raise ValueError("code_modifier_columns must be a list of non-empty strings")
44
+ if len(modifiers) != len(set(modifiers)):
45
+ raise ValueError("code_modifier_columns contains duplicates")
46
+ if set(modifiers) & {"subject_id", "time", "code"}:
47
+ raise ValueError("code_modifier_columns contains a core MEDS column")
48
+ return modifiers
49
+
50
+
51
+ def event_files(root: str | Path) -> list[Path]:
52
+ """List event shards beneath the standard MEDS data directory."""
53
+
54
+ data = dataset_root(root) / meds.data_subdirectory
55
+ files = sorted(data.rglob("*.parquet"))
56
+ if not files:
57
+ raise FileNotFoundError(f"No MEDS Parquet shards found below {data}")
58
+ return files
59
+
60
+
61
+ def scan_events(root: str | Path, modifiers: Sequence[str] = ()) -> pl.LazyFrame:
62
+ """Validate shard footers and lazily project fields needed for counting."""
63
+
64
+ files = event_files(root)
65
+ for path in files:
66
+ schema = pq.read_schema(path)
67
+ meds.DataSchema.validate(schema)
68
+ for column in modifiers:
69
+ if column not in schema.names:
70
+ raise ValueError(
71
+ f"Declared code modifier '{column}' is missing from {path}"
72
+ )
73
+ if not (
74
+ pa.types.is_string(schema.field(column).type)
75
+ or pa.types.is_large_string(schema.field(column).type)
76
+ ):
77
+ raise ValueError(f"Code modifier '{column}' must be a string in {path}")
78
+ columns = ["subject_id", "code", *modifiers]
79
+ return pl.scan_parquet([str(path) for path in files]).select(columns)
80
+
81
+
82
+ def scan_code_metadata(root: str | Path) -> pl.LazyFrame:
83
+ """Validate and lazily scan the canonical code metadata table.
84
+
85
+ Optional standard columns are added as typed nulls so the summary output
86
+ always includes descriptions and parent-code relationships.
87
+ """
88
+
89
+ path = dataset_root(root) / meds.code_metadata_filepath
90
+ if not path.is_file():
91
+ raise FileNotFoundError(f"MEDS code metadata not found: {path}")
92
+ meds.CodeMetadataSchema.validate(pq.read_schema(path))
93
+ frame = pl.scan_parquet(path)
94
+ names = set(frame.collect_schema().names())
95
+ missing = []
96
+ if "description" not in names:
97
+ missing.append(pl.lit(None, dtype=pl.String).alias("description"))
98
+ if "parent_codes" not in names:
99
+ missing.append(pl.lit(None, dtype=pl.List(pl.String)).alias("parent_codes"))
100
+ return frame.with_columns(missing) if missing else frame
101
+
102
+
103
+ def scan_subject_splits(root: str | Path) -> pl.LazyFrame:
104
+ """Validate and lazily scan canonical subject-to-split assignments."""
105
+
106
+ path = dataset_root(root) / meds.subject_splits_filepath
107
+ if not path.is_file():
108
+ raise FileNotFoundError(f"MEDS subject splits not found: {path}")
109
+ meds.SubjectSplitSchema.validate(pq.read_schema(path))
110
+ return pl.scan_parquet(path).select("subject_id", "split")
@@ -0,0 +1,352 @@
1
+ """Lazy code occurrence summaries with rare-code masking."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import re
6
+ from pathlib import Path
7
+ from typing import Sequence
8
+ from uuid import uuid4
9
+
10
+ import polars as pl
11
+
12
+ from .config import SummaryConfig
13
+ from .scan import (
14
+ code_modifier_columns,
15
+ dataset_root,
16
+ read_dataset_metadata,
17
+ scan_code_metadata,
18
+ scan_events,
19
+ scan_subject_splits,
20
+ )
21
+
22
+ _COUNT_COLUMNS = ("event_count", "subject_count")
23
+ _RESERVED_COLUMNS = {*_COUNT_COLUMNS, "is_masked", "split"}
24
+ _OUTPUT_FORMATS = {
25
+ ".parquet": "parquet",
26
+ ".csv": "csv",
27
+ ".json": "json",
28
+ ".jsonl": "ndjson",
29
+ ".ndjson": "ndjson",
30
+ }
31
+
32
+
33
+ def code_occurrences(
34
+ events: pl.LazyFrame,
35
+ keys: Sequence[str],
36
+ splits: Sequence[str] = (),
37
+ ) -> pl.LazyFrame:
38
+ """Count total and optional named-split events and unique subjects lazily."""
39
+
40
+ return _code_occurrences(events, keys, _split_columns(splits))
41
+
42
+
43
+ def _code_occurrences(
44
+ events: pl.LazyFrame,
45
+ keys: Sequence[str],
46
+ split_columns: Sequence[tuple[str, str, str]],
47
+ ) -> pl.LazyFrame:
48
+ expressions = _count_expressions(split_columns)
49
+ if split_columns:
50
+ expressions.append(
51
+ pl.col("split")
52
+ .is_null()
53
+ .sum()
54
+ .cast(pl.UInt64)
55
+ .alias("_missing_split_count")
56
+ )
57
+ return events.group_by(list(keys)).agg(expressions)
58
+
59
+
60
+ def summarize(
61
+ root: str | Path,
62
+ output: str | Path,
63
+ config: SummaryConfig | None = None,
64
+ ) -> Path:
65
+ """Write a privacy-aware MEDS code catalog and return its path.
66
+
67
+ Event rows remain in Polars lazy plans and are streamed to bounded-memory
68
+ outputs. Bucket mode performs a second projected event scan so its subject
69
+ count is the exact union of subjects across rare codes.
70
+ """
71
+
72
+ config = config or SummaryConfig()
73
+ root = dataset_root(root)
74
+ output = Path(output).expanduser().resolve()
75
+ output_format = _OUTPUT_FORMATS.get(output.suffix.lower())
76
+ if output_format is None:
77
+ raise ValueError(
78
+ f"unsupported output suffix {output.suffix!r}; use Parquet, CSV, or JSON"
79
+ )
80
+ if output.is_relative_to(root / "data") or output.is_relative_to(root / "metadata"):
81
+ raise ValueError("output must not overwrite MEDS data or metadata")
82
+ output.parent.mkdir(parents=True, exist_ok=True)
83
+
84
+ metadata_json = read_dataset_metadata(root)
85
+ modifiers = code_modifier_columns(metadata_json)
86
+ metadata = scan_code_metadata(root)
87
+ events = scan_events(root, modifiers)
88
+ scope: list[str] = []
89
+ split_columns: list[tuple[str, str, str]] = []
90
+ if config.per_split or config.split_columns:
91
+ splits = scan_subject_splits(root)
92
+ duplicate = (
93
+ splits.group_by("subject_id").len().filter(pl.col("len") > 1).limit(1)
94
+ )
95
+ _raise_if_any(duplicate, "subject_splits contains duplicate subject IDs")
96
+ if config.split_columns:
97
+ split_values = (
98
+ splits.select("split")
99
+ .unique()
100
+ .sort("split")
101
+ .collect(engine="streaming")
102
+ .get_column("split")
103
+ .to_list()
104
+ )
105
+ if not split_values:
106
+ raise ValueError("subject_splits contains no splits")
107
+ split_columns = _split_columns(split_values)
108
+ events = events.join(splits, on="subject_id", how="left")
109
+ if config.per_split:
110
+ scope.append("split")
111
+
112
+ base_keys = ["code", *modifiers]
113
+ count_keys = [*scope, *base_keys]
114
+ split_count_names = _split_count_names(split_columns)
115
+ _validate_metadata(metadata, modifiers, config, split_count_names)
116
+ token = uuid4().hex
117
+ counts_path = output.with_name(f".{output.name}.{token}.counts.parquet")
118
+ staged_path = output.with_name(f".{output.name}.{token}.tmp{output.suffix}")
119
+
120
+ try:
121
+ _code_occurrences(events, count_keys, split_columns).sink_parquet(counts_path)
122
+ counts = pl.scan_parquet(counts_path)
123
+ if config.per_split:
124
+ missing_split = counts.filter(pl.col("split").is_null()).limit(1)
125
+ elif config.split_columns:
126
+ missing_split = counts.filter(pl.col("_missing_split_count") > 0).limit(1)
127
+ counts = counts.drop("_missing_split_count")
128
+ else:
129
+ missing_split = None
130
+ if missing_split is not None:
131
+ _raise_if_any(
132
+ missing_split,
133
+ "subject_splits does not cover every event subject",
134
+ )
135
+ safe = counts.filter(pl.col("subject_count") >= config.min_subjects)
136
+ if split_columns:
137
+ safe = _mask_split_counts(safe, split_columns, config.min_split_subjects)
138
+ common = safe.join(
139
+ metadata, on=base_keys, how="left", nulls_equal=True
140
+ ).with_columns(pl.lit(False).alias("is_masked"))
141
+ count_columns = [*_COUNT_COLUMNS, *split_count_names]
142
+ columns = [
143
+ *scope,
144
+ *metadata.collect_schema().names(),
145
+ *count_columns,
146
+ "is_masked",
147
+ ]
148
+ common = _round_counts(common, config.round_counts_to, count_columns).select(
149
+ columns
150
+ )
151
+
152
+ frames = [common]
153
+ if config.rare_code_action == "bucket":
154
+ rare_keys = counts.filter(
155
+ pl.col("subject_count") < config.min_subjects
156
+ ).select(count_keys)
157
+ rare_events = events.join(
158
+ rare_keys, on=count_keys, how="semi", nulls_equal=True
159
+ )
160
+ bucket = _rare_bucket(
161
+ rare_events, metadata, scope, split_columns, config
162
+ ).select(columns)
163
+ frames.append(bucket)
164
+
165
+ result = pl.concat(frames, how="vertical_relaxed").sort(
166
+ [*scope, "code", *modifiers]
167
+ )
168
+ _sink_output(result, staged_path, output_format)
169
+ staged_path.replace(output)
170
+ finally:
171
+ counts_path.unlink(missing_ok=True)
172
+ staged_path.unlink(missing_ok=True)
173
+
174
+ return output
175
+
176
+
177
+ def _split_columns(values: Sequence[str]) -> list[tuple[str, str, str]]:
178
+ result = []
179
+ used = set()
180
+ for value in values:
181
+ suffix = re.sub(r"\W+", "_", value).strip("_").lower() or "unnamed"
182
+ names = (f"event_count_{suffix}", f"subject_count_{suffix}")
183
+ if used & set(names):
184
+ raise ValueError("split names produce duplicate output columns")
185
+ used.update(names)
186
+ result.append((value, *names))
187
+ return result
188
+
189
+
190
+ def _split_count_names(split_columns: Sequence[tuple[str, str, str]]) -> list[str]:
191
+ return [name for columns in split_columns for name in columns[1:]]
192
+
193
+
194
+ def _count_expressions(
195
+ split_columns: Sequence[tuple[str, str, str]],
196
+ ) -> list[pl.Expr]:
197
+ expressions = [
198
+ pl.len().cast(pl.UInt64).alias("event_count"),
199
+ pl.col("subject_id").n_unique().cast(pl.UInt64).alias("subject_count"),
200
+ ]
201
+ for value, event_name, subject_name in split_columns:
202
+ in_split = pl.col("split") == value
203
+ expressions.extend(
204
+ [
205
+ in_split.sum().cast(pl.UInt64).alias(event_name),
206
+ pl.col("subject_id")
207
+ .filter(in_split)
208
+ .n_unique()
209
+ .cast(pl.UInt64)
210
+ .alias(subject_name),
211
+ ]
212
+ )
213
+ return expressions
214
+
215
+
216
+ def _mask_split_counts(
217
+ frame: pl.LazyFrame,
218
+ split_columns: Sequence[tuple[str, str, str]],
219
+ min_subjects: int,
220
+ ) -> pl.LazyFrame:
221
+ if min_subjects == 1:
222
+ return frame
223
+ return frame.with_columns(
224
+ expression
225
+ for _, event_name, subject_name in split_columns
226
+ for expression in (
227
+ pl.when(pl.col(subject_name) >= min_subjects)
228
+ .then(pl.col(event_name))
229
+ .otherwise(None)
230
+ .alias(event_name),
231
+ pl.when(pl.col(subject_name) >= min_subjects)
232
+ .then(pl.col(subject_name))
233
+ .otherwise(None)
234
+ .alias(subject_name),
235
+ )
236
+ )
237
+
238
+
239
+ def _sink_output(frame: pl.LazyFrame, path: Path, output_format: str) -> None:
240
+ if output_format == "parquet":
241
+ frame.sink_parquet(path)
242
+ elif output_format == "csv":
243
+ schema = frame.collect_schema()
244
+ nested = []
245
+ for name, dtype in schema.items():
246
+ column = pl.col(name)
247
+ if isinstance(dtype, pl.Array):
248
+ column = column.arr.to_list()
249
+ dtype = pl.List(dtype.inner)
250
+ if isinstance(dtype, pl.List):
251
+ nested.append(
252
+ column.list.eval(pl.element().cast(pl.String))
253
+ .list.join("|")
254
+ .alias(name)
255
+ )
256
+ elif isinstance(dtype, pl.Struct):
257
+ nested.append(column.struct.json_encode().alias(name))
258
+ frame.with_columns(nested).sink_csv(path)
259
+ elif output_format == "ndjson":
260
+ frame.sink_ndjson(path)
261
+ else:
262
+ lines = path.with_suffix(path.suffix + ".ndjson")
263
+ try:
264
+ frame.sink_ndjson(lines)
265
+ with lines.open("rb") as source, path.open("wb") as target:
266
+ target.write(b"[")
267
+ separator = b""
268
+ for line in source:
269
+ line = line.strip()
270
+ if line:
271
+ target.write(separator + line)
272
+ separator = b",\n"
273
+ target.write(b"]\n")
274
+ finally:
275
+ lines.unlink(missing_ok=True)
276
+
277
+
278
+ def _validate_metadata(
279
+ metadata: pl.LazyFrame,
280
+ modifiers: Sequence[str],
281
+ config: SummaryConfig,
282
+ extra_reserved: Sequence[str] = (),
283
+ ) -> None:
284
+ names = set(metadata.collect_schema().names())
285
+ missing = set(modifiers) - names
286
+ if missing:
287
+ raise ValueError(
288
+ f"Code metadata is missing declared modifiers: {sorted(missing)}"
289
+ )
290
+ collisions = names & (_RESERVED_COLUMNS | set(extra_reserved))
291
+ if collisions:
292
+ raise ValueError(
293
+ f"Code metadata uses reserved output columns: {sorted(collisions)}"
294
+ )
295
+ keys = ["code", *modifiers]
296
+ duplicate = metadata.group_by(keys).len().filter(pl.col("len") > 1).limit(1)
297
+ _raise_if_any(duplicate, f"Code metadata keys are not unique: {keys}")
298
+ if config.rare_code_action == "bucket":
299
+ collision = metadata.filter(pl.col("code") == config.rare_code_label).limit(1)
300
+ _raise_if_any(
301
+ collision, f"rare_code_label already exists: {config.rare_code_label!r}"
302
+ )
303
+
304
+
305
+ def _rare_bucket(
306
+ events: pl.LazyFrame,
307
+ metadata: pl.LazyFrame,
308
+ scope: Sequence[str],
309
+ split_columns: Sequence[tuple[str, str, str]],
310
+ config: SummaryConfig,
311
+ ) -> pl.LazyFrame:
312
+ expressions = _count_expressions(split_columns)
313
+ if scope:
314
+ result = events.group_by(list(scope)).agg(expressions)
315
+ else:
316
+ result = events.select(expressions)
317
+ result = result.filter(pl.col("subject_count") >= config.min_subjects)
318
+ if split_columns:
319
+ result = _mask_split_counts(result, split_columns, config.min_split_subjects)
320
+
321
+ schema = metadata.collect_schema()
322
+ values = [
323
+ pl.lit(config.rare_code_label, dtype=dtype).alias(name)
324
+ if name == "code"
325
+ else pl.lit(None, dtype=dtype).alias(name)
326
+ for name, dtype in schema.items()
327
+ ]
328
+ return _round_counts(
329
+ result.with_columns(*values, pl.lit(True).alias("is_masked")),
330
+ config.round_counts_to,
331
+ [*_COUNT_COLUMNS, *_split_count_names(split_columns)],
332
+ )
333
+
334
+
335
+ def _round_counts(
336
+ frame: pl.LazyFrame,
337
+ multiple: int | None,
338
+ columns: Sequence[str] = _COUNT_COLUMNS,
339
+ ) -> pl.LazyFrame:
340
+ if not multiple or multiple == 1:
341
+ return frame
342
+ return frame.with_columns(
343
+ (((pl.col(name) + multiple // 2) // multiple) * multiple)
344
+ .cast(pl.UInt64)
345
+ .alias(name)
346
+ for name in columns
347
+ )
348
+
349
+
350
+ def _raise_if_any(frame: pl.LazyFrame, message: str) -> None:
351
+ if frame.collect(engine="streaming").height:
352
+ raise ValueError(message)