Intake wdl

docker/Dockerfile

@@ -3,12 +3,16 @@ FROM eu.gcr.io/finngen-refinery-dev/bioinformatics:0.8
 RUN apt-get update && apt-get upgrade --yes
 RUN curl https://clickhouse.com/ | sh && ./clickhouse install
 RUN apt install pigz
-ADD ./docker/requirements.txt requirements.txt
-RUN pip3 install --upgrade scipy
-RUN  pip3 install -r requirements.txt
 
-ADD finngen_qc finngen_qc
-ADD core core
-ADD sb_release sb_release
-ADD qc_scripts qc_scripts
+# Install kanta package
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh
+ENV PATH="/root/.local/bin:$PATH"
+RUN uv python install 3.14
+ADD pyproject.toml pyproject.toml
+ADD README.md README.md
+ADD uv.lock uv.lock
+ADD src src
+ENV UV_PROJECT_ENVIRONMENT="/opt/kanta-venv"
+RUN uv sync --python 3.14 --locked
+ENV PATH="/opt/kanta-venv/bin:$PATH"
 

src/kanta/intake/__main__.py

@@ -53,6 +53,9 @@
     )
     output.check_safe_write(output_file_assemble_stage)
     output.check_safe_write(output_file_tidyup_stage)
+    output.check_safe_write(
+        output_file_tidyup_stage.with_name(f"{output_file_tidyup_stage.stem}_duplicates.parquet")
+    )
 
     tmp_dir = output.create_tmp_dir()
 

src/kanta/intake/tidyup.py

@@ -5,7 +5,6 @@
 
 Differences from the WDL implementation
 =======================================
-- No logging of duplicates/err lines.
 - Outputs to a single parquet file, no .txt.gz, as this is very slow.
 
 
@@ -18,19 +17,26 @@
 
 Lowest tested working spec: 8 CPUs / 8 GB RAM, 32 buckets. Runs in 6-12 min.
 
-If failing due to OOM in the sort+dedup stage, try increasing the bucket count.
+If failing due to OOM, try increasing the bucket count — both the sort+dedup
+stage and the final join+sanitize stage process one bucket at a time.
 
 The GCP VM type appears to matter. N2D is about 2x faster than E2.
 """
 
 from argparse import ArgumentParser
+from datetime import datetime
 from pathlib import Path
 
 import polars as pl
+import pyarrow.parquet as pq
 
 from kanta import output
 
 
+def log_step(message: str) -> None:
+    print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}] {message}", flush=True)
+
+
 COLUMNS_UNIQUENESS_SORT = [
     "FINNGENID",
     "APPROX_EVENT_DAY",
@@ -40,6 +46,8 @@
     "tutkimusvastauksentila",
     "tutkimustulosarvo",
     "tutkimustulosyksikko",
+    "tuloksenpoikkeavuus",
+    "tutkimustulosteksti"    
 ]
 
 
@@ -66,45 +74,48 @@ def main(
     tmp_dir_sort_dedup = tmp_dir / "sort_dedup"
     tmp_dir_sort_dedup.mkdir()
 
-    print("# Consolidate")
+    tmp_dir_duplicates = tmp_dir / "duplicates"
+    tmp_dir_duplicates.mkdir()
+
+    log_step("# Consolidate: start")
     consolidated_file = consolidate_columns(assembled_file, tmp_file_consolidate)
+    log_step("# Consolidate: done")
 
-    print("# Partition")
+    log_step("# Partition: start")
     partition(consolidated_file, tmp_dir_partition, partition_n_buckets)
-
-    print("# Sort + Dedup")
-    for bucket_file in tmp_dir_partition.glob("bucket_id__*.parquet"):
-        (
-            pl.scan_parquet(bucket_file)
-            .pipe(sort_dedup)
-            .sink_parquet(tmp_dir_sort_dedup / bucket_file.name)
-        )
-
-    print("# Concatenate + join SEX")
-    bucket_files = []
-    for bucket_id in range(partition_n_buckets):
-        bucket_files.append(tmp_dir_sort_dedup / f"bucket_id__{bucket_id}.parquet")
+    log_step("# Partition: done")
+
+    log_step("# Sort + Dedup: start")
+    bucket_files_to_process = sorted(tmp_dir_partition.glob("bucket_id__*.parquet"))
+    total_n_before = 0
+    total_n_after = 0
+    bucket_row_counts = {}
+    for i, bucket_file in enumerate(bucket_files_to_process, start=1):
+        log_step(f"# Sort + Dedup: bucket {i}/{len(bucket_files_to_process)} ({bucket_file.name})")
+        n_before = pl.scan_parquet(bucket_file).select(pl.len()).collect().item()
+        output_bucket_file = tmp_dir_sort_dedup / bucket_file.name
+        kept, dropped = sort_dedup(pl.scan_parquet(bucket_file))
+        kept.sink_parquet(output_bucket_file)
+        dropped.sink_parquet(tmp_dir_duplicates / bucket_file.name)
+        n_after = pl.scan_parquet(output_bucket_file).select(pl.len()).collect().item()
+        total_n_before += n_before
+        total_n_after += n_after
+        bucket_row_counts[bucket_file.name] = n_after
+    n_removed = total_n_before - total_n_after
+    pct_removed = 100 * n_removed / total_n_before if total_n_before else 0
+    log_step(f"# Sort + Dedup: done, removed {n_removed} duplicate rows total ({pct_removed:.2f}% of {total_n_before})")
+
+    duplicates_output_file = output_file.with_name(f"{output_file.stem}_duplicates.parquet")
+    log_step(f"# Sort + Dedup: writing duplicate rows to {duplicates_output_file}")
+    duplicate_bucket_files = sorted(tmp_dir_duplicates.glob("bucket_id__*.parquet"))
+    concatenate_parquet_files(duplicate_bucket_files, duplicates_output_file)
 
     df_pheno = pl.scan_csv(
         phenotype_file,
         infer_schema=False,
         separator="\t",
     ).select("FINNGENID", "SEX")
 
-    df_concat = (
-        pl.scan_parquet(bucket_files)
-        # Join SEX
-        .join(
-            df_pheno,
-            left_on="FINNGENID",
-            right_on="FINNGENID",
-            how="left",
-            maintain_order="left",
-        )
-        .with_row_index(name="_rowid", offset=1)
-    )
-
-    print("# Sanitize text fields")
     # Unicode "SYMBOL FOR NEWLINE", displayed as: ␤
     unicode_newline = "\u2424"
     # Unicode "SYMBOL FOR HORIZONTAL TABULATION", displayed as: ␉
@@ -120,42 +131,56 @@ def main(
         "load_id_pseudo",
         "file_name_pseudo",
         "laboratoriotutkimusoid",
-        "_rowid",
+        "ROWID",
         "_rowid_source",
         "SEX",
     ]
-    (
-        df_concat.with_columns(
-            pl.selectors.exclude(*trusted_columns)
-            .str.replace_all(pattern="\r\n|\r|\n", value=unicode_newline)
-            .str.replace_all(pattern="\t", value=unicode_tab, literal=True)
+
+    tmp_dir_final = tmp_dir / "final"
+    tmp_dir_final.mkdir()
+
+    # Process one bucket at a time (rather than one join + with_row_index +
+    # sanitize + sink over the whole dataset) because maintain_order="left" on
+    # the join and the global counter in with_row_index both disable Polars'
+    # streaming engine, forcing full materialization of the joined/indexed/
+    # sanitized frame in memory. Per-bucket, each chunk is small enough to
+    # materialize safely; ROWID stays globally monotonic via a running offset.
+    log_step("# Join SEX + Sanitize: start (per bucket)")
+    row_index_offset = 1
+    final_bucket_files = []
+    for bucket_id in range(partition_n_buckets):
+        bucket_name = f"bucket_id__{bucket_id}.parquet"
+        bucket_file = tmp_dir_sort_dedup / bucket_name
+        n_rows = bucket_row_counts[bucket_name]
+        log_step(
+            f"# Join SEX + Sanitize: bucket {bucket_id + 1}/{partition_n_buckets} "
+            f"({bucket_name}), {n_rows} rows, ROWID offset {row_index_offset}"
         )
-        # Re-order column to be somewhat backward compatible with previous implementation
-        .select(
-            "_rowid",
-            "_rowid_source",
-            "FINNGENID",
-            "EVENT_AGE",
-            "APPROX_EVENT_DAY",
-            "TIME",
-            "laboratoriotutkimusnimike",
-            "paikallinentutkimusnimike_koodi",
-            "paikallinentutkimusnimike_selite",
-            "tutkimuskoodistonjarjestelma",
-            "tutkimusvastauksentila",
-            "tutkimustulosarvo",
-            "tutkimustulosyksikko",
-            "tuloksenpoikkeavuus",
-            "viitearvoryhma",
-            "viitevalialkuarvo",
-            "viitevalialkuyksikko",
-            "viitevaliloppuarvo",
-            "viitevaliloppuyksikko",
-            "tutkimustulosteksti",
-            "SEX",
+        output_bucket_file = tmp_dir_final / bucket_name
+        (
+            pl.scan_parquet(bucket_file)
+            .join(
+                df_pheno,
+                left_on="FINNGENID",
+                right_on="FINNGENID",
+                how="left",
+                maintain_order="left",
+            )
+            .with_row_index(name="ROWID", offset=row_index_offset)
+            .with_columns(
+                pl.selectors.exclude(*trusted_columns)
+                .str.replace_all(pattern="\r\n|\r|\n", value=unicode_newline)
+                .str.replace_all(pattern="\t", value=unicode_tab, literal=True)
+            )
+            .sink_parquet(output_bucket_file)
         )
-        .sink_parquet(output_file)
-    )
+        final_bucket_files.append(output_bucket_file)
+        row_index_offset += n_rows
+    log_step("# Join SEX + Sanitize: done")
+
+    log_step("# Concatenate final bucket files: start")
+    concatenate_parquet_files(final_bucket_files, output_file)
+    log_step("# Concatenate final bucket files: done")
 
 
 def init_cli():
@@ -196,27 +221,16 @@ def init_cli():
 
 
 def consolidate_columns(assembled_file: Path, output_file: Path) -> Path:
-    """Remove unnecessary columns form the assembled file and rename the ones we will keep."""
+    """Keep all main.* columns (stripped of prefix) plus freetext text field and _rowid_source."""
+    schema = pl.scan_parquet(assembled_file).collect_schema()
+
+    internal_columns = {"main._rowid", "main._filename"}
     rename_columns = {
-        "main.FINNGENID": "FINNGENID",
-        "main.EVENT_AGE": "EVENT_AGE",
-        "main.APPROX_EVENT_DAY": "APPROX_EVENT_DAY",
-        "main.TIME": "TIME",
-        "main.laboratoriotutkimusnimike": "laboratoriotutkimusnimike",
-        "main.paikallinentutkimusnimike_koodi": "paikallinentutkimusnimike_koodi",
-        "main.paikallinentutkimusnimike_selite": "paikallinentutkimusnimike_selite",
-        "main.tutkimuskoodistonjarjestelma": "tutkimuskoodistonjarjestelma",
-        "main.tutkimusvastauksentila": "tutkimusvastauksentila",
-        "main.tutkimustulosarvo": "tutkimustulosarvo",
-        "main.tutkimustulosyksikko": "tutkimustulosyksikko",
-        "main.tuloksenpoikkeavuus": "tuloksenpoikkeavuus",
-        "main.viitearvoryhma": "viitearvoryhma",
-        "main.viitevalialkuarvo": "viitevalialkuarvo",
-        "main.viitevalialkuyksikko": "viitevalialkuyksikko",
-        "main.viitevaliloppuarvo": "viitevaliloppuarvo",
-        "main.viitevaliloppuyksikko": "viitevaliloppuyksikko",
-        "freetext.tutkimustulosteksti": "tutkimustulosteksti",
+        col: col.removeprefix("main.")
+        for col in schema.names()
+        if col.startswith("main.") and col not in internal_columns
     }
+    rename_columns["freetext.tutkimustulosteksti"] = "tutkimustulosteksti"
 
     out_columns = list(rename_columns.keys()) + ["_rowid_source"]
 
@@ -232,14 +246,37 @@ def consolidate_columns(assembled_file: Path, output_file: Path) -> Path:
 
 def partition(assembled_file: Path, tmp_dir: Path, n_buckets):
     for bucket_id in range(n_buckets):
+        log_step(f"# Partition: bucket {bucket_id + 1}/{n_buckets}")
         (
             pl.scan_parquet(assembled_file)
             .filter(pl.col("FINNGENID").hash() % n_buckets == bucket_id)
             .sink_parquet(tmp_dir / f"bucket_id__{bucket_id}.parquet")
         )
 
 
+def concatenate_parquet_files(input_files: list[Path], output_file: Path) -> None:
+    """Concatenate parquet files by copying row batches directly, one at a time.
+
+    Polars' `sink_parquet` over a multi-file `scan_parquet` is not reliably
+    streaming in all versions/engines, and has been observed to materialize
+    the whole dataset in memory. Using pyarrow directly guarantees a bounded
+    memory footprint regardless of total row count.
+    """
+    writer = None
+    try:
+        for input_file in input_files:
+            reader = pq.ParquetFile(input_file)
+            if writer is None:
+                writer = pq.ParquetWriter(output_file, reader.schema_arrow, compression="zstd")
+            for batch in reader.iter_batches():
+                writer.write_batch(batch)
+    finally:
+        if writer is not None:
+            writer.close()
+
+
 def sort_dedup(frame: pl.LazyFrame | pl.DataFrame):
+    """Sort by the full column order, then split into kept rows (first per duplicate key) and dropped duplicates."""
     all_columns = frame.collect_schema().names()
     sort_subset_columns = set(COLUMNS_UNIQUENESS_SORT)
     other_columns = []
@@ -249,15 +286,19 @@ def sort_dedup(frame: pl.LazyFrame | pl.DataFrame):
 
     sort_full_columns = COLUMNS_UNIQUENESS_SORT + other_columns
 
-    return frame.sort(by=sort_full_columns).unique(
-        subset=COLUMNS_UNIQUENESS_SORT, keep="first", maintain_order=True
-    )
+    sorted_frame = frame.sort(by=sort_full_columns)
+    is_first = pl.struct(COLUMNS_UNIQUENESS_SORT).is_first_distinct()
+
+    kept = sorted_frame.filter(is_first)
+    dropped = sorted_frame.filter(~is_first)
+    return kept, dropped
 
 
 if __name__ == "__main__":
     args = init_cli()
 
     output.check_safe_write(args.output_file)
+    output.check_safe_write(args.output_file.with_name(f"{args.output_file.stem}_duplicates.parquet"))
     tmp_dir = output.create_tmp_dir()
 
     main(