setup clean base for Pietro

pyproject.toml

@@ -7,6 +7,7 @@ requires-python = ">=3.12"
 dependencies = [
     "pandas>=3.0.2",
     "polars>=1.40.0",
+    "pyarrow>=24.0.0",
 ]
 
 [dependency-groups]

src/kanta/config.py

@@ -0,0 +1,19 @@
+# Shared config across the project.
+
+# Aliases mapping column name from input (dict keys) to a name used in the code (dict values).
+# The purpose is to expose easier column names that can be referenced in the data processing code.
+# For more column names, check: https://github.com/FINNGEN/Kanta_lab_QC#column-description
+COLUMN_ALIASES = {
+    "tutkimuskoodistonjarjestelma": "CODING_SYSTEM",
+    "paikallinentutkimusnimike_selite": "TEST_NAME_ABBREVIATION",
+    "tutkimustulosarvo": "MEASUREMENT_VALUE",
+    "tutkimustulosyksikko": "MEASUREMENT_UNIT",
+    "tutkimusvastauksentila": "MEASUREMENT_STATUS",
+    "tuloksenpoikkeavuus": "TEST_OUTCOME",
+    "viitearvoryhma": "REFERENCE_RANGE_GROUP",
+    "viitevalialkuarvo": "REFERENCE_RANGE_LOWER_VALUE",
+    "viitevalialkuyksikko": "REFERENCE_RANGE_LOWER_UNIT",
+    "viitevaliloppuarvo": "REFERENCE_RANGE_UPPER_VALUE",
+    "viitevaliloppuyksikko": "REFERENCE_RANGE_UPPER_UNIT",
+    "tutkimustulosteksti": "MEASUREMENT_FREE_TEXT",
+}

src/kanta/engine/__main__.py

@@ -0,0 +1,129 @@
+import multiprocessing as mp
+import os
+from argparse import ArgumentParser
+from functools import partial
+from pathlib import Path
+
+from kanta import output
+from kanta.engine import chunking, processing
+
+
+def main(
+    input_file: Path,
+    output_file: Path,
+    tmp_dir: Path,
+    *,
+    is_test_run=False,
+    n_workers=1,
+):
+    # Setup
+    processing.configure_pandas()
+
+    chunks_dir = tmp_dir / "chunks"
+    chunks_dir.mkdir()
+
+    # Iterate over each chunk
+    iter_indexed_chunks = chunking.chunk_iterator(input_file, is_test_run=is_test_run)
+
+    process_func = partial(processing.process_chunk, chunks_dir=chunks_dir)
+
+    if n_workers > 1:
+        process_in_parallel(process_func, iter_indexed_chunks, n_workers=n_workers)
+    else:
+        for indexed_chunk in iter_indexed_chunks:
+            process_func(indexed_chunk)
+
+    chunking.concatenate_chunks(chunks_dir, output_file)
+
+
+def process_in_parallel(func, indexed_chunks, *, n_workers: int):
+    """Process the chunks in parallel using `n_workers` spawned processes."""
+    # Explicitly use the "spawn" method to create workers for consistent behavior across OSes
+    # and Python versions.
+    ctx = mp.get_context("spawn")
+
+    # Setting the `initializer` here is required since we used the "spawn" method above to
+    # start workers: they start with no configuration, so we must provide it.
+    #
+    # NOTE(Vincent 2026-06-17):
+    # It looks like `multiprocessing.Pool` has a non-trivial silent failure mode: if a
+    # worker process get OOM-killed (i.e. because available memory is low) then the killed
+    # worker will fail silently and the Pool will hang forever.
+    # IMHO we should leave as it is for now and make sure to monitor memory usage. The future
+    # rewrite to Polars will remove the use of `multiprocessing` and this problem.
+    # See: https://bugs.python.org/issue22393
+    with ctx.Pool(n_workers, initializer=processing.configure_pandas) as pool:
+        for _result in pool.imap_unordered(func, indexed_chunks):
+            # Consume the iterator, discard the result since it's written to disk.
+            pass
+
+
+def init_cli():
+    parser = ArgumentParser(
+        description="Kanta Lab preprocessing pipeline: raw data ⇒ clean data."
+    )
+    parser.add_argument(
+        "--input-file",
+        type=Path,
+        help="Path to the Kanta Lab data file coming from the intake stage (Parquet)",
+        required=True,
+    )
+    parser.add_argument(
+        "--test",
+        action="store_true",
+        help=(
+            "Process only the first chunk (useful for debugging). "
+            "This overwrites --n-workers to 1."
+        ),
+        required=False,
+    )
+    parser.add_argument(
+        "--output-file",
+        type=Path,
+        help="Output file path (Parquet)",
+        required=True,
+    )
+    parser.add_argument(
+        "--n-workers",
+        type=int,
+        default=os.process_cpu_count() or 1,
+        help=(
+            "Number of worker processes used to process chunks in parallel. "
+            "Defaults to the number of available CPUs. Use 1 to run serially "
+            "(useful for debugging)."
+        ),
+        required=False,
+    )
+    parser.add_argument(
+        "--keep-intermediate-files",
+        help="Keep intermediate files, useful for debugging.",
+        action="store_true",
+    )
+
+    args = parser.parse_args()
+
+    if args.n_workers < 1:
+        raise ValueError("--n-workers must be 1 or more")
+
+    if args.test:
+        args.n_workers = 1
+
+    return args
+
+
+if __name__ == "__main__":
+    args = init_cli()
+
+    output.check_safe_write(args.output_file)
+    tmp_dir = output.create_tmp_dir()
+
+    main(
+        args.input_file,
+        args.output_file,
+        tmp_dir,
+        is_test_run=args.test,
+        n_workers=args.n_workers,
+    )
+
+    if not args.keep_intermediate_files:
+        output.teardown_dir(tmp_dir)

src/kanta/engine/chunking.py

@@ -0,0 +1,113 @@
+import shutil
+from collections.abc import Iterator
+from pathlib import Path
+
+import pandas as pd
+import pyarrow.parquet as pq
+
+
+# Which columns to read from the input file. This is used very early in the pipeline to limit the
+# amount of data read, so the column renaming has not been done yet, hence using the original
+# column names.
+# TODO(Vincent 2026-06-17)  This should use the canonical column aliases from the config. A good
+# time to do that would be when the data processing is in place so that all the column reference
+# the canonical column aliases.
+READ_COLUMNS = [
+    "FINNGENID",
+    "EVENT_AGE",
+    "tutkimuskoodistonjarjestelma",
+    "paikallinentutkimusnimike_selite",
+    "tutkimustulosarvo",
+    "tutkimustulosyksikko",
+    "tutkimusvastauksentila",
+    "tuloksenpoikkeavuus",
+    "viitearvoryhma",
+    "viitevalialkuarvo",
+    "viitevalialkuyksikko",
+    "viitevaliloppuarvo",
+    "viitevaliloppuyksikko",
+    "tutkimustulosteksti",
+    "paikallinentutkimusnimike_koodi",
+    "laboratoriotutkimusnimike",
+    "APPROX_EVENT_DAY",
+    "TIME",
+    "_rowid",
+    "_rowid_source",
+    "SEX",
+]
+
+# Number of rows per chunk when streaming the input Parquet file.
+# The value is independent of the number of CPUs: the memory used by the engine
+# is already proportional to the number of workers, so scaling the number of
+# rows per chunk by the number of workers would make the memory use scale by
+# (N workers × N workers).
+N_LINES_PER_CHUNK = 200_000
+CHUNKS_FILE_TEMPLATE = "chunk_{index:06d}.parquet"
+CHUNKS_FILE_GLOB = "chunk_*.parquet"
+
+
+def chunk_iterator(
+    input_file: Path, *, is_test_run: bool = False
+) -> Iterator[tuple[int, pd.DataFrame]]:
+    # Use pyarrow to read the Parquet file in chunks.
+    parquet_file = pq.ParquetFile(input_file)
+
+    chunk_index = 0
+    for batch in parquet_file.iter_batches(
+        batch_size=N_LINES_PER_CHUNK,
+        # Select only the given columns, this speeds up the read quite a lot for Parquet files.
+        columns=READ_COLUMNS,
+    ):
+        yield (
+            chunk_index,
+            batch.to_pandas(
+                # Use nullable data-types backed by pyarrow.
+                types_mapper=pd.ArrowDtype
+            ),
+        )
+
+        chunk_index += 1
+
+        # Yield only the first item when test run is ON
+        if is_test_run:
+            break
+
+
+def write_chunk(dataframe: pd.DataFrame, chunks_dir: Path, chunk_index: int) -> Path:
+    """Write one processed chunk to its own Parquet chunk file.
+
+    Use `chunk_index` to keep track of order for later in-order concatenation.
+    """
+    chunk_path = chunks_dir / CHUNKS_FILE_TEMPLATE.format(index=chunk_index)
+    dataframe.to_parquet(chunk_path, engine="pyarrow", compression="zstd", index=False)
+    return chunk_path
+
+
+def concatenate_chunks(chunks_dir: Path, output_file: Path, cleanup: bool = True):
+    """Concatenate chunks in order so no sorting is required afterwards.
+
+    The order relies on the filename, which holds the chunk index.
+    """
+    chunks = sorted(chunks_dir.glob(CHUNKS_FILE_GLOB), key=get_chunk_index)
+
+    writer = None
+    try:
+        for chunk_path in chunks:
+            table = pq.read_table(chunk_path)
+
+            # Initialize the ParquetWriter, needs the schema.
+            if writer is None:
+                writer = pq.ParquetWriter(output_file, table.schema)
+            writer.write_table(table)
+
+    finally:
+        if writer is not None:
+            writer.close()
+
+    if cleanup:
+        shutil.rmtree(chunks_dir)
+
+
+def get_chunk_index(chunk_path: Path) -> int:
+    """Extract the integer chunk index from a chunk file name (e.g. chunk_000007 -> 7)."""
+    return int(chunk_path.stem.split("_")[-1])

src/kanta/engine/pipes.py

@@ -0,0 +1,23 @@
+import pandas as pd
+
+from kanta import config
+
+
+def noop_filter(df):
+    return df
+
+
+def rename_cols(df: pd.DataFrame) -> pd.DataFrame:
+    return df.rename(columns=config.COLUMN_ALIASES)
+
+
+def run_all(df: pd.DataFrame) -> pd.DataFrame:
+    return (
+        df.pipe(rename_cols)
+        .pipe(noop_filter)
+        .pipe(noop_filter)
+        .pipe(noop_filter)
+        .pipe(noop_filter)
+        .pipe(noop_filter)
+        .pipe(noop_filter)
+    )

src/kanta/engine/processing.py

@@ -0,0 +1,41 @@
+import warnings
+from pathlib import Path
+
+import pandas as pd
+
+from kanta.engine import chunking, pipes
+
+
+def configure_pandas():
+    """Set preferred Pandas behavior via options.
+
+    IMPORTANT: This function must be called when initializing workers for multiprocessing, since
+    creating them with the 'spawn' method doesn't carry over the Pandas configuration.
+    """
+    # Treat NaN (a real float value) and NA (a missing value) as distinct.
+    pd.options.future.distinguish_nan_and_na = True
+
+    # Default to the pyarrow engine for all Parquet reads/writes, so we don't have to pass
+    # engine="pyarrow" on every call.
+    pd.options.io.parquet.engine = "pyarrow"
+
+    # Turn chained assignment into a hard error.
+    # Chained assignments have weird behavior in that they would turn operations into no-ops.
+    # For example this chained assignment *does not change any value* of df:
+    #   `df[df["A"] > 0]["B"] = 1`
+    # Instead, use `.loc` or `.iloc` to get correct behavior:
+    #   `df.loc[df["A"] > 0, "B"] = 1`
+    # We change Pandas default from just throwing a warning (ChainedAssignmentError), to actually
+    # raising an error.
+    # Note that unfortunately we can't just set the option
+    # `mode.chained_assignment` to `"raise"` because it has no effect under
+    # Copy-on-Write, so we have to resort to promoting the warning to an error.
+    warnings.filterwarnings("error", category=pd.errors.ChainedAssignmentError)
+
+
+def process_chunk(indexed_chunk: tuple[int, pd.DataFrame], chunks_dir: Path) -> Path:
+    chunk_index, df_chunk = indexed_chunk
+
+    df_chunk = pipes.run_all(df_chunk)
+
+    return chunking.write_chunk(df_chunk, chunks_dir, chunk_index)

src/kanta/intake/__main__.py

@@ -3,6 +3,7 @@
     from datetime import date
     from pathlib import Path
 
+    from kanta import output
     from kanta.intake import assemble
     from kanta.intake import tidyup
 
@@ -34,32 +35,38 @@
         default=24,
     )
     parser.add_argument(
-        "--debug",
-        help="Increase verbosity and keep intermediate files",
-        required=False,
+        "--keep-intermediate-files",
+        help="Keep intermediate files, useful for debugging.",
         action="store_true",
     )
 
     args = parser.parse_args()
 
-    # Assemble stage
+    # Setup
     output_file_assemble_stage = (
         args.output_dir
         / f"finngen_R14_kanta_laboratory_responses.assemble-stage.{date.today()}.parquet"
     )
-    post_assemble_file = assemble.main(
-        args.source_list_file, output_file_assemble_stage
-    )
-
-    # Tidy-up stage
     output_file_tidyup_stage = (
         args.output_dir
         / f"finngen_R14_kanta_laboratory_responses_internal_1.0_{date.today()}.parquet"
     )
+    output.check_safe_write(output_file_assemble_stage)
+    output.check_safe_write(output_file_tidyup_stage)
+
+    tmp_dir = output.create_tmp_dir()
+
+    # Assemble stage
+    assemble.main(args.source_list_file, output_file_assemble_stage)
+
+    # Tidy-up stage
     tidyup.main(
         output_file_assemble_stage,
         args.phenotype_file,
         output_file_tidyup_stage,
+        tmp_dir=tmp_dir,
         partition_n_buckets=args.partition_n_buckets,
-        keep_intermediate_files=args.debug,
     )
+
+    if not args.keep_intermediate_files:
+        output.teardown_dir(tmp_dir)