sunchao commented on code in PR #6098:
URL: https://github.com/apache/datafusion-comet/pull/6098#discussion_r4075215243
##########
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
+ .len()
+ .saturating_add(std::mem::size_of_val(schema.as_ref()))
+ .saturating_add(schema.fields().size())
+ .saturating_add(
+ schema
+ .metadata()
+ .capacity()
+ .saturating_mul(std::mem::size_of::<(String, String)>()),
+ );
+ for (key, value) in schema.metadata() {
+ retained_size = retained_size
Review Comment:
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`.
##########
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:
Addressed in feccb4df: 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.
##########
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:
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.
##########
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:
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.
##########
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:
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.
##########
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:
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.
--
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]
