Metadata-Version: 2.4
Name: echtvar-hail
Version: 0.7.0
Summary: Pure-Python and Hail-expression implementations of the echtvar var32 and kmer16 variant encodings
Project-URL: Homepage, https://github.com/broadinstitute/echtvar-hail
Project-URL: Repository, https://github.com/broadinstitute/echtvar-hail
Project-URL: Issues, https://github.com/broadinstitute/echtvar-hail/issues
Author-email: Ben Blankenmeister <bblanken@broad.mit.edu>
License: MIT
Keywords: echtvar,encoding,genomics,hail,variant
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
Requires-Python: <3.13,>=3.11
Requires-Dist: hail==0.2.138
Requires-Dist: pyarrow>=14
Description-Content-Type: text/markdown

# echtvar-hail

[![Tests](https://github.com/broadinstitute/echtvar-hail/actions/workflows/test.yml/badge.svg)](https://github.com/broadinstitute/echtvar-hail/actions/workflows/test.yml)
[![PyPI](https://img.shields.io/pypi/v/echtvar-hail.svg)](https://pypi.org/project/echtvar-hail/)
[![Python](https://img.shields.io/badge/python-3.11%20%7C%203.12-blue.svg)](https://www.python.org/downloads/)

[echtvar](https://github.com/brentp/echtvar)'s variant encodings in Hail, and an
export that writes them to Parquet instead of echtvar's zip container.

Two things live here:

* **`var32` / `kmer16`** — Python ports of `src/lib/{var32,kmer16}.rs`, each with
  an `hl_*` twin so the encodings run inside a Hail pipeline rather than by
  collecting rows to Python.
* **`parquet`** — an export in echtvar's per-`(chrom, chunk)` shape, as standard
  Parquet. The archive becomes a self-describing columnar file that DuckDB,
  Polars, pandas or Arrow can read without echtvar.

## Install

```bash
pip install echtvar-hail
```

Python 3.11 or 3.12. Hail brings PySpark, which needs a JVM — Hail is built and tested
against **Java 11**.

## Quick start

One Hail table, one field config, covering every shape `Field` supports —
copied as-is, quantized, renamed, and a nested struct — plus the routing that
splits variants across the two output tables.

```python
import hail as hl
import pyarrow.parquet as pq
from echtvar_hail import var32
from echtvar_hail.parquet import Field, export_parquet, finalize_parquet

hl.init()

ht = hl.Table.parallelize(
    [
        {"locus": hl.Locus("chr1", 100, "GRCh38"), "alleles": ["A", "T"],
         "filters": "PASS", "af": 0.0031, "cadd_raw": 3.21,
         "vep": [hl.Struct(consequence="missense_variant", biotype="protein_coding", polyphen=0.91),
                 hl.Struct(consequence="intron_variant", biotype="lncRNA", polyphen=0.05)]},
        {"locus": hl.Locus("chr1", 200, "GRCh38"), "alleles": ["AC", "A"],
         "filters": "PASS", "af": 0.41, "cadd_raw": -1.44,
         "vep": [hl.Struct(consequence="intron_variant", biotype="protein_coding", polyphen=0.01)]},
        {"locus": hl.Locus("chr1", 300, "GRCh38"), "alleles": ["A", "ACGTACGT"],
         "filters": "low_qual", "af": 0.0001, "cadd_raw": 22.6,
         "vep": [hl.Struct(consequence="stop_gained", biotype="protein_coding", polyphen=1.0)]},
    ],
    hl.tstruct(
        locus=hl.tlocus("GRCh38"), alleles=hl.tarray(hl.tstr),
        filters=hl.tstr, af=hl.tfloat64, cadd_raw=hl.tfloat64,
        vep=hl.tarray(hl.tstruct(consequence=hl.tstr, biotype=hl.tstr, polyphen=hl.tfloat64)),
    ),
)

fields = [
    Field("filters"),                                  # plain: copied as-is
    Field("af", precision=1e-6),                        # quantized: to a stated accuracy
    Field("cadd_raw", alias="cadd", precision=0.01),    # quantized + renamed
    Field("vep", fields=(                               # struct: member-by-member
        Field("consequence", categorical=True),
        Field("biotype", categorical=True),
        Field("polyphen", precision=1e-3),
    )),
]
```

Every Hail column not named in `fields` is dropped — `locus`/`alleles` become
the key, and anything else absent from this list simply doesn't reach the
file. What each `Field` does to its column:

| Hail column | `Field` | on disk |
|---|---|---|
| `filters` | `Field("filters")` | copied through unchanged |
| `af` | `Field("af", precision=1e-6)` | `float64` → `int64`, `0.0031` → `3100` |
| `cadd_raw` | `Field("cadd_raw", alias="cadd", precision=0.01)` | renamed `cadd`, `-1.44` → `-144` |
| `vep` | `Field("vep", fields=(...))` | stays `array<struct<...>>`, each member its own leaf |

```python
# Stage 1 (Hail + Spark, distributed): route short/long, encode, partition.
short, long = export_parquet(ht, "/tmp/staging", fields=fields)

# Stage 2 (pyarrow, one node): choose encodings, consolidate.
finalize_parquet(short, "/tmp/echtvar/short", fields=fields)
finalize_parquet(long, "/tmp/echtvar/long", fields=fields)
```

Reading it back — a key plus its `chrom`/`chunk` is a whole variant, quantized
values divide back out by their multiplier, and the struct comes back as a
list of dicts:

```python
t = pq.read_table("/tmp/echtvar/short/chr1.parquet")
print(t.schema.names)
# ['key', 'filters', 'af', 'cadd', 'vep', 'chrom', 'chunk']

for i in range(t.num_rows):
    d = var32.decode(t["key"][i].as_py(), t["chrom"][i].as_py(), t["chunk"][i].as_py())
    print(f"{d.chrom}:{d.position + 1} {d.reference}>{d.alternate}")
    print(f"  filters={t['filters'][i].as_py()!r}  af={t['af'][i].as_py() / 1_000_000}"
          f"  cadd={t['cadd'][i].as_py() / 100}")
    print(f"  vep={t['vep'][i].as_py()}")

# chr1:100 A>T
#   filters='PASS'  af=0.0031  cadd=3.21
#   vep=[{'consequence': 'missense_variant', 'biotype': 'protein_coding', 'polyphen': 910},
#        {'consequence': 'intron_variant',   'biotype': 'lncRNA',         'polyphen': 50}]
# chr1:200 AC>A
#   filters='PASS'  af=0.41  cadd=-1.44
#   vep=[{'consequence': 'intron_variant', 'biotype': 'protein_coding', 'polyphen': 10}]
```

The third variant is missing from that output because `A>ACGTACGT` exceeds
var32's four-base budget, so it routed to the long table instead:

```python
t = pq.read_table("/tmp/echtvar/long/chr1.parquet")
print(t.schema.field("key").type, t.num_rows)
# list<element: uint32> 1
```

## Why two stages

No single tool can write this file. Hail and Spark are needed for the encoding
and the shuffle, at scale. But Spark writes Parquet through **parquet-mr**,
which cannot be told a column's encoding at all — it infers one from a writer
version and a dictionary setting. The choices that make the format worth having
are unavailable there.

| | stage 1 — `export_parquet` | stage 2 — `finalize_parquet` |
|---|---|---|
| engine | Hail + Spark | pyarrow |
| runs | distributed | one node, streaming a chunk at a time |
| does | routes short/long, encodes keys, applies value transforms, partitions, sorts | picks encodings, narrows the key to `uint32`, writes the footer config, consolidates |
| output | `staging/short/chrom=chr1/chunk=0/…` | `out/short/chr1.parquet` |

Stage 2 output is **one file per chromosome, several row groups per chunk**, so a
chunk lookup is a row-group read rather than opening a file. Measured over
400 chunks: same bytes to within 1%, and 2.9× faster on random chunk reads
locally — the gap widens on object storage, where opening a file is a round
trip.

## API

### `Field`

How one annotation is carried. Three shapes, and any of them may be
array-valued — VEP-style annotations usually are, and an array stays an array
rather than being flattened or joined into a delimited string.

```python
Field("af")                              # carried as-is
Field("af", precision=1e-6)              # float -> int, to a stated accuracy
Field("consequence", categorical=True)   # dictionary-encoded (echtvar's string table)
Field("cadd_raw", alias="cadd", precision=0.01)   # renamed on the way out
```

#### Nested fields

The `vep` field above is the common case — nest `fields` to describe a
struct's members, an array of them being one per transcript. Here's why it's
shaped that way:

```python
Field("csq", fields=(
    Field("consequence", categorical=True),
    Field("biotype", alias="tx_biotype", categorical=True),
    Field("polyphen", precision=1e-3),
))
```

Given `[hl.struct(consequence="missense_variant", biotype="protein_coding", polyphen=0.912), ...]`
that writes one Parquet leaf per member:

```
csq.list.element.consequence   BYTE_ARRAY  RLE_DICTIONARY
csq.list.element.tx_biotype    BYTE_ARRAY  RLE_DICTIONARY
csq.list.element.polyphen      INT64       PLAIN            # 0.912 -> 912
```

Three things follow from each member being its own column:

* **Each takes its own encoding.** The two string members get their own
  dictionaries; the quantized one correctly gets none.
* **`fields` is an allowlist.** Members you don't name are dropped, which is
  how you take three columns out of a VEP annotation with forty.
* **`precision` applies per member**, and the export's sentinel check reaches
  inside:

  ```python
  export_parquet(ht, staging, fields=fields)
  # [{'column': 'csq.polyphen', 'multiplier': 1000, 'min': 0.05, 'max': 1.0, ...}]
  ```

The struct survives to disk as a struct — nothing is flattened or exploded, so
a row reads back as a list of dicts:

```python
t["csq"][0].as_py()
# [{'consequence': 'missense_variant', 'tx_biotype': 'protein_coding', 'polyphen': 912},
#  {'consequence': 'intron_variant',   'tx_biotype': 'lncRNA',         'polyphen': 50}]
```

The array is incidental. A plain struct nests the same way, to any depth —
`fields` describes the shape, and whether it sits inside a list only changes
the leaf paths:

```python
Field("a", fields=(
    Field("b", fields=(Field("c", categorical=True), Field("score", precision=1e-3))),
    Field("n"),
))
```

```
a.b.c       BYTE_ARRAY  RLE_DICTIONARY
a.b.score   INT64       PLAIN            # 0.125 -> 125
a.n         INT32       PLAIN

# arrow type: struct<b: struct<c: string, score: int64>, n: int32>
# the nested leaf a.b.score is quantized with multiplier 1000
```

`precision` is an **accuracy budget**, not a performance knob. Quantization
shrinks a column by making values *repeat*, so what you keep is what you don't
save. On 300k gnomAD-like allele frequencies, against 1.09MB as float32:

| precision | size | vs float32 | distinct values |
|---|---|---|---|
| 1e-8 | 753KB | 1.45× | 112,963 |
| 1e-7 | 628KB | 1.74× | 58,461 |
| 1e-6 | 520KB | 2.10× | 12,943 |
| 1e-5 | 362KB | 3.02× | 1,338 |
| 1e-4 | 230KB | 4.75× | 135 |

An allele frequency is `AC/AN`, so it cannot mean anything below `1/AN` — about
6.7e-7 at AN≈1.5M. Asking for finer spends bytes on the denominator wobbling
between sites.

### `export_parquet(ht, base_path, *, locus="locus", alleles="alleles", fields=(), mode="overwrite")`

Stage 1. Returns `(short_path, long_path)`. Before writing, it refuses any
`precision` whose observed range would reach the missing sentinel — that failure
is otherwise silent, since real values would simply read back as absent. One
aggregation pass, and free when no field is quantized. Routes each variant on
`len(ref) + len(alt) > 4`: short variants take a 32-bit `var32` key, long ones a
`kmer16` array. Multi-allelic rows contribute only their first alternate — split
them (`hl.split_multi_hts`) first.

### `finalize_parquet(staging_path, output_path, *, fields=(), filesystem=None)`

Stage 2. Pass the *same* `fields`. On GCS, hand it a `pyarrow.fs.GcsFileSystem`
— pyarrow does not use Hadoop's GCS connector.

### Decoding

`var32.decode(key, chrom, chunk)` and `kmer16.decode(key, chrom)` return a
`PRA(chrom, chunk, position, reference, alternate)`, 0-based. A var32 key stores
position truncated to 20 bits, so it needs its `chunk` to reconstruct an
absolute coordinate — which is why `chunk` is a column.

## Output layout

```
out/short/chr1.parquet      key: uint32          64k-row row groups, nested in chunks
out/long/chr1.parquet       key: list<uint32>
```

Columns are `key`, `chrom`, `chunk`, and one per field. Position and alleles are
not stored — the key plus `chunk` reconstructs them.

Row groups hold ~64,000 rows and **nest inside chunks** — a chunk becomes
several of them, and none ever spans a chunk boundary. A lookup is two stages,
and the order matters:

```python
f = pq.ParquetFile("out/short/chr1.parquet")
md = f.metadata
ci, ki = md.schema.names.index("chunk"), md.schema.names.index("key")
want, key = chunk_of(pos0), var32.encode(pos0, ref, alt)

def holds(i):
    rg = md.row_group(i)
    k = rg.column(ki).statistics
    return rg.column(ci).statistics.min == want and k.min <= key <= k.max

rg = next(i for i in range(md.num_row_groups) if holds(i))   # footer only, no I/O
keys = f.read_row_group(rg, columns=["key"])["key"].to_pylist()
i = bisect.bisect_left(keys, key)
```

**Stage 1 cannot be skipped.** var32 keys overlap across chunks by construction
(position is truncated to `POSITION_BITS`), so matching on key statistics alone
lands in another chunk's row group. Within a chunk the rows stay key-ordered, so
its row groups have disjoint ascending key ranges and exactly one can hold the
key. And never index positionally — a chunk with no variants produces no row
group at all.

Row-group size is what sets lookup cost, because pyarrow decodes a whole row
group either way (it prunes on the page index for neither `read_row_group()` nor
a filtered scan). Measured **locally** on 2.5M real gnomAD chr1 rows:

| rows/row group | file size | footer | lookup | MB decoded |
|---|---|---|---|---|
| 312,746 (a whole chunk) | 1.00× | 0.5 ms | 89.8 ms | 4.34 |
| 64,000 (current) | 1.01× | 0.6 ms | 13.4 ms | 0.68 |
| 16,000 | 1.02× | 1.1 ms | 3.5 ms | 0.16 |
| 4,000 | 1.11× | 3.6 ms | 1.6 ms | 0.05 |
| 1,000 | 1.36× | 12.1 ms | 2.0 ms | 0.01 |

Below 16,000 compression collapses — encodings and zstd's window reset per row
group — and footer parsing starts to dominate.

**Over a network the tradeoff inverts**, and that is what 64,000 is chosen for.
A remote read has a fixed round-trip cost that shrinking a row group cannot
touch, and it swamps the decode. Fitted on real gnomAD chr1 (68.2M rows) read
from GCS: the key column costs 208 ms + 0.319 ms per 1k rows, all 11 columns
284 ms + 0.24 ms per 1k. A lookup is both reads:

| rows/row group | remote lookup | chr1 footer | chr1 row groups | genome index |
|---|---|---|---|---|
| 310,000 (a whole chunk) | 665 ms | 0.29 MB | 220 | 22 KiB |
| **64,000** (current) | **528 ms** | **1.4 MB** | **1,093** | **108 KiB** |
| 16,000 | 501 ms | 5.7 MB | 4,372 | 437 KiB |

64,000 takes 95% of the available speedup for a quarter of the footer. The index
addressing those row groups is 8 bytes each and trivial either way; the footer is
the real cost, paid on every cold open, and it grows ~20× going from whole-chunk
to 16,000. Reading these files off local disk instead? Prefer 16,000 — the local
table shows it 3.8× faster there, with no round trip to hide the decode.

`chrom` identifies rows once several
files are read together (`ds.dataset(root)` concatenates them and does not
surface filenames). Both are constant over their scope, so RLE reduces them to
almost nothing.

Quantized columns are unreadable without their multiplier, so it is written to
the Parquet footer under the `fields` key:

```python
import json
md = pq.ParquetFile("/tmp/echtvar/short/chr1.parquet").metadata
json.loads(md.metadata[b"fields"])
# [{'name': 'filters', 'column': 'filters'},
#  {'name': 'af', 'column': 'af', 'multiplier': 1000000, 'missing': 4294967295},
#  {'name': 'cadd_raw', 'column': 'cadd', 'multiplier': 100, 'missing': 4294967295},
#  {'name': 'vep', 'column': 'vep', 'fields': [
#      {'name': 'consequence', 'column': 'consequence'},
#      {'name': 'biotype', 'column': 'biotype'},
#      {'name': 'polyphen', 'column': 'polyphen', 'multiplier': 1000, 'missing': 4294967295}]}]
```

A struct field's own entry has no `multiplier` — it's a container, so the
multiplier sits on the member that needs one, and the nesting in the config is
mirrored in the footer.

Dictionary encoding is deliberately *not* recorded there — Parquet already
reports each column chunk's encoding, and a footer note would be the less
reliable of the two.

## Development

```bash
uv sync            # dev tools are a PEP 735 dependency group
uv run pytest
uv run ruff check
```

The suite starts a real Hail/Spark session; it needs a JVM.
