Metadata-Version: 2.4
Name: echtvar-hail
Version: 0.12.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>=18
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", scale=6),  # quantized: to a stated accuracy
    Field("cadd_raw", alias="cadd", scale=2),  # quantized + renamed
    Field(
        "vep",
        fields=(  # struct: member-by-member
            Field("consequence", categorical=True),
            Field("biotype", categorical=True),
            Field("polyphen", scale=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", scale=6)` | `float64` → `decimal128(18, 6)`, reads back `0.003100` |
| `cadd_raw` | `Field("cadd_raw", alias="cadd", scale=2)` | renamed `cadd`, reads back `-1.44` |
| `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 arrive as `Decimal`s at their own scale, 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()}"
        f"  cadd={t['cadd'][i].as_py()}"
    )
    print(f"  vep={t['vep'][i].as_py()}")

# chr1:100 A>T
#   filters='PASS'  af=0.003100  cadd=3.21
#   vep=[{'consequence': 'missense_variant', 'biotype': 'protein_coding', 'polyphen': Decimal('0.910')},
#        {'consequence': 'intron_variant',   'biotype': 'lncRNA',         'polyphen': Decimal('0.050')}]
# chr1:200 AC>A
#   filters='PASS'  af=0.410000  cadd=-1.44
#   vep=[{'consequence': 'intron_variant', 'biotype': 'protein_coding', 'polyphen': Decimal('0.010')}]

# No division and nothing to look up: the values ARE Decimals, at the scale
# the schema records.
```

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, partitions, sorts | picks encodings, casts quantized columns to `DECIMAL`, narrows the key to `uint32`, 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", scale=6)  # float -> DECIMAL, to a stated accuracy
Field("consequence", categorical=True)  # dictionary-encoded (echtvar's string table)
Field("cadd_raw", alias="cadd", scale=2)  # 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", scale=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.
* **`scale` applies per member**, and the DECIMAL cast reaches inside —
  `csq.polyphen` at `scale=3` becomes `decimal128(18, 3)` while its
  sibling members are untouched.

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", scale=3))),
        Field("n"),
    ),
)
```

```
a.b.c       BYTE_ARRAY  RLE_DICTIONARY
a.b.score   INT64       PLAIN            # DECIMAL(18,3), reads back as 0.125
a.n         INT32       PLAIN

# arrow type: struct<b: struct<c: string, score: decimal128(18, 3)>, n: int32>
# the nested leaf a.b.score is quantized; the scale is in the schema
```

`scale` is an **accuracy budget**, not a performance knob — it is Parquet's
DECIMAL scale, a count of decimal places, passed straight through. 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:

| scale | size | vs float32 | distinct values |
|---|---|---|---|
| 8 | 753KB | 1.45× | 112,963 |
| 7 | 628KB | 1.74× | 58,461 |
| 6 | 520KB | 2.10× | 12,943 |
| 5 | 362KB | 3.02× | 1,338 |
| 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, and it applies no encoding — `scale` and `categorical` are both
acted on by `finalize_parquet`, so staging carries values as they were, and a
missing value is a Parquet null. There is no pre-flight aggregation over your
data; `Field` rejects an impossible `scale` at construction, and the DECIMAL
cast in stage 2 raises on a value that will not fit. 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          16k-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 ~16,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 | 1.01× | 0.6 ms | 13.4 ms | 0.68 |
| **16,000** (current) | **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.

Confirmed on a whole real chromosome rather than a slice — chr21, 11.7M rows,
rewritten both ways, reading the key column as numpy:

| rows/row group | row groups | file | footer | parse | lookup |
|---|---|---|---|---|---|
| **16,000** (current) | 749 | 117.9 MB | 0.99 MB | 3.59 ms | **0.58 ms** |
| 64,000 | 208 | 116.6 MB | 0.29 MB | 1.06 ms | 2.07 ms |

3.6× faster lookups for 3.4× more footer and 1.1% more bytes. **This is tuned
for locally cached reads**, where the bytes are nearly free and a footer cache
miss costs a few milliseconds off local disk.

**Over a network the tradeoff inverts.** 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, so remotely 16,000 is only ~5% faster than 64,000 (501 ms vs 528 ms)
while costing 4× the footer — and behind a CDN that footer is fetched before any
lookup can start, by every cold client. **If these move to a CDN, 64,000 is the
constant to return to.**

One caveat that dominates both regimes: **do not read the key column with
`to_pylist()`.** It builds one Python int per row under the GIL — 5.6× on a
single lookup, and it serialized concurrent lookups across files (10 threads
gave 1.15×). Read `.to_numpy()`, which is zero-copy for `uint32`, and keep the
row in Arrow; that gave 4.8× from concurrency instead.

`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.

**The files carry no metadata of ours.** A quantized column is a Parquet
DECIMAL, so its scale is in the schema; a categorical one announces itself
through its column-chunk encoding; a missing value is a Parquet null. Nothing
has to know what wrote the file to read it correctly:

```python
schema = pq.ParquetFile("/tmp/echtvar/short/chr1.parquet").schema_arrow
schema.field("af").type  # decimal128(18, 6)  -> reads back as 0.010000
schema.field("cadd").type  # list<decimal128(18, 2)>
```

`store_decimal_as_integer` keeps the physical column an `INT64`, so being
self-describing costs nothing on disk.

## 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.
