fix: bound shuffle schema cache retention and preserve eviction order

[sunchao]

Which issue does this PR close?

Closes #6097.

Rationale for this change

The schema cache introduced by #5809 can evict a recently used schema because promoting a hit swaps another entry into its old slot. For example, A B C D A E D reparses D even though B is older. Limiting the cache to four entries also leaves retained schema metadata unrestricted.

What changes are included in this PR?

• Rotate the prefix when promoting a hit, preserving the relative recency of other entries.
• Bound each thread's cache to four entries and a shared 4 MiB estimated serialized-plus-parsed byte budget. Evict least-recently-used entries until a new schema fits. A schema exceeding the entire budget still decodes successfully without evicting existing entries or copying its serialized key into the cache.
• Extract the retained-size estimator, including recursive fields and metadata string capacities. Store each admitted entry's estimate so hits and eviction do not walk schemas again. This is a retention estimate excluding allocator overhead, not a peak process-memory limit.
• Document the budget rationale and the separate 1 MiB scratch-capacity retention limit.
• Add wide-schema reuse, byte-budget eviction, exact-boundary, metadata-capacity, and reset/release coverage. Name the oversized-schema cases and include codec/decode-mode context in failures.
• Add an 8,000-column case to the shuffle reader's warm, validated, and forced-cold benchmark arms.

Both decoder entry points retain their validation and complete-stream checks.

How are these changes tested?

At feccb4df, with Rust 1.98.1 and the upstream lockfile (Arrow 59.3.0, DataFusion 55.1.0):

• cargo test --locked --offline --profile ci -p datafusion-comet-shuffle: 156 passed, including existing warm-allocation checks.
• cargo clippy --locked --offline --profile ci -p datafusion-comet-shuffle --all-targets -- -D warnings: passed.
• Rustfmt checks on both changed files and git diff --check: passed.
• Wide-schema reuse and oversized bypass cover all four codecs and both decoder entry points. Weak references verify release after eviction/reset and lack of retention for oversized schemas.

The new benchmark ran with cargo bench --locked --offline --profile ci -p datafusion-comet-shuffle --bench shuffle_reader -- 8000col_64row --sample-size 10 --warm-up-time 1 --measurement-time 2 --noplot. For 8,000 alternating Int64/string columns and 64 rows, local Criterion point estimates were:

Codec	Warm decode	Forced-cold decode
None	2.30 ms	5.27 ms
LZ4	6.71 ms	9.82 ms

These compare warm reuse with clearing the cache every iteration in the current implementation; they are local microbenchmark results, not whole-query performance measurements. TPC and JVM suites were not rerun locally.

[peterxcli]

+1

[andygrove]

Thanks for the quick follow-up. The eviction order fix is clearly right, and I like that the rotate keeps the cost to a handful of pointer moves.

I want to push on the 1 MiB admission limit. I measured what actually trips it and it turns out to be column count rather than pathological field names. With plain Int32 columns the estimate crosses 1 MiB at around 5,800 columns. With 30-character column names it is around 4,100. With struct columns of 10 Utf8 subfields it is around 585 of them. Those are wide schemas but they are real ones.

At that point the cache turns off for that schema and every block reparses. Building this locally in the ci profile, a 6,000 column by 1,000 row block decodes in 2152 us with this change. At 5,000 columns the same shape is 877 us warm against 1778 us cold, so the cache is worth about 2x right up to the cutoff. It does not amortize as rows grow either, since the parse scales with columns and only the body scales with rows. That is the win #5809 was merged for, and the guard removes it from precisely the schemas where a reparse costs the most.

Would you consider making the limit a budget for the cache as a whole rather than a per entry admission test? An oversized schema would then displace other entries rather than being refused, which keeps the working schema cached in the case that matters while still bounding what a thread retains. As written the aggregate is SCHEMA_CACHE_CAPACITY * SCHEMA_CACHE_ENTRY_RETAIN_LIMIT, so 4 MiB per thread, and that seems worth stating in the code either way.

If we do keep a refusal, could the doc comment on SCHEMA_CACHE_ENTRY_RETAIN_LIMIT say roughly how many columns 1 MiB buys? That matters more now that the sentence explaining the value of SCRATCH_RETAIN_LIMIT is gone, and someone later deciding whether to raise it has nothing to go on. A 2x decode regression with no log and no metric behind it is hard to diagnose from the outside.

[comphead]

does it makes sense to

retained_size = retained_size.saturating_add(key.capacity() + value.capacity());

?

Or the key+value can overflow too?

[viirya]

Your instinct is right, and I think the current form should stay.

key.capacity() + value.capacity() is a plain addition, so it is the one spot in this computation that could overflow — panicking in a debug build, wrapping in release. Wrapping is the bad case here: the sum would come out small, retained_size would land under the limit, and an oversized schema would be admitted, which is exactly the outcome the check exists to prevent.

Reaching it would need two String capacities summing past usize::MAX, so on a 64-bit target it is not reachable in practice. But the two chained saturating_adds cost nothing, match the style of the rest of the expression, and mean the reader does not have to work out that it is unreachable. Folding them back into one addition would make this the only unguarded arithmetic in the function.

For what it is worth, I checked the composition itself against arrow-schema 59.3.0 while reviewing: Fields::size() recurses through DataType::size() and covers field names and field-level metadata, so this loop is picking up the schema-level metadata it does not reach. No double counting between the two.

[sunchao]

I kept the two calls for consistency with the surrounding size accounting. Combining these particular additions is also safe: each valid String capacity is at most isize::MAX, so the pair cannot overflow usize.

[comphead]

Thanks @sunchao it looks good to me

[viirya]

Thanks for the quick follow-up on #5809. I verified both fixes independently — the LRU change is correct and the test that pins it is well chosen. I do think the 1 MiB retention limit needs another look before this lands, because of what it ends up pricing.

The LRU fix is right. I simulated both policies over the A B C D A E D sequence from #6097: swap gives 1 hit / 6 misses, rotate_right gives 2 hits / 5 misses, matching the asserted stats(2, 5). Reading B afterwards stays a miss, which confirms B — the genuinely oldest entry — is the one evicted, rather than D. rotate_right is also in-bounds at every edge, since hit comes from position(); the if hit != 0 guard is a useful short-circuit but not load-bearing, as [..=0].rotate_right(1) is already a no-op.

I also checked the sizing against arrow-schema 59.3.0: Fields::size() does recurse through DataType::size() and already covers each field's name.capacity() and field-level metadata, there is no Schema::size() upstream to reuse, and the manual loop covers schema-level metadata that Fields::size() genuinely does not. So the composition is correct and nothing is double-counted. All three oversized test schemas really do cross the limit, and cases 2 and 3 have a wire message under the limit with only the parsed copy pushing it over — exactly the gap the early return alone would miss. Running both new tests against the old logic first is the right way to land a regression test.

The retention limit is where I'd push back. The budget is consumed by column count, not by the oversized strings the check is aimed at. A four-character Int32 field costs about 124 bytes of estimated retained size, of which only 4 bytes are the name — the rest is fixed per-field struct overhead. That puts the cutoff at roughly 5,800 flat Int32 columns, or about 800 struct columns (10 utf8 children each):

columns (Int32, short names)	estimated retained	admitted?	re-parse cost per block
4,000	719,676	yes	—
6,000	1,079,676	no	~875 µs
8,000	1,439,676	no	~1,168 µs

That inverts the performance cliff: schema parse cost grows with column count, so wide schemas are where #5809's caching is worth the most, and they are precisely what gets dropped. It is also silent — SchemaCacheStats is #[cfg(test)], so there is no hit-rate signal in a real deployment to explain the slowdown. And shuffle_reader.rs only benchmarks 5 and 50 columns, two orders of magnitude below the cutoff, so nothing in CI would notice. Five thousand columns is within range for wide feature tables and flattened nested schemas; this isn't only reachable by corrupt input, which is the distinction the SCRATCH_RETAIN_LIMIT comment used to draw.

The threat model worth defending against is a single pathological field name or metadata value — attacker-influenced strings — rather than a legitimately wide business schema. A few ways to get that without the cliff, any one is fine:

• Preferred: price only string capacity, excluding the fixed per-field struct overhead. Still blocks pathological strings, stops penalising width.
• Raise the limit substantially (16 MiB, say). Four entries at that size is still far below what shuffle itself holds, and every realistic schema stays cached.
• Budget the cache as a whole and evict down to fit, instead of an all-or-nothing per-entry admission.

Whichever route, a test asserting that a wide schema (8,000 columns) still hits the cache would pin the invariant.

Since the two halves are independent, splitting them is a reasonable option if you'd like the LRU fix to land now.

A few smaller things, left inline: extracting the size computation so it can be tested directly, documenting where the 1 MiB figure comes from, and the two unrelated limits that happen to share a value.

[viirya]

This is the main thing I'd like to resolve before merge (details and numbers in the top-level comment).

The budget here is spent on column count rather than on the oversized strings the check targets. Field::size() for a 4-char Int32 field is ~116 bytes (~124 including the FieldRef), of which only 4 bytes are the name — the rest is fixed struct overhead. So the cutoff lands at ~5,800 flat Int32 columns, or ~800 struct columns. At 6,000 columns the cache goes silently off and every block pays ~875 µs to re-parse, which is the cost #5809 existed to remove.

Suggestion: price only the string capacities (field names, metadata keys/values) and leave out the fixed per-field overhead — that still blocks a pathological name or metadata value without making width the thing that disqualifies a schema. Alternatively raise the constant to something like 16 MiB, or budget the cache as a whole with eviction-to-fit rather than per-entry admission.

Separately, the doc comment explains what the constant is but not why 1 MiB, nor that it implies a ceiling on column count. Worth stating, since that's the part a future reader would need in order to change it safely.

[sunchao]

Addressed in feccb4d: the cache now has a shared 4 MiB budget, keeping the four-entry cap and evicting the least recently used entries until both limits are satisfied. This preserves the previous maximum estimated cache retention while allowing wider schemas to use unused space. I kept full serialized-plus-parsed accounting because the per-field allocations are real retained memory. Added 8,000-column cache/release coverage, byte-budget eviction and boundary tests, and an 8,000-column benchmark case; all 156 shuffle tests pass.

[viirya]

Minor: this and SCHEMA_CACHE_ENTRY_RETAIN_LIMIT now both read 1 << 20 but mean unrelated things — a scratch buffer's capacity versus an estimated retained size across two representations. Worth a note on one of them that the shared value is coincidental, so nobody later assumes they have to move together.

Also, the line dropped from this comment ("Real metadata is a few KiB even for wide schemas; only a corrupt length gets anywhere near") carried useful intent. Removing it from here is right since it described the scratch limit, but the new constant has no equivalent "what actually reaches this" note — and per the other comment, its answer is different: normal wide schemas do reach it.

[sunchao]

Clarified that scratch has its own buffer-capacity retention limit, independent of the cache's estimated serialized-plus-parsed byte budget. The cache comment now explains the shared 4 MiB budget and gives an approximate 8,000-column example; scratch remains at 1 MiB. These are retention limits, not peak decode-memory limits.

[viirya]

Confirmed correct. I simulated both policies over A B C D A E D: swap yields 1 hit / 6 misses, rotate_right yields 2 / 5, matching the assertion in the new test — and a following read of B is still a miss, which is what proves B rather than D was evicted.

Bounds are fine at every edge since hit comes from position(). Note if hit != 0 is a short-circuit rather than a safety requirement — [..=0].rotate_right(1) is already a no-op — so it can stay or go on readability grounds alone.

[viirya]

Good call checking the serialized length before computing retained_size — it avoids walking a huge field list recursively just to reject it. Worth keeping that ordering intentional if this block gets refactored.

[viirya]

Consider extracting this into fn estimated_retained_size(schema_message: &[u8], schema: &Schema) -> usize.

Two reasons: it keeps cache_schema about admission policy rather than arithmetic, and more usefully it makes the estimate directly unit-testable. Right now it can only be exercised end-to-end through a decode, and this is the part most likely to drift quietly on an Arrow upgrade — Fields::size() and DataType::size() are upstream implementation details, so a change there silently moves the cutoff with no failing test.

For what it's worth, the composition itself checks out against arrow-schema 59.3.0: Fields::size() recurses via DataType::size() and covers field names and field-level metadata, there's no Schema::size() to reuse, and the loop below correctly adds the schema-level metadata that Fields::size() omits. No double counting.

[sunchao]

Extracted estimated_retained_size and added direct coverage for spare string capacity in schema metadata and nested field metadata. The nested case also exercises recursive field sizing. Each cache entry stores its estimate so hits and eviction do not walk the schema again.

[viirya]

Coverage here is good — 3 schemas x 4 codecs x 2 entry points, plus the Arc::downgrade / upgrade().is_none() check, which is the most direct way to assert the schema really isn't retained.

The cost is that a failure in the innermost assertion doesn't say which combination produced it. Adding index and codec to the assertion messages (or splitting the schema cases into separate tests) would make that cheaper to diagnose.

[sunchao]

Added the named schema case, readable codec, and validation mode to decode failures and per-decode assertions. Cache and scratch checks also identify the case and codec.

[viirya]

This if index > 0 encodes knowledge that isn't stated anywhere: index 0 is the message-exceeds-limit case, while 1 and 2 are the message-fits-but-parsed-copy-exceeds cases. A one-line comment saying so would help, since that distinction is the whole reason the second check in cache_schema exists rather than just the early return.

(I confirmed the split holds: case 0's stream is 1,049,032 bytes, cases 1 and 2 are 524,744 and 524,808 with estimated retained sizes of 1,049,216 and 1,049,428.)

[sunchao]

Replaced index > 0 with named cases and an explicit wire_fits expectation. The wire-size assertion now covers all three cases, distinguishing an oversized serialized schema from schemas whose parsed copies push total retention over the budget.

[sunchao]

@andygrove @viirya Addressed in feccb4df. The cache now shares a 4 MiB estimated serialized-plus-parsed budget across at most four entries, evicting least-recently-used entries until a new schema fits. Schemas exceeding the whole budget still decode without disturbing existing entries. The LRU promotion fix and early serialized-size check are preserved.

I kept the full retained-size estimate because field/type allocations are real retained memory. The shared budget lets a wide working schema use space left by other entries without raising the previous maximum estimated cache retention. The docs now explain the budget and its independence from the 1 MiB scratch limit.

Added regression tests for 8,000-column reuse across every codec and both decoder entry points, multi-entry eviction, exact budget boundaries, nested/schema metadata capacity, and release after reset. Also addressed the helper extraction and test diagnostics comments. All 156 shuffle tests and all-target Clippy with -D warnings pass; formatting and diff checks pass.

The new 8,000-column / 64-row benchmark compares warm and forced-cold decoding on the current code. Local ci profile point estimates were 2.30 ms vs 5.27 ms uncompressed and 6.71 ms vs 9.82 ms with LZ4 (10 samples, 1 s warm-up, 2 s measurement). These are local microbenchmark results, not whole-query measurements. The PR description now contains the current validation results and commands; TPC/JVM suites were not rerun locally.