feat(nlp): add a variety of --cohort-* args to filter notes

Three main new arguments:
* --cohort-csv: a csv with a column like patient_id or note_ref
* --cohort-anon-csv: same but with anonymized IDs
* --cohort-athena-table: same but points at an Athena table

To support the Athena option, we have some other new args:
* --athena-workgroup: specify the workgroup to use
* --athena-database: specify the database to use
* --allow-large-cohort: if the table is gigantic, use it anyway
  (this is here because we have a typo-guard in there - if you
  accidentally point at the base observation table, we're gonna
  stop you from downloading a terabyte of data)

You can specify the database in the table arg with a period.
And the workgroup can be specified via env var or CLI.

If we find a docref/dxreport ID/ref column, we'll use that. Otherwise,
we'll use a patient ID column and grab all notes for those patients.

This cohort filtering replaces instead of augments the previous default
filtering of "final" status notes (i.e. skipping draft or superceded
notes). But if the user is specifying the IDs manually for us, they
must know what they want and we don't need to do the status check for
them.

Author: mikix

compose.yaml

@@ -33,8 +33,11 @@ services:
     environment:
       # Environment variobles to pull in from the host
       - AWS_ACCESS_KEY_ID
+      - AWS_ATHENA_WORK_GROUP
       - AWS_DEFAULT_PROFILE
+      - AWS_DEFAULT_REGION
       - AWS_PROFILE
+      - AWS_REGION
       - AWS_SECRET_ACCESS_KEY
       - AWS_SESSION_TOKEN
       - AZURE_OPENAI_API_KEY

cumulus_etl/cli_utils.py

@@ -39,7 +39,7 @@ def add_auth(parser: argparse.ArgumentParser, *, use_fhir_url: bool = True):
     group.add_argument("--smart-jwks", metavar="PATH", help=argparse.SUPPRESS)
 
 
-def add_aws(parser: argparse.ArgumentParser) -> None:
+def add_aws(parser: argparse.ArgumentParser, athena: bool = False) -> None:
     group = parser.add_argument_group("AWS")
     group.add_argument(
         "--s3-region",
@@ -51,6 +51,22 @@ def add_aws(parser: argparse.ArgumentParser) -> None:
         metavar="KEY",
         help="if using S3 paths (s3://...), this is the KMS key ID to use",
     )
+    if athena:
+        group.add_argument(
+            "--athena-region",
+            metavar="REGION",
+            help="the region of your Athena workgroup (default is us-east-1)",
+        )
+        group.add_argument(
+            "--athena-workgroup",
+            metavar="GROUP",
+            help="the name of your Athena workgroup",
+        )
+        group.add_argument(
+            "--athena-database",
+            metavar="DB",
+            help="the name of your Athena database",
+        )
 
 
 def add_bulk_export(parser: argparse.ArgumentParser, *, as_subgroup: bool = True):
@@ -85,15 +101,13 @@ def add_bulk_export(parser: argparse.ArgumentParser, *, as_subgroup: bool = True
     return parser
 
 
-def add_nlp(parser: argparse.ArgumentParser):
-    group = parser.add_argument_group("NLP")
-    group.add_argument(
+def add_ctakes_override(parser: argparse.ArgumentParser):
+    parser.add_argument(
         "--ctakes-overrides",
         metavar="DIR",
         default="/ctakes-overrides",
         help="path to cTAKES overrides dir (default is /ctakes-overrides)",
     )
-    return group
 
 
 def add_output_format(parser: argparse.ArgumentParser) -> None:
@@ -105,20 +119,14 @@ def add_output_format(parser: argparse.ArgumentParser) -> None:
     )
 
 
-def add_task_selection(parser: argparse.ArgumentParser):
-    task = parser.add_argument_group("task selection")
-    task.add_argument(
+def add_task_selection(parser: argparse.ArgumentParser, *, etl_mode: bool):
+    default = ", default is all supported FHIR resources" if etl_mode else ""
+    required = not etl_mode
+    parser.add_argument(
         "--task",
         action="append",
-        help="only consider these tasks (comma separated, "
-        "default is all supported FHIR resources, "
-        "use '--task help' to see full list)",
-    )
-    task.add_argument(
-        "--task-filter",
-        action="append",
-        choices=["covid_symptom", "irae", "cpu", "gpu"],
-        help="restrict tasks to only the given sets (comma separated)",
+        help=f"only run these tasks (comma separated{default}, use '--task help' to see full list)",
+        required=required,
     )
 
 

cumulus_etl/errors.py

@@ -17,8 +17,8 @@
 TASK_UNKNOWN = 17
 CNLPT_MISSING = 18
 TASK_FAILED = 19
-TASK_FILTERED_OUT = 20
-TASK_SET_EMPTY = 21
+# TASK_FILTERED_OUT = 20 # Obsolete, we don't have --task-filter anymore
+# TASK_SET_EMPTY = 21 # Obsolete, we don't have --task-filter anymore
 ARGS_CONFLICT = 22
 ARGS_INVALID = 23
 # FHIR_URL_MISSING = 24 # Obsolete, it's no longer fatal
@@ -39,8 +39,13 @@
 INLINE_TASK_FAILED = 39
 INLINE_WITHOUT_FOLDER = 40
 WRONG_PHI_FOLDER = 41
-TASK_NOT_PROVIDED = 42
+# TASK_NOT_PROVIDED = 42  # checked now by argparse
 TASK_MISMATCH = 43
+ATHENA_TABLE_TOO_BIG = 44
+ATHENA_TABLE_NAME_INVALID = 45
+ATHENA_DATABASE_MISSING = 46
+MULTIPLE_COHORT_ARGS = 47
+COHORT_NOT_FOUND = 48
 
 
 class FatalError(Exception):

cumulus_etl/etl/cli.py

@@ -48,6 +48,14 @@ def define_etl_parser(parser: argparse.ArgumentParser) -> None:
     parser.add_argument(
         "--philter", action="store_true", help="run philter on all freeform text fields"
     )
+    parser.add_argument(
+        "--allow-missing-resources",
+        action="store_true",
+        help="run tasks even if their resources are not present",
+    )
+    cli_utils.add_task_selection(parser, etl_mode=True)
+
+    cli_utils.add_aws(parser)
 
     export = cli_utils.add_bulk_export(parser)
     export.add_argument(

cumulus_etl/etl/config.py

@@ -2,11 +2,12 @@
 
 import datetime
 import os
+from collections.abc import Callable
 from socket import gethostname
 
 import cumulus_fhir_support as cfs
 
-from cumulus_etl import common, errors, formats, store
+from cumulus_etl import common, deid, errors, formats, store
 
 
 class JobConfig:
@@ -39,6 +40,7 @@ def __init__(
         export_datetime: datetime.datetime | None = None,
         export_url: str | None = None,
         deleted_ids: dict[str, set[str]] | None = None,
+        resource_filter: Callable[[deid.Codebook, dict], bool] | None = None,
     ):
         self._dir_input_orig = dir_input_orig
         self.dir_input = dir_input_deid
@@ -59,6 +61,7 @@ def __init__(
         self.export_datetime = export_datetime
         self.export_url = export_url
         self.deleted_ids = deleted_ids or {}
+        self.resource_filter = resource_filter
 
         # initialize format class
         self._output_root = store.Root(self._dir_output, create=True)

cumulus_etl/etl/nlp/cli.py

@@ -9,8 +9,12 @@
 """
 
 import argparse
+import string
+from collections.abc import Callable
 
-from cumulus_etl import cli_utils, deid, loaders
+import pyathena
+
+from cumulus_etl import cli_utils, deid, errors, id_handling, loaders
 from cumulus_etl.etl import pipeline
 
 
@@ -19,12 +23,131 @@ def define_nlp_parser(parser: argparse.ArgumentParser) -> None:
     parser.usage = "%(prog)s [OPTION]... INPUT OUTPUT PHI"
 
     pipeline.add_common_etl_args(parser)
-    cli_utils.add_nlp(parser)
+    cli_utils.add_ctakes_override(parser)
+    cli_utils.add_task_selection(parser, etl_mode=False)
+
+    cli_utils.add_aws(parser, athena=True)
+
+    group = parser.add_argument_group("cohort selection")
+    group.add_argument(
+        "--cohort-csv",
+        metavar="FILE",
+        help="path to a .csv file with original patient and/or note IDs",
+    )
+    group.add_argument(
+        "--cohort-anon-csv",
+        metavar="FILE",
+        help="path to a .csv file with anonymized patient and/or note IDs",
+    )
+    group.add_argument(
+        "--cohort-athena-table",
+        metavar="DB.TABLE",
+        help="name of an Athena table with patient and/or note IDs",
+    )
+    group.add_argument(
+        "--allow-large-cohort",
+        action="store_true",
+        help="allow a larger-than-normal cohort",
+    )
+
+
+def get_cohort_filter(args: argparse.Namespace) -> Callable[[deid.Codebook, dict], bool] | None:
+    """Returns (patient refs to match, resource refs to match)"""
+    # Poor man's add_mutually_exclusive_group(), which we don't use because we have additional
+    # flags for the group, like "--allow-large-cohort".
+    has_csv = bool(args.cohort_csv)
+    has_anon_csv = bool(args.cohort_anon_csv)
+    has_athena_table = bool(args.cohort_athena_table)
+    arg_count = int(has_csv) + int(has_anon_csv) + int(has_athena_table)
+    if not arg_count:
+        return None
+    elif arg_count > 1:
+        errors.fatal(
+            "Multiple cohort arguments provided. Please specify just one.",
+            errors.MULTIPLE_COHORT_ARGS,
+        )
+
+    if has_athena_table:
+        if "." in args.cohort_athena_table:
+            parts = args.cohort_athena_table.split(".", 1)
+            database = parts[0]
+            table = parts[-1]
+        else:
+            database = args.athena_database
+            table = args.cohort_athena_table
+        if not database:
+            errors.fatal(
+                "You must provide an Athena database with --athena-database.",
+                errors.ATHENA_DATABASE_MISSING,
+            )
+        if set(table) - set(string.ascii_letters + string.digits + "-_"):
+            errors.fatal(
+                f"Athena table name '{table}' has invalid characters.",
+                errors.ATHENA_TABLE_NAME_INVALID,
+            )
+        cursor = pyathena.connect(
+            region_name=args.athena_region,
+            work_group=args.athena_workgroup,
+            schema_name=database,
+        ).cursor()
+        count = cursor.execute(f'SELECT count(*) FROM "{table}"').fetchone()[0]  # noqa: S608
+        if int(count) > 20_000 and not args.allow_large_cohort:
+            errors.fatal(
+                f"Athena cohort in '{table}' is very large ({int(count):,} rows).\n"
+                "If you want to use it anyway, pass --allow-large-cohort",
+                errors.ATHENA_TABLE_TOO_BIG,
+            )
+        csv_file = cursor.execute(f'SELECT * FROM "{table}"').output_location  # noqa: S608
+    else:
+        csv_file = args.cohort_anon_csv or args.cohort_csv
+
+    is_anon = has_anon_csv or has_athena_table
+
+    dxreport_ids = id_handling.get_ids_from_csv(csv_file, "DiagnosticReport", is_anon=is_anon)
+    docref_ids = id_handling.get_ids_from_csv(csv_file, "DocumentReference", is_anon=is_anon)
+    patient_ids = id_handling.get_ids_from_csv(csv_file, "Patient", is_anon=is_anon)
+
+    if not dxreport_ids and not docref_ids and not patient_ids:
+        errors.fatal("No patient or note IDs found in cohort.", errors.COHORT_NOT_FOUND)
+
+    def res_filter(codebook: deid.Codebook, resource: dict) -> bool:
+        match resource["resourceType"]:
+            # TODO: uncomment once we support DxReport NLP (coming soon)
+            # case "DiagnosticReport":
+            #     id_pool = dxreport_ids
+            #     patient_ref = resource.get("subject", {}).get("reference")
+            case "DocumentReference":
+                id_pool = docref_ids
+                patient_ref = resource.get("subject", {}).get("reference")
+            case _:  # pragma: no cover
+                # shouldn't happen
+                return False  # pragma: no cover
+
+        # Check if we have an exact resource ID match (if the user defined exact IDs, we only use
+        # them, and don't do any patient matching)
+        if id_pool:
+            res_id = resource["id"]
+            if is_anon:
+                res_id = codebook.fake_id(resource["resourceType"], res_id, caching_allowed=False)
+            return res_id in id_pool
+
+        # Else match on patients if no resource IDs were defined
+        if not patient_ref:
+            return False
+        patient_id = patient_ref.removeprefix("Patient/")
+        if is_anon:
+            patient_id = codebook.fake_id("Patient", patient_id, caching_allowed=False)
+        return patient_id in patient_ids
+
+    return res_filter
 
 
 async def nlp_main(args: argparse.Namespace) -> None:
+    res_filter = get_cohort_filter(args)
+
     async def prep_scrubber(_results: loaders.LoaderResults) -> tuple[deid.Scrubber, dict]:
-        return deid.Scrubber(args.dir_phi), {"ctakes_overrides": args.ctakes_overrides}
+        config_args = {"ctakes_overrides": args.ctakes_overrides, "resource_filter": res_filter}
+        return deid.Scrubber(args.dir_phi), config_args
 
     await pipeline.run_pipeline(args, prep_scrubber=prep_scrubber, nlp=True)
 

cumulus_etl/etl/pipeline.py

@@ -63,15 +63,8 @@ def add_common_etl_args(parser: argparse.ArgumentParser) -> None:
     parser.add_argument(
         "--errors-to", metavar="DIR", help="where to put resources that could not be processed"
     )
-    parser.add_argument(
-        "--allow-missing-resources",
-        action="store_true",
-        help="run tasks even if their resources are not present",
-    )
 
-    cli_utils.add_aws(parser)
     cli_utils.add_auth(parser)
-    cli_utils.add_task_selection(parser)
     cli_utils.add_debugging(parser)
 
 
@@ -137,7 +130,8 @@ async def check_available_resources(
     # Reconciling is helpful for performance reasons (don't need to finalize untouched tables),
     # UX reasons (can tell user if they made a CLI mistake), and completion tracking (don't
     # mark a resource as complete if we didn't even export it)
-    if args.allow_missing_resources:
+    has_allow_missing = hasattr(args, "allow_missing_resources")
+    if has_allow_missing and args.allow_missing_resources:
         return requested_resources
 
     detected = await loader.detect_resources()
@@ -158,7 +152,8 @@ async def check_available_resources(
                 )
         else:
             msg = "Required resources not found.\n"
-            msg += "Add --allow-missing-resources to run related tasks anyway with no input."
+            if has_allow_missing:
+                msg += "Add --allow-missing-resources to run related tasks anyway with no input."
             errors.fatal(msg, errors.MISSING_REQUESTED_RESOURCES)
 
     return requested_resources
@@ -182,8 +177,8 @@ async def run_pipeline(
     job_context = context.JobContext(root_phi.joinpath("context.json"))
     job_datetime = common.datetime_now()  # grab timestamp before we do anything
 
-    selected_tasks = task_factory.get_selected_tasks(args.task, args.task_filter, nlp=nlp)
-    is_default_tasks = not args.task and not args.task_filter
+    selected_tasks = task_factory.get_selected_tasks(args.task, nlp=nlp)
+    is_default_tasks = not args.task
 
     # Print configuration
     print_config(args, job_datetime, selected_tasks)

cumulus_etl/etl/tasks/base.py

@@ -416,8 +416,10 @@ async def read_entries(self, *, progress: rich.progress.Progress = None) -> Entr
         Something like "yield x, y" or "yield x, [y, z]" - these streams of entries will be kept
         separated into two different DataFrames.
         """
+        resource_filter = self.task_config.resource_filter
         for x in filter(self.scrubber.scrub_resource, self.read_ndjson(progress=progress)):
-            yield x
+            if not resource_filter or resource_filter(self.scrubber.codebook, x):
+                yield x
 
     def table_batch_cleanup(self, table_index: int, batch_index: int) -> None:
         """Override to add any necessary cleanup from writing a batch out (releasing memory etc)"""

cumulus_etl/etl/tasks/nlp_task.py

@@ -84,10 +84,12 @@ async def read_notes(
         """
         warned_connection_error = False
 
+        note_filter = self.task_config.resource_filter or nlp.is_docref_valid
+
         for docref in self.read_ndjson(progress=progress):
             orig_docref = copy.deepcopy(docref)
             can_process = (
-                nlp.is_docref_valid(docref)
+                note_filter(self.scrubber.codebook, docref)
                 and (doc_check is None or doc_check(docref))
                 and self.scrubber.scrub_resource(docref, scrub_attachments=False, keep_stats=False)
             )