dwsmith1983 commented on code in PR #5568:
URL: https://github.com/apache/datafusion-comet/pull/5568#discussion_r3913226639
##########
native/shuffle/src/writers/buf_batch_writer.rs:
##########
@@ -160,3 +209,140 @@ impl<S: Borrow<ShuffleBlockWriter>, W: Write + Seek>
BufBatchWriter<S, W> {
self.writer.stream_position().map_err(Into::into)
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::{read_ipc_compressed, CompressionCodec};
+ use arrow::array::Int64Array;
+ use arrow::datatypes::{DataType, Field, Schema};
+ use std::sync::Arc;
+
+ fn test_batch(seed: i64) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new("a",
DataType::Int64, false)]));
+ let values: Vec<i64> = (0..100).map(|i| seed * 1_000 + i).collect();
+ RecordBatch::try_new(schema,
vec![Arc::new(Int64Array::from(values))]).unwrap()
+ }
+
+ fn write_one_partition(seed: i64, scratch: &mut Vec<u8>) -> Vec<u8> {
+ let batch = test_batch(seed);
+ let block_writer =
+ ShuffleBlockWriter::try_new(batch.schema().as_ref(),
CompressionCodec::Zstd(1))
+ .unwrap();
+ let mut output = Vec::new();
+ let time = Time::default();
+ let mut writer = BufBatchWriter::new(block_writer, &mut output, 1 <<
20, 8192);
+ writer.write(&batch, scratch, &time, &time).unwrap();
+ writer.flush(scratch, &time, &time).unwrap();
+ output
+ }
+
+ /// A scratch buffer recycled across partitions must produce
byte-identical output to
+ /// fresh per-partition buffers, come back drained, and keep its grown
capacity.
+ #[test]
+ #[cfg_attr(miri, ignore)] // miri can't call zstd's C FFI
+ fn recycled_scratch_matches_fresh_buffers_and_keeps_capacity() {
+ let fresh: Vec<Vec<u8>> = (0..3)
+ .map(|p| write_one_partition(p, &mut Vec::new()))
+ .collect();
+
+ let mut scratch = Vec::new();
+ let mut recycled = Vec::new();
+ for p in 0..3 {
+ let output = write_one_partition(p, &mut scratch);
+ assert!(
+ scratch.is_empty(),
+ "recycled scratch must come back drained"
+ );
+ recycled.push(output);
+ }
+
+ assert_eq!(fresh, recycled);
+ assert!(
+ scratch.capacity() > 0,
+ "capacity grown in one partition must survive into the next"
+ );
+ for output in &recycled {
+ let decoded = read_ipc_compressed(&output[16..]).unwrap();
+ assert_eq!(decoded.num_rows(), 100);
+ }
+ }
+
+ /// Handing a non-empty scratch to a fresh writer would silently prepend
stale bytes
+ /// to the first block; debug builds must catch it.
+ #[cfg(debug_assertions)]
+ #[test]
+ #[should_panic(expected = "non-empty scratch")]
+ fn fresh_writer_rejects_dirty_scratch() {
+ let batch = test_batch(0);
+ let block_writer =
+ ShuffleBlockWriter::try_new(batch.schema().as_ref(),
CompressionCodec::None).unwrap();
+ let mut output = Vec::new();
+ let time = Time::default();
+ let mut writer = BufBatchWriter::new(block_writer, &mut output, 1 <<
20, 8192);
+ let mut dirty = vec![0xAB, 0xCD];
+ let _ = writer.write(&batch, &mut dirty, &time, &time);
+ }
+
+ /// Swapping in a different scratch mid-writer would silently abandon any
bytes still
+ /// buffered in the first one; the identity check has to catch it in debug
builds.
+ #[test]
+ #[cfg(debug_assertions)]
+ #[should_panic(expected = "same scratch buffer")]
+ fn writer_rejects_swapped_scratch() {
+ let batch = test_batch(0);
+ let block_writer =
+ ShuffleBlockWriter::try_new(batch.schema().as_ref(),
CompressionCodec::None).unwrap();
+ let mut output = Vec::new();
+ let time = Time::default();
+ let mut writer = BufBatchWriter::new(block_writer, &mut output, 1 <<
20, 8192);
+ let mut first = Vec::new();
+ writer.write(&batch, &mut first, &time, &time).unwrap();
+ let mut second = Vec::new();
+ let _ = writer.write(&batch, &mut second, &time, &time);
+ }
+
+ /// A block that crosses `buffer_max_size` grows the scratch past the cap;
`flush`
+ /// must shrink retained capacity back to the configured buffer size,
while a
+ /// normally-sized run keeps its (sub-cap) capacity untouched.
+ #[test]
+ fn flush_caps_retained_scratch_capacity() {
+ let batch = test_batch(0); // 100 rows of Int64: block is far larger
than 64 bytes
+ let buffer_max_size = 64usize;
+ // batch_size below the row count so the batch bypasses the coalescer
and is
+ // serialized into the scratch during `write`.
+ let batch_size = 10usize;
+ let block_writer =
+ ShuffleBlockWriter::try_new(batch.schema().as_ref(),
CompressionCodec::None).unwrap();
+ let mut output = Vec::new();
+ let time = Time::default();
+ let mut scratch = Vec::new();
+ let mut writer =
+ BufBatchWriter::new(block_writer, &mut output, buffer_max_size,
batch_size);
+ writer.write(&batch, &mut scratch, &time, &time).unwrap();
+ assert!(
+ scratch.capacity() > buffer_max_size,
+ "oversized block must have grown the scratch past the cap"
+ );
+ writer.flush(&mut scratch, &time, &time).unwrap();
+ assert!(scratch.is_empty());
+ assert!(
+ scratch.capacity() <= buffer_max_size,
+ "retained capacity {} exceeds cap {}",
+ scratch.capacity(),
+ buffer_max_size
+ );
+
+ // With a roomy cap the grown capacity is retained (shrink_to never
grows the
+ // target below the cap, so no over-shrinking).
+ let large_cap = 1 << 20;
+ let block_writer =
+ ShuffleBlockWriter::try_new(batch.schema().as_ref(),
CompressionCodec::None).unwrap();
+ let mut output = Vec::new();
+ let mut scratch = Vec::new();
+ let mut writer = BufBatchWriter::new(block_writer, &mut output,
large_cap, 8192);
+ writer.write(&batch, &mut scratch, &time, &time).unwrap();
+ writer.flush(&mut scratch, &time, &time).unwrap();
+ assert!(scratch.capacity() > 0 && scratch.capacity() <= large_cap);
Review Comment:
You are right, the coalescer was buffering everything so the second half was
not proving anything. The test now uses a batch size below the row count so the
writes actually serialize, captures the capacity right after them, and asserts
flush leaves it exactly unchanged. I checked it fails if the flush path shrinks
the buffer.
--
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]
