r"""Join ESM protein-atlas SAE embeddings to folded structures, reading only the columns
needed, and write (protein_hash, max_pooled, structure_blob) parquet.
Both sources are columnar, and each is dominated by one very large column we do not want
(`per_position` on the SAE side, `pae` on the folds side), so reading only the needed columns
skips most of the source bytes. See IMPLEMENTATION NOTES.
-------------------------------------------------------------------------------
WHAT IT DOES
-------------------------------------------------------------------------------
One invocation processes one batch of SAE parquet shards: it reads their protein_hash and
max_pooled, fetches the matching structures from the folds Lance datasets, inner-joins them,
and writes a single parquet file. Batches are independent and can all run in parallel.
output columns: protein_hash (string), max_pooled (binary), structure_blob (binary)
output file: <out-uri>/batch_<id:05d>.parquet
The join keeps only proteins that have a predicted structure. Roughly 1 in 6 SAE proteins
has one, so the output has far fewer rows than the SAE input.
-------------------------------------------------------------------------------
REQUIREMENTS
-------------------------------------------------------------------------------
pip install duckdb pylance pyarrow
Reads need NO AWS credentials -- the source bucket is public (AWS Open Data). You need write
credentials only for your own --out-uri. Pass --signed-reads to sign reads anyway.
-------------------------------------------------------------------------------
USAGE
-------------------------------------------------------------------------------
python extract_sae_folds_join.py --batch-id N --batch-size 16 --out-uri s3://<you>/out
Batch ids 0..4095 at --batch-size 16 cover all 65,536 shards. Run them however you like:
# SLURM array
sbatch --array=0-4095%64 -c 12 --mem 32G -p <partition> -t 240 -J sae_folds \
--wrap "python extract_sae_folds_join.py \
--batch-id \$SLURM_ARRAY_TASK_ID --batch-size 16 \
--out-uri s3://<you>/out"
# no scheduler, 64-way local parallelism
seq 0 4095 | xargs -P 64 -I{} python extract_sae_folds_join.py \
--batch-id {} --batch-size 16 --out-uri s3://<you>/out
Each task is idempotent: it skips a batch whose output already exists, and writes via a .tmp
plus rename, so interrupted or preempted tasks can simply be re-run.
-------------------------------------------------------------------------------
RUN IT IN us-west-2
-------------------------------------------------------------------------------
The data lives in us-west-2. This is about time, not money -- it is an AWS Open Data bucket,
so AWS covers egress. In-region a full run takes hours; from outside the region you are
capped by internet bandwidth, which stretches the full job to months.
-------------------------------------------------------------------------------
RESOURCES PER TASK
-------------------------------------------------------------------------------
RAM 32 GB request is comfortable
CPU 12 cores is plenty
Memory is bounded because the script streams one shard at a time. Do not change it to
materialise a whole batch: that holds every shard's max_pooled in memory at once and
multiplies resident memory by the batch size.
-------------------------------------------------------------------------------
HOW LONG IT TAKES
-------------------------------------------------------------------------------
Calibrate: run batch 0 alone, take its wall time T, then
total_hours ~ 4096 * T_seconds / concurrency / 3600
S3 request rate, not bandwidth, is likely the real ceiling: all shards share one prefix, and
AWS guidance is ~5,500 GET/s per partitioned prefix, so past some point more concurrency
does not help. Early `503 SlowDown` is normal while S3 auto-partitions -- back off and retry,
do not treat it as fatal. Raising --batch-size does not reduce the total number of GETs.
-------------------------------------------------------------------------------
COVERAGE: both structure datasets are read
-------------------------------------------------------------------------------
There are two structure datasets under v1/folds/ and this script reads BOTH, unioning them
for each batch range: folds_1B.lance and the much smaller folds_atlas.lance.
They are disjoint, not nested -- sampled folds_atlas protein_hashes from separated prefixes
had no matches in folds_1B -- so together they are the complete structure set and the union
needs no dedup. Their projected schemas are identical (protein_hash: string, structure_blob:
binary). Including folds_atlas adds under 1% more bytes and rows.
Reading both in one pass matters: the expensive side is the SAE read, so making a second
pass with a different structure dataset would repeat all of it to pick up little data.
The per-batch log prints each dataset's contribution separately, then the union total.
===============================================================================
IMPLEMENTATION NOTES -- background, not needed to run the script
===============================================================================
WHERE THE BYTES ARE
-------------------
DuckDB and Lance fetch only the byte ranges for the columns you select (confirmed on the
wire: DuckDB issues HTTP range GETs and receives `206 Partial Content`), so projection
really does reduce bytes transferred rather than downloading whole objects.
Approximate share of bytes by column:
s3://esm-protein-atlas/v1/sae/data_shards/
per_position ~89% <-- not wanted
max_pooled ~11%
header, protein_hash <1%
s3://esm-protein-atlas/v1/folds/folds_1B.lance
pae ~80% <-- not wanted
structure_blob ~19%
all other columns ~1%
WHY WE SELECT `header` AND NOT `protein_hash`
---------------------------------------------
Purely to halve the request count; it saves no bytes. A Parquet column is not one contiguous
run -- each shard has many row groups, each storing all four chunks back to back:
[RG0: header | max_pooled | per_position | protein_hash][RG1: ...] ... [footer]
So `max_pooled` is one byte range per row group. `header` is adjacent to it (one coalesced
range GET per row group), while `protein_hash` sits past the `per_position` chunk we do not
want, so it cannot be coalesced and costs a second GET per row group.
`header` is formatted `{protein_hash}|{id}|{source}`, and `split_part(header,'|',1) ==
protein_hash` was verified for every row of full shards. Selecting `protein_hash` directly
is fine too -- it just costs ~10-15% wall time. The load-bearing decision is dropping
`per_position`, not this.
DATA LAYOUT: ordering, sharding, and one trap
----------------------------------------------
* SAE rows are GLOBALLY SORTED by protein_hash. Global order is bucket char in
"0123456789abcdef", then the 0000-4095 index (verified: 0_4095 ends 0ffffff0..., 1_0000
starts 10000006...; 9_4095 ends 9ffffff7..., a_0000 starts a0000001...).
* Lance fragments are RANGE-PARTITIONED on protein_hash, ascending and disjoint. Rows are
not sorted *within* a fragment, which does not matter for a range join.
* Both Lance datasets carry BTree scalar indices on BOTH `protein_hash` and `header`.
* Because both sides are hash-ordered, this is a partitioned merge join: no shuffle, bounded
memory, 4096 independent work units.
TRAP: shards are equal-row-count splits, so shard boundaries do NOT align with hash-prefix
boundaries -- e.g. a_0095 spans a05->a06 and a_0111 spans a06->a07. NEVER derive a file list
from a hash prefix by arithmetic; you would silently drop rows. This script batches
consecutive shards and reads each batch's true range from the footers instead.
protein_hash is globally unique (verified across overlapping shards), so batches never emit
duplicate rows even though adjacent batches' hash ranges overlap slightly at the seams.
NO ROW-GROUP PRUNING is available for the full join. At ~100 rows/row group and ~1 in 6
selectivity, P(a row group has zero matches) ~ (5/6)^100 ~ 1e-8. Sortedness buys the merge
join, not fewer bytes -- reading all of `max_pooled` is the floor for an all-rows join.
IF YOU ONLY NEED A SUBSET OF PROTEINS
--------------------------------------
If you want a specific protein set / a sample / top-K by SAE feature rather than all
matches, invert this: drive from your key list and use the protein_hash BTree for indexed
lookups on the Lance side, plus a narrow range predicate on the parquet side (which DOES
prune row groups when the selection is narrow). That drops total bytes by orders of
magnitude and is the only variant reasonable to run outside us-west-2.
"""
import argparse
import duckdb
import lance
import pyarrow as pa
import pyarrow.fs as pafs
import pyarrow.parquet as pq
REGION = "us-west-2"
BUCKET = "esm-protein-atlas"
# The bucket is public: unsigned LIST and ranged GET both succeed and requester-pays is off.
# So reads need NO credentials -- you only need AWS creds for writing your own --out-uri.
# DuckDB reads unsigned via plain https:// URLs; Lance via aws_skip_signature.
S3_PREFIX = f"s3://{BUCKET}/v1/sae/data_shards"
HTTPS_PREFIX = f"https://{BUCKET}.s3.{REGION}.amazonaws.com/v1/sae/data_shards"
# Both structure datasets are read and unioned. They are DISJOINT, not nested (sampled
# folds_atlas hashes had no matches in folds_1B), so together they are the complete structure
# set and the union needs no dedup. Their projected schemas are identical (protein_hash:
# string, structure_blob: binary), so concat is safe.
FOLDS_URIS = (
f"s3://{BUCKET}/v1/folds/folds_1B.lance",
f"s3://{BUCKET}/v1/folds/folds_atlas.lance",
)
# Shards are equal-row-count splits of a globally protein_hash-sorted stream.
# Global order is bucket char in "0123456789abcdef", then the 0000-4095 index.
# NOTE: shard boundaries do NOT align with hash-prefix boundaries (e.g. a_0095 spans
# a05->a06), so never derive a file list from a hash prefix by arithmetic. Batch
# consecutive shards and read their true range from the footers instead.
# protein_hash is globally unique (verified across overlapping shards), so batches never
# emit duplicate rows even though adjacent batches' hash ranges overlap slightly at the seams.
BUCKETS = "0123456789abcdef"
SHARDS_PER_BUCKET = 4096
def all_shards(anonymous: bool = True) -> list[str]:
"""All 65,536 shard URIs in global protein_hash order."""
prefix = HTTPS_PREFIX if anonymous else S3_PREFIX
return [
f"{prefix}/deduplicated_{b}_{n:04d}.parquet"
for b in BUCKETS
for n in range(SHARDS_PER_BUCKET)
]
def _arrow(res):
"""duckdb renamed fetch_arrow_table() -> to_arrow_table() in 1.5; support both."""
fn = getattr(res, "to_arrow_table", None) or res.fetch_arrow_table
return fn()
def connect(anonymous: bool = True) -> duckdb.DuckDBPyConnection:
con = duckdb.connect()
con.execute("INSTALL httpfs; LOAD httpfs")
if not anonymous:
con.execute(
f"CREATE SECRET (TYPE s3, PROVIDER credential_chain, REGION '{REGION}')"
)
con.execute("SET preserve_insertion_order=false")
return con
def folds_storage_options(anonymous: bool = True) -> dict[str, str]:
opts = {"region": REGION}
if anonymous:
opts["aws_skip_signature"] = "true"
return opts
def hash_range(con: duckdb.DuckDBPyConnection, files: list[str]) -> tuple[str, str]:
"""Read the true protein_hash range of `files` from Parquet footers only.
Footer-only: no column data is transferred.
Raises if any protein_hash chunk lacks min/max stats. Without this guard a single
stats-less row group would silently narrow [lo, hi], so the matching folds rows would
never be fetched and those joins would vanish with no error across a 4096-task job.
"""
lo, hi, n_null = con.execute(
"""
SELECT min(stats_min_value), max(stats_max_value),
count(*) FILTER (stats_min_value IS NULL OR stats_max_value IS NULL)
FROM parquet_metadata($files)
WHERE path_in_schema = 'protein_hash'
""",
{"files": files},
).fetchone()
if n_null:
raise ValueError(
f"{n_null} protein_hash chunks have no footer stats; range unsafe"
)
if lo is None or hi is None:
raise ValueError("no protein_hash stats found in footers")
return lo, hi
def run_batch(
batch_id: int, batch_size: int, out_uri: str, anonymous: bool = True
) -> int:
# Explicit region: pyarrow's S3 region autodetect needs an extra call and can pick wrong.
fs, out_path = pafs.FileSystem.from_uri(out_uri)
if isinstance(fs, pafs.S3FileSystem):
fs = pafs.S3FileSystem(region=REGION)
dest = f"{out_path}/batch_{batch_id:05d}.parquet"
# Idempotent: a 4096-task array job will get preempted and requeued.
if fs.get_file_info(dest).type != pafs.FileType.NotFound:
print(f"batch {batch_id}: {dest} exists, skipping", flush=True)
return 0
con = connect(anonymous)
shards = all_shards(anonymous)
files = shards[batch_id * batch_size : (batch_id + 1) * batch_size]
if not files:
raise ValueError(f"batch_id {batch_id} is out of range")
lo, hi = hash_range(con, files)
print(
f"batch {batch_id}: {len(files)} shards, protein_hash [{lo} .. {hi}]",
flush=True,
)
# Folds side, fetched ONCE for the whole batch range: protein_hash + structure_blob only
# (~19% of dataset bytes; `pae` alone is ~80% and we never touch it). The range predicate
# hits the protein_hash BTree index on each dataset, and because fragments are
# range-partitioned on protein_hash only the aligned fragments are read.
predicate = f"protein_hash >= '{lo}' AND protein_hash <= '{hi}'"
parts = []
for uri in FOLDS_URIS:
part = lance.dataset(
uri, storage_options=folds_storage_options(anonymous)
).to_table(columns=["protein_hash", "structure_blob"], filter=predicate)
print(
f"batch {batch_id}: {uri.rsplit('/', 1)[-1]} rows in range = {part.num_rows}",
flush=True,
)
parts.append(part)
fold = pa.concat_tables(parts)
con.register("fold", fold)
print(
f"batch {batch_id}: folds rows in range = {fold.num_rows} (union)", flush=True
)
# Stream one shard at a time; materialising the whole batch would hold every shard's
# max_pooled in memory at once.
tmp = dest + ".tmp"
writer = None
total = 0
try:
for i, f in enumerate(files):
sae = _arrow(
con.execute(
"""
SELECT split_part(header, '|', 1) AS protein_hash, max_pooled
FROM read_parquet($files)
""",
{"files": [f]},
)
)
con.register("sae", sae)
# Guard: every SAE row must fall inside the footer-derived range, else the folds
# fetch above was too narrow and we would silently drop joins.
outside = con.execute(
"SELECT count(*) FROM sae WHERE protein_hash < $lo OR protein_hash > $hi",
{"lo": lo, "hi": hi},
).fetchone()[0]
if outside:
raise ValueError(
f"batch {batch_id} shard {f}: {outside} rows outside [{lo}, {hi}]"
)
out = _arrow(
con.execute(
"""
SELECT s.protein_hash, s.max_pooled, f.structure_blob
FROM sae s JOIN fold f USING (protein_hash)
"""
)
)
con.unregister("sae")
if writer is None:
writer = pq.ParquetWriter(
tmp, out.schema, filesystem=fs, compression="zstd"
)
writer.write_table(out)
total += out.num_rows
print(
f" [{i + 1}/{len(files)}] {f.split('/')[-1]}: "
f"sae={sae.num_rows} joined={out.num_rows}",
flush=True,
)
finally:
if writer is not None:
writer.close()
# Rename only after a clean close, so a preempted task never leaves a half-written file
# that the idempotency check would mistake for a completed batch.
fs.move(tmp, dest)
print(f"batch {batch_id}: wrote {total} rows -> {dest}", flush=True)
return total
def main() -> None:
p = argparse.ArgumentParser()
p.add_argument("--batch-id", type=int, required=True)
p.add_argument("--batch-size", type=int, default=16)
p.add_argument("--out-uri", required=True, help="s3://... destination, same region")
p.add_argument(
"--signed-reads",
action="store_true",
help="sign read requests with the AWS credential chain (default: unsigned, "
"since the source bucket is public)",
)
a = p.parse_args()
run_batch(a.batch_id, a.batch_size, a.out_uri, anonymous=not a.signed_reads)
if __name__ == "__main__":
main()
ESM Atlas
Join SAE features and structures
Join per-protein SAE features with predicted structures by protein hash.
Join per-protein SAE features with predicted structures by protein hash.
Download extract_sae_folds_join.py