diff --git a/polars_bio/io.py b/polars_bio/io.py index d1cdeb33..13bbda26 100644 --- a/polars_bio/io.py +++ b/polars_bio/io.py @@ -76,6 +76,95 @@ _VALID_SAM_TYPE_CODES = _VALID_SAM_SCALAR_TYPE_CODES | {"B"} +def _split_genotypes_to_per_sample( + lf: pl.LazyFrame, sample_names: list[str] +) -> pl.LazyFrame: + """Explode the multi-sample `genotypes` struct into flat per-sample columns. + + The ``genotypes`` column is a ``Struct{: List[T], ...}`` where each + list holds one value per sample in declaration order. This function converts + that layout into individual columns named ``_`` + (e.g. ``TUMOR_GT``, ``TUMOR_DP``, ``NORMAL_GT``, …) and drops the original + ``genotypes`` column. + + If the LazyFrame has no ``genotypes`` column (single-sample VCF) the frame + is returned unchanged. + + Args: + lf: LazyFrame produced by :func:`scan_vcf`. + sample_names: Ordered list of sample names matching the list positions + inside ``genotypes`` fields. Must equal the samples in the VCF + (or the subset requested via the ``samples`` parameter). + + Returns: + LazyFrame with per-sample flat columns replacing ``genotypes``. + """ + schema = lf.collect_schema() + if "genotypes" not in schema.names(): + return lf + + genotypes_dtype = schema["genotypes"] + if not isinstance(genotypes_dtype, pl.Struct): + return lf + + format_fields = [f.name for f in genotypes_dtype.fields] + exprs = [ + pl.col("genotypes").struct.field(fmt).list.get(i).alias(f"{sample}_{fmt}") + for i, sample in enumerate(sample_names) + for fmt in format_fields + ] + return lf.with_columns(exprs).drop("genotypes") + + +def _resolve_sample_names_for_split( + lf: pl.LazyFrame, samples: list[str] | None, path: str | None = None +) -> list[str]: + """Return the ordered sample names to use when splitting genotypes. + + Prefers the caller-supplied *samples* list. Falls back to reading the + ``#CHROM`` header line from a local VCF file to extract sample names. + + Args: + lf: LazyFrame from ``scan_vcf`` (used to verify genotypes column exists). + samples: The ``samples`` argument forwarded from ``read_vcf``/``scan_vcf``. + path: Path to the VCF file for header-based fallback. + + Returns: + Ordered list of sample names. + + Raises: + ValueError: If sample names cannot be determined. + """ + if samples is not None: + return list(samples) + + # Fallback: parse #CHROM line from local VCF header (lightweight, header-only read) + if path is not None and not path.startswith(("gs://", "s3://", "az://", "http")): + try: + import gzip + + opener = ( + gzip.open + if path.endswith(".gz") or path.endswith(".bgz") + else open + ) + with opener(path, "rt") as fh: + for line in fh: + if line.startswith("#CHROM"): + parts = line.strip().split("\t") + # Standard VCF columns: CHROM POS ID REF ALT QUAL FILTER INFO FORMAT … + if len(parts) > 9: + return parts[9:] + break + except Exception: + pass + + raise ValueError( + "split_samples=True requires sample names, but they could not be inferred. " + "Pass samples=[...] explicitly." + ) + + def _supported_sam_type_message() -> str: return ( "Supported scalar types: A, c, C, s, S, i, I, f, Z, H. " @@ -358,6 +447,7 @@ def read_vcf( predicate_pushdown: bool = True, use_zero_based: Optional[bool] = None, samples: Union[list[str], None] = None, + split_samples: bool = False, genotype_encoding_raw: bool = True, ) -> pl.DataFrame: """ @@ -374,6 +464,7 @@ def read_vcf( info_fields: List of INFO field names to include. If *None*, all INFO fields from the VCF header are included by default. Use this to limit fields for better performance. format_fields: List of FORMAT field names to include (per-sample genotype data). If *None*, all FORMAT fields are included by default. For **single-sample** VCFs, FORMAT fields are top-level columns (e.g., `GT`, `DP`). For **multi-sample** VCFs, FORMAT data is exposed as a nested `genotypes` column (`struct`) with sample names in `meta["header"]["sample_names"]`. samples: Optional list of sample names to include from the VCF header. Matching is exact and case-sensitive. Missing sample names are skipped with a warning. The output follows the requested sample order. + split_samples: If *True* and the VCF is multi-sample, explode the nested `genotypes` struct into individual flat columns named `_` (e.g. `NA12878_GT`, `NA12878_DP`). Has no effect on single-sample VCFs. chunk_size: The size in MB of a chunk when reading from an object store. The default is 8 MB. For large scale operations, it is recommended to increase this value to 64. concurrent_fetches: [GCS] The number of concurrent fetches when reading from an object store. The default is 1. For large scale operations, it is recommended to increase this value to 8 or even more. allow_anonymous: [GCS, AWS S3] Whether to allow anonymous access to object storage. @@ -411,9 +502,14 @@ def read_vcf( # │ 1 ┆ 10015 ┆ A ┆ . ┆ null ┆ 0/0 ┆ 17 ┆ 35 │ # └───────┴───────┴─────┴─────┴──────┴─────┴─────┴─────┘ - # Multi-sample VCF: FORMAT data is nested in "genotypes" - df = pb.read_vcf("multisample.vcf", format_fields=["GT", "DP"]) - print(df.select(["chrom", "start", "genotypes"])) + # Multi-sample VCF: split into per-sample columns + df = pb.read_vcf( + "multisample.vcf", + samples=["TUMOR", "NORMAL"], + format_fields=["GT", "DP"], + split_samples=True, + ) + print(df.select(["chrom", "start", "TUMOR_GT", "TUMOR_DP", "NORMAL_GT", "NORMAL_DP"])) ``` """ lf = IOOperations.scan_vcf( @@ -431,6 +527,7 @@ def read_vcf( predicate_pushdown=predicate_pushdown, use_zero_based=use_zero_based, samples=samples, + split_samples=split_samples, ) # Get metadata before collecting (polars-config-meta doesn't preserve through collect) zero_based = lf.config_meta.get_metadata().get("coordinate_system_zero_based") @@ -456,6 +553,7 @@ def scan_vcf( predicate_pushdown: bool = True, use_zero_based: Optional[bool] = None, samples: Union[list[str], None] = None, + split_samples: bool = False, ) -> pl.LazyFrame: """ Lazily read a VCF file into a LazyFrame. @@ -471,6 +569,7 @@ def scan_vcf( info_fields: List of INFO field names to include. If *None*, all INFO fields from the VCF header are included by default. Use this to limit fields for better performance. format_fields: List of FORMAT field names to include (per-sample genotype data). If *None*, all FORMAT fields are included by default. For **single-sample** VCFs, FORMAT fields are top-level columns (e.g., `GT`, `DP`). For **multi-sample** VCFs, FORMAT data is exposed as a nested `genotypes` column (`struct`) with sample names in `meta["header"]["sample_names"]`. samples: Optional list of sample names to include from the VCF header. Matching is exact and case-sensitive. Missing sample names are skipped with a warning. The output follows the requested sample order. + split_samples: If *True* and the VCF is multi-sample, explode the nested `genotypes` struct into individual flat columns named `_` (e.g. `NA12878_GT`, `NA12878_DP`). Has no effect on single-sample VCFs. chunk_size: The size in MB of a chunk when reading from an object store. The default is 8 MB. For large scale operations, it is recommended to increase this value to 64. concurrent_fetches: [GCS] The number of concurrent fetches when reading from an object store. The default is 1. For large scale operations, it is recommended to increase this value to 8 or even more. allow_anonymous: [GCS, AWS S3] Whether to allow anonymous access to object storage. @@ -527,7 +626,7 @@ def scan_vcf( zero_based=zero_based, ) read_options = ReadOptions(vcf_read_options=vcf_read_options) - return _read_file( + lf = _read_file( path, InputFormat.Vcf, read_options, @@ -535,6 +634,13 @@ def scan_vcf( predicate_pushdown, zero_based=zero_based, ) + if split_samples: + # Only apply if a "genotypes" struct column is actually present + # (multi-sample). Single-sample VCFs have no such column. + if "genotypes" in lf.collect_schema().names(): + resolved_samples = _resolve_sample_names_for_split(lf, samples, path) + lf = _split_genotypes_to_per_sample(lf, resolved_samples) + return lf @staticmethod def read_vcf_zarr( diff --git a/tests/test_io_vcf.py b/tests/test_io_vcf.py index 172c6b57..c38dca2b 100644 --- a/tests/test_io_vcf.py +++ b/tests/test_io_vcf.py @@ -238,3 +238,71 @@ def test_sink_vcf_bgz(self, tmp_path): df_back = pb.read_vcf(out_path) assert len(df_back) == 2 + + +class TestSplitSamples: + """Tests for read_vcf / scan_vcf split_samples feature (issue #394).""" + + vcf = f"{DATA_DIR}/io/vcf/multisample.vcf" + single = f"{DATA_DIR}/io/vcf/single_sample_collision.vcf" + + def test_split_samples_explicit_produces_flat_columns(self): + df = pb.read_vcf( + self.vcf, + samples=["NA12878", "NA12879", "NA12880"], + split_samples=True, + ) + assert "genotypes" not in df.columns + for sample in ("NA12878", "NA12879", "NA12880"): + for field in ("GT", "DP", "GQ"): + assert f"{sample}_{field}" in df.columns + + def test_split_samples_auto_detect_all_samples(self): + """Without explicit samples=, names are read from the VCF header.""" + df = pb.read_vcf(self.vcf, split_samples=True) + assert "genotypes" not in df.columns + assert "NA12878_GT" in df.columns + assert "NA12879_DP" in df.columns + assert "NA12880_GQ" in df.columns + + def test_split_samples_subset(self): + df = pb.read_vcf(self.vcf, samples=["NA12878", "NA12880"], split_samples=True) + assert "NA12878_GT" in df.columns + assert "NA12880_DP" in df.columns + assert "NA12879_GT" not in df.columns + + def test_split_samples_values_correct(self): + df = pb.read_vcf( + self.vcf, + samples=["NA12878", "NA12879", "NA12880"], + format_fields=["GT", "DP"], + split_samples=True, + ) + row = df.row(0, named=True) + assert row["NA12878_GT"] == "0/1" + assert row["NA12878_DP"] == 25 + assert row["NA12879_GT"] == "1/1" + assert row["NA12879_DP"] == 30 + assert row["NA12880_GT"] == "0/0" + assert row["NA12880_DP"] == 20 + + def test_split_samples_default_false_unchanged(self): + df = pb.read_vcf(self.vcf) + assert "genotypes" in df.columns + assert "NA12878_GT" not in df.columns + + def test_split_samples_single_sample_noop(self): + """split_samples=True on a single-sample VCF is a no-op.""" + df = pb.read_vcf(self.single, split_samples=True) + assert "genotypes" not in df.columns + assert "GT" in df.columns + + def test_scan_vcf_split_samples_lazy_filter(self): + """split_samples works lazily; per-sample columns usable in filter.""" + import polars as pl + + lf = pb.scan_vcf(self.vcf, split_samples=True) + assert "NA12878_GT" in lf.collect_schema().names() + df = lf.filter(pl.col("NA12878_DP") > 20).collect() + assert len(df) == 2 + assert all(dp > 20 for dp in df["NA12878_DP"].to_list())