Add split sample feature

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{<field>: List[T], ...}`` where each
+    list holds one value per sample in declaration order.  This function converts
+    that layout into individual columns named ``<SAMPLE>_<FIELD>``
+    (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<GT: list, DP: list, ...>`) 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 `<SAMPLE>_<FIELD>` (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<GT: list, DP: list, ...>`) 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 `<SAMPLE>_<FIELD>` (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,14 +626,21 @@ 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,
             projection_pushdown,
             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(

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())