Skip to content

Parse the NumPy arrays of pickles (#1241) - #1241

Open
Cyan4973 wants to merge 2 commits into
meta:devfrom
Cyan4973:export-D128229489
Open

Cyan4973 wants to merge 2 commits into
meta:devfrom
Cyan4973:export-D128229489

Conversation

@Cyan4973

@Cyan4973 Cyan4973 commented Oct 30, 2026 •

Copy link
Copy Markdown
Contributor

Summary:

Pickles often hold NumPy arrays: pickle.dump of dicts of arrays, pandas.to_pickle, model weights. On Zenodo, .pkl files total 3.0 TB, and NumPy arrays make about 74% of their bytes. Today the numpy profile compresses pickles with the generic graph.

This diff adds a pickle lexer (custom_parsers/pickle_lexer.h), which finds the bytes of NumPy arrays in pickles of protocols 2 to 5 by following their opcodes, without executing anything. The NumPy parser uses it: arrays go to the Numeric Transformer, the rest to zstd. A new pickle profile selects this graph.

26 real pickles zstd -19 today -p pickle
16 with arrays (1.33 GB) x1.565 x1.488 x1.779
10 without arrays (1.00 GB) x1.925 x1.777 x1.777

Pickles without arrays compress as today. Arrays smaller than 16 KiB stay with the opcodes, where they compress better. Object arrays of .npy files come in the next diff.

Differential Revision: D128229489

@meta-cla meta-cla Bot added the cla signed label Oct 30, 2026
@meta-codesync

meta-codesync Bot commented Oct 30, 2026

Copy link
Copy Markdown

@Cyan4973 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D128229489.

@meta-codesync meta-codesync Bot changed the title Parse the NumPy arrays of pickles Parse the NumPy arrays of pickles (#1241) Nov 2, 2026
Cyan4973 added a commit to Cyan4973/openzl that referenced this pull request Nov 2, 2026
Summary:
Pull Request resolved: meta#1241

Pickles often hold NumPy arrays: `pickle.dump` of dicts of arrays, `pandas.to_pickle`, model weights. On Zenodo, `.pkl` files total 3.0 TB, and NumPy arrays make about 74% of their bytes. Today the `numpy` profile compresses pickles with the generic graph.

This diff adds a pickle lexer (`custom_parsers/pickle_lexer.h`), which finds the bytes of NumPy arrays in pickles of protocols 2 to 5 by following their opcodes, without executing anything. The NumPy parser uses it: arrays go to the Numeric Transformer, the rest to zstd. A new `pickle` profile selects this graph.

| 26 real pickles | zstd -19 | today | `-p pickle` |
| --- | --- | --- | --- |
| 16 with arrays (1.33 GB) | x1.565 | x1.488 | x1.779 |
| 10 without arrays (1.00 GB) | x1.925 | x1.777 | x1.777 |

Pickles without arrays compress as today. Arrays smaller than 16 KiB stay with the opcodes, where they compress better. Object arrays of `.npy` files come in the next diff.

Differential Revision: D128229489
Summary:
NumPy files hold the arrays of many scientific datasets and some model checkpoints, such as the JAX checkpoints of ViT and GenCast. A `.npy` file holds one array, and a `.npz` archive is a zip of `.npy` members. OpenZL had no profile for them, so they went through generic compression, which does no better than zstd on them.

This diff adds a `numpy` CLI profile and `ZL_NumPy_registerGraph()`.
- A small lexer reads the header of each `.npy` file, or of each stored `.npy` member of a `.npz` archive (through `zip_lexer`), to find the type and shape of its array.
- Arrays of numbers (booleans, integers, floats, complex numbers, datetimes) go to the Numeric Transformer, which picks the transforms for each array. Other arrays, headers and zip records are compressed apart.
- Chunking reuses the tensor chunker of the model parsers.

| File | zstd -19 | `-p serial` | `-p numpy` | vs zstd -19 |
|---|---|---|---|---|
| ViT-B/16 JAX checkpoint, float32 `.npz` (413 MB) | x1.289 | x1.286 | **x1.427** | +10.7% |
| GenCast JAX checkpoint, float32 `.npz` (230 MB) | x1.075 | x1.073 | **x1.190** | +10.7% |
| rMD17 aspirin trajectories, float64 `.npz` (154 MB) | x1.168 | x1.048 | **x1.589** | +36.0% |
| rMD17 benzene trajectories, float64 `.npz` (89 MB) | x1.607 | x1.348 | **x2.533** | +57.7% |
| Neuropixels templates, float32 `.npy` (15 MB) | x11.25 | x10.30 | **x13.77** | +22.4% |
| QuickDraw cat bitmaps, uint8 `.npy` (97 MB) | x2.699 | x2.460 | **x2.784** | +3.2% |

Trade-offs:
- The Numeric Transformer compresses at 25-50 MB/s on these files; decompression runs at 0.3-2.5 GB/s. On the float32 weights of ViT, the tensor graphs of the model parsers compress 14 to 18 times faster, for a ratio within 1%. Speeding up the Transformer is a separate effort, so the model parsers keep their tuned graphs.
- Members deflated by `np.savez_compressed` (MedMNIST, Keras MNIST) are not inflated, so they gain nothing.
- In a `.npz` archive, arrays smaller than 16 KiB share a stream with the headers around them, which the generic graph compresses: 0.1-0.2% of the bytes of ViT and GenCast.
- Strings, objects, structured types and big-endian arrays are compressed as raw bytes.

Differential Revision: D127896258
Summary:
Pull Request resolved: meta#1241

Pickles often hold NumPy arrays: `pickle.dump` of dicts of arrays, `pandas.to_pickle`, model weights. On Zenodo, `.pkl` files total 3.0 TB, and NumPy arrays make about 74% of their bytes. Today the `numpy` profile compresses pickles with the generic graph.

This diff adds a pickle lexer (`custom_parsers/pickle_lexer.h`), which finds the bytes of NumPy arrays in pickles of protocols 2 to 5 by following their opcodes, without executing anything. The NumPy parser uses it: arrays go to the Numeric Transformer, the rest to zstd. A new `pickle` profile selects this graph.

| 26 real pickles | zstd -19 | today | `-p pickle` |
| --- | --- | --- | --- |
| 16 with arrays (1.33 GB) | x1.565 | x1.488 | x1.779 |
| 10 without arrays (1.00 GB) | x1.925 | x1.777 | x1.777 |

Pickles without arrays compress as today. Arrays smaller than 16 KiB stay with the opcodes, where they compress better. Object arrays of `.npy` files come in the next diff.

Differential Revision: D128229489

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant