Add lightweight PDB (Protein Data Bank) file support

[behroozazarkhalili]

Summary

This PR adds support for loading PDB (Protein Data Bank) files with load_dataset(), following the ImageFolder pattern where one row = one structure.

Based on feedback from @lhoestq in #7930, this approach makes datasets more practical for ML workflows:

• Each row is independent, enabling train/test splits and shuffling
• Easy to add labels (folder-based) and metadata (metadata.jsonl)
• Compatible with Dataset Viewer (one 3D render per row)

Architecture

Uses FolderBasedBuilder pattern (like ImageFolder, AudioFolder):

class PdbFolder(FolderBasedBuilder):
    BASE_FEATURE = ProteinStructure
    BASE_COLUMN_NAME = "structure"
    EXTENSIONS = [".pdb", ".ent"]

New ProteinStructure Feature Type

# Arrow schema for lazy loading
pa.struct({"bytes": pa.binary(), "path": pa.string()})

# Decoded: returns structure file content as string
dataset = load_dataset("pdb", data_dir="structures/")
print(dataset[0]["structure"])  # Full PDB file content

Supported Extensions

.pdb, .ent

Usage

from datasets import load_dataset

# Load from directory
dataset = load_dataset("pdb", data_dir="protein_structures/")

# Load with folder-based labels
# structures/
#   enzymes/
#     1abc.pdb
#   receptors/
#     2def.pdb
dataset = load_dataset("pdb", data_dir="structures/")
print(dataset[0])  # {"structure": "HEADER...", "label": "enzymes"}

# Load with metadata
# structures/
#   1abc.pdb
#   metadata.jsonl  # {"file_name": "1abc.pdb", "resolution": 2.5}
dataset = load_dataset("pdb", data_dir="structures/")
print(dataset[0])  # {"structure": "HEADER...", "resolution": 2.5}

# Drop labels or metadata
dataset = load_dataset("pdb", data_dir="structures/", drop_labels=True)
dataset = load_dataset("pdb", data_dir="structures/", drop_metadata=True)

Test Results

All 28 PDB tests + 15 ProteinStructure feature tests pass.

Related PRs

• feat: Add mmCIF file support for macromolecular structures #7925 - mmCIF support (same pattern)
• Proposal: Protein 3D Structure Visualization for Dataset Viewer #7930 - Protein 3D visualization proposal

cc @lhoestq @georgia-hf

[lhoestq]

Hi ! same comment as in #7925: is there a lib in python that parses the text of the files ? Maybe it would be more useful to return a parsed object rather than just the data as string

[behroozazarkhalili]

Hi @lhoestq, thanks for the review!

Same plan as #7925 (cross-linking so the discussion stays in one place):

On a parsing library: mature options exist — gemmi, biotite, and Biopython all parse PDB/mmCIF — but they're all compiled/heavyweight dependencies (Biopython ~150 MB installed; gemmi and biotite ship C/C++ extensions), which is a lot to add to datasets for one loader. I kept this consistent with the other lightweight bio loaders (FASTA/FASTQ, GenBank), which use small pure-Python parsers with zero dependencies.

On returning a parsed object instead of a string: agreed, and it's a regression I introduced. The original version parsed ATOM/HETATM records into typed columns (atom_name, element, residue, chain, x/y/z, occupancy, B-factor, …) with a record_types filter. When I moved to the one-row-per-structure layout you suggested in #7930, I replaced that with a ProteinStructure feature that currently returns the raw file text — so the parsing (and the record_types filter) was dropped.

Proposal — combine both so each row is one structure and a parsed object, with no new dependency:

• one row = one structure (independence for splits/shuffling/metadata, per Proposal: Protein 3D Structure Visualization for Dataset Viewer #7930)
• the structure column becomes a parsed struct of atom-level arrays from the existing lightweight parser, e.g.

structure: {
  atom_name: List[str], element: List[str],
  residue:   List[str], chain_id: List[str], seq_id: List[int],
  x: List[float], y: List[float], z: List[float],
  occupancy: List[float], b_factor: List[float],
}

• reintroduce the record_types (ATOM/HETATM) filter as a config option on top of the parsed schema
• optionally keep the raw text available too (ProteinStructure(decode=False) for the bytes, or a raw field) for users who want to parse it themselves

That keeps it queryable and ML-friendly while staying dependency-free. Does that direction work for you?

[behroozazarkhalili]

Hi @lhoestq — same update as #7925, applied here so the two PRs stay consistent.

The structure column is now parsed on decode into a struct-of-arrays of atom-level columns (one row = one structure preserved). PDB fixed-width ATOM/HETATM records are mapped onto the same PDBx/mmCIF _atom_site dictionary names the mmCIF loader uses, so a PDB dataset and an mmCIF dataset expose identical columns:

>>> ds = load_dataset("pdb", data_dir="structures")
>>> ds[0]["structure"]
{'type_symbol': ['N', 'C', ...], 'label_atom_id': ['N', 'CA', ...],
 'label_comp_id': ['ALA', 'ALA', ...], 'Cartn_x': [0.0, 1.458, ...], ...}

• Dependency-free, shared pure-Python parser (features/protein_parsing.py) — one canonical column/dtype table used by both formats and the feature, no duplication.
• Lazy parse on decode, mirroring Image/Audio/Mesh; storage stays {bytes, path} so ProteinStructure(decode=False) still returns the raw file.
• Config: include_hetatm (default True) and columns=[...] subset selection.

Existing tests pass; added coverage for the parsed schema, HETATM filtering, and column selection, and loading.mdx is updated. Ready for another look — happy to tweak the schema if you'd like a different set of columns.

[lhoestq]

Hi ! same comment as in the mmCIF PR:

I feel like it would be better for people to extend the parsing and to add useful processing methods if the decoding returns an object like a biopython structure (esp. since the community has been contributing to biopython mostly so they can contribute there). IMO It's ok if the dependency is not that lightweight, what matters the most is to get an object that users can easily process/convert/prepare to train models

lmk what you think