Skip to content

gh-155460: Pre-allocate the output buffer in compression.zstd decompression - #155464

Open
ThomasWaldmann wants to merge 1 commit into
python:mainfrom
ThomasWaldmann:zstd-decompress-prealloc
Open

gh-155460: Pre-allocate the output buffer in compression.zstd decompression#155464
ThomasWaldmann wants to merge 1 commit into
python:mainfrom
ThomasWaldmann:zstd-decompress-prealloc

Conversation

@ThomasWaldmann

Copy link
Copy Markdown
Contributor

compression.zstd decompression always builds its output through the growing
output buffer (initial 32 KiB, then progressively larger growth steps, finished
with a resize to the actual size), even though most zstd frames record the exact
decompressed size in the frame header: the one-shot compression APIs
(zstd.compress(), ZSTD_compress2()) write it by default.

When the input starts at such a frame header, decompress_lock_held() now reads
the recorded size with ZSTD_getFrameContentSize() and allocates the output
buffer at its exact size via _OutputBuffer_InitWithSize(). An exactly filled
buffer is returned without any resize or copy. Everything else falls back to the
existing _OutputBuffer_InitAndGrow() path: frames without a recorded size,
continuation calls in the middle of a frame (ZSTD_CONTENTSIZE_ERROR), and
skippable frames (recorded size 0).

This is only a sizing hint — decompression does not rely on it:

  • a hand-crafted header recording a wrong size gives exactly the same results
    and exceptions as before (libzstd validates the recorded size itself; the
    buffer grows further if the hint was too small, or is shrunk to the actual
    size on finish);
  • the recorded size is trusted only up to 1 GiB (OUTPUT_PREALLOC_MAX), so a
    crafted header cannot request an arbitrarily large allocation;
  • a zstd block cannot expand to more than 128 KiB from less than 4 bytes of
    input, so recorded sizes claiming more than a 32768x expansion
    (OUTPUT_MAX_EXPANSION) of the available input are unfulfillable and are
    ignored;
  • _OutputBuffer_InitWithSize() clamps to max_length, so the max_length
    argument keeps working unchanged.

This covers both ZstdDecompressor.decompress() and the module-level
zstd.decompress().

Tests

Six tests are added to DecompressorTestCase, including a helper that rewrites
the Frame_Content_Size field of a frame per RFC 8878 §3.1.1.1 so that wrong
and absurd recorded sizes can be tested: sizes that are too big, too small or
zero, 0xFFFFFFFF declared on a 130 KiB frame (hits the 1 GiB guard) and on a
tiny RLE frame (hits the expansion-ratio guard), round-trips at block
boundaries, the unknown-size fallback path, max_length, and split input.

They all pass on unpatched main as well — deliberately, since they pin
behaviour rather than the implementation.

Benchmarks

zstd.decompress() of level-1 compressed frames of moderately compressible
data; 10 timed runs after an untimed warmup, median run. Full tables, per-size
results and the profiling analysis are in gh-155460.

A mixed-size workload, 1000 chunks with sizes drawn from a normal distribution
in log2(size) centred on 2 MiB spanning 512 KiB to 8 MiB (2269 MiB total):

main patched speedup
Linux/x86-64 (Ryzen 5 8500GE, glibc 2.41) 1.407 s 1.135 s +19.3%
macOS/arm64 (M3 Pro) 1.245 s 1.185 s +4.9%

Single-size frames, same two machines: on Linux the patch is faster at every
size measured (+5–8.6% up to 5 MiB, +41–55% for 6–16 MiB, +0.7–2.7% at
32–256 MiB); on macOS +10% at 1 MiB, +18–19% at 6–8 MiB, +4.6–12% at 9–16 MiB,
and a 2–8.5% regression in the 3–5 MiB band, where the unpatched path needs no
growth step and its cache-resident internal window beats direct writes spanning
the whole output.

Prior art: python-zstandard's decompress() allocates the output buffer from
the frame's recorded content size, and pyzstd (from which this module descends)
did the same before the stdlib port.


Disclosure: this change was prepared with AI assistance (Claude Code); the
benchmarks and the analysis above were run and reviewed by me.

…ecompression

Frames produced by the one-shot compression APIs record the decompressed
size in the frame header.  When the input starts at such a frame header,
read that size with ZSTD_getFrameContentSize() and allocate the output
buffer at its exact size, instead of growing it progressively and
shrinking it on finish.  An exactly filled buffer is then returned
without a copy.

This is only a sizing hint: decompression does not rely on it, so a
hand-crafted header recording a wrong size produces exactly the same
results and exceptions as before.  Two guards bound the allocation: the
recorded size is trusted only up to 1 GiB, and recorded sizes claiming
more than a 32768x expansion of the available input -- more than the
format can produce -- are ignored.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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