viirya commented on code in PR #6098:
URL: https://github.com/apache/datafusion-comet/pull/6098#discussion_r4074374416
##########
native/shuffle/src/ipc.rs:
##########
@@ -108,9 +111,9 @@ fn cached_schema(
let hit = schemas
.iter()
.position(|(message, _)| message.as_ref() == schema_message)?;
- // most recently used first, so an alternating pair stays resident
+ // Promote the hit without changing the relative recency of the other
entries.
if hit != 0 {
- schemas.swap(0, hit);
+ schemas[..=hit].rotate_right(1);
Review Comment:
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.
##########
native/shuffle/src/ipc.rs:
##########
@@ -46,11 +46,14 @@ const CONTINUATION_MARKER: [u8; 4] = [0xff; 4];
/// from several shuffles, and a single entry would thrash.
const SCHEMA_CACHE_CAPACITY: usize = 4;
+/// Maximum estimated serialized-plus-parsed size per cached schema, excluding
allocator overhead.
+const SCHEMA_CACHE_ENTRY_RETAIN_LIMIT: usize = 1 << 20;
+
/// Metadata scratch larger than this is released after the block rather than
kept for the thread.
-/// Real metadata is a few KiB even for wide schemas; only a corrupt length
gets anywhere near.
const SCRATCH_RETAIN_LIMIT: usize = 1 << 20;
Review Comment:
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.
##########
native/shuffle/src/ipc.rs:
##########
@@ -120,6 +123,29 @@ fn cache_schema(
schema_message: &[u8],
schema: SchemaRef,
) {
+ // Admission only affects reuse. Large valid schemas still decode, without
evicting useful
+ // entries or retaining their serialized and parsed copies for the
lifetime of the thread.
+ if schema_message.len() > SCHEMA_CACHE_ENTRY_RETAIN_LIMIT {
Review Comment:
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.
##########
native/shuffle/src/ipc.rs:
##########
@@ -120,6 +123,29 @@ fn cache_schema(
schema_message: &[u8],
schema: SchemaRef,
) {
+ // Admission only affects reuse. Large valid schemas still decode, without
evicting useful
+ // entries or retaining their serialized and parsed copies for the
lifetime of the thread.
+ if schema_message.len() > SCHEMA_CACHE_ENTRY_RETAIN_LIMIT {
+ return;
+ }
+ let mut retained_size = schema_message
Review Comment:
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.
##########
native/shuffle/src/ipc.rs:
##########
@@ -695,6 +722,85 @@ mod tests {
assert_eq!(schema_cache_stats(), stats(5, 6), "most recent stays");
}
+ #[test]
+ fn promoting_a_schema_preserves_eviction_order() {
+ let blocks: Vec<_> = (1..=5)
+ .map(|columns| block_for(&n_column_batch(columns), b"NONE"))
+ .collect();
+ reset_schema_cache();
+ // A B C D A E D: promoting A must keep D newer than B and C, so E
evicts B.
+ for index in [0, 1, 2, 3, 0, 4, 3] {
+ assert_eq!(
+ read_ipc_compressed(&blocks[index]).unwrap(),
+ n_column_batch(index + 1)
+ );
+ }
+ assert_eq!(schema_cache_stats(), stats(2, 5));
+ read_ipc_compressed(&blocks[1]).unwrap();
+ assert_eq!(schema_cache_stats(), stats(2, 6));
+ }
+
+ #[test]
+ #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI.
+ fn oversized_schemas_decode_without_retention_or_eviction() {
+ let normal_blocks: Vec<_> = (1..=SCHEMA_CACHE_CAPACITY)
+ .map(|columns| block_for(&n_column_batch(columns), b"NONE"))
+ .collect();
+ let schemas = [
+ Schema::new(vec![Field::new(
+ "x".repeat(SCHEMA_CACHE_ENTRY_RETAIN_LIMIT + 1),
+ DataType::Int32,
+ false,
+ )]),
+ // Each wire message fits the limit, but its parsed copy pushes
retention over it.
+ Schema::new(vec![Field::new(
+ "x".repeat(SCHEMA_CACHE_ENTRY_RETAIN_LIMIT / 2),
+ DataType::Int32,
+ false,
+ )]),
+ Schema::new(vec![Field::new("c", DataType::Int32,
false)]).with_metadata(
+ HashMap::from([(
+ "key".into(),
+ "x".repeat(SCHEMA_CACHE_ENTRY_RETAIN_LIMIT / 2),
+ )]),
+ ),
+ ];
+ for (index, schema) in schemas.into_iter().enumerate() {
+ let batch = RecordBatch::try_new(
+ Arc::new(schema),
+ vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
+ )
+ .unwrap();
+ let ipc = ipc_bytes(&batch);
+ if index > 0 {
Review Comment:
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.)
##########
native/shuffle/src/ipc.rs:
##########
@@ -695,6 +722,85 @@ mod tests {
assert_eq!(schema_cache_stats(), stats(5, 6), "most recent stays");
}
+ #[test]
+ fn promoting_a_schema_preserves_eviction_order() {
+ let blocks: Vec<_> = (1..=5)
+ .map(|columns| block_for(&n_column_batch(columns), b"NONE"))
+ .collect();
+ reset_schema_cache();
+ // A B C D A E D: promoting A must keep D newer than B and C, so E
evicts B.
+ for index in [0, 1, 2, 3, 0, 4, 3] {
+ assert_eq!(
+ read_ipc_compressed(&blocks[index]).unwrap(),
+ n_column_batch(index + 1)
+ );
+ }
+ assert_eq!(schema_cache_stats(), stats(2, 5));
+ read_ipc_compressed(&blocks[1]).unwrap();
+ assert_eq!(schema_cache_stats(), stats(2, 6));
+ }
+
+ #[test]
+ #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI.
+ fn oversized_schemas_decode_without_retention_or_eviction() {
Review Comment:
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.
##########
native/shuffle/src/ipc.rs:
##########
@@ -46,11 +46,14 @@ const CONTINUATION_MARKER: [u8; 4] = [0xff; 4];
/// from several shuffles, and a single entry would thrash.
const SCHEMA_CACHE_CAPACITY: usize = 4;
+/// Maximum estimated serialized-plus-parsed size per cached schema, excluding
allocator overhead.
+const SCHEMA_CACHE_ENTRY_RETAIN_LIMIT: usize = 1 << 20;
Review Comment:
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
