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 +21 -0
- sumeds-0.1.0/PKG-INFO +100 -0
- sumeds-0.1.0/README.md +80 -0
- sumeds-0.1.0/pyproject.toml +39 -0
- sumeds-0.1.0/src/sumeds/__init__.py +26 -0
- sumeds-0.1.0/src/sumeds/cli.py +87 -0
- sumeds-0.1.0/src/sumeds/config.py +90 -0
- sumeds-0.1.0/src/sumeds/py.typed +0 -0
- sumeds-0.1.0/src/sumeds/scan.py +110 -0
- sumeds-0.1.0/src/sumeds/summary.py +352 -0
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)
|