andygrove commented on code in PR #5568:
URL: https://github.com/apache/datafusion-comet/pull/5568#discussion_r3910254634
##########
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:
The second half of `flush_caps_retained_scratch_capacity` cannot fail as
written. With `large_cap` at 1 MiB and a block around a kilobyte,
`scratch.capacity() > 0 && scratch.capacity() <= large_cap` holds for any
implementation, including one where `flush` shrank to a single byte or did not
shrink at all. The comment says it is checking that there is no over-shrinking,
which is the property worth having, but with `batch_size` at 8192 the coalescer
defers serialization to `flush` and `shrink_to` ends up a no-op, so nothing is
being observed.
Could you use a `batch_size` below the row count the way the first half
does, so `write` actually serializes into the scratch, then capture the
capacity and assert `flush` leaves it alone? `assert_eq!(scratch.capacity(),
cap_after_write)` would fail if this ever became `shrink_to_fit`, which is the
regression the case exists to catch.
##########
native/shuffle/src/writers/local/local_partition_writer.rs:
##########
@@ -243,9 +260,18 @@ impl PartitionWriter for LocalPartitionWriter {
);
for batch in iter.by_ref() {
let batch = batch?;
- buf_batch_writer.write(&batch, &metrics.encode_time,
&metrics.write_time)?;
+ buf_batch_writer.write(
Review Comment:
The borrow change fixed the `mem::take` wart I raised, but I think it opened
a different one on the same path. If the iterator or the writer errors part way
through a partition, `finish_partition` returns before `flush` runs, so
`recycled_buffer` keeps whatever was already encoded. I confirmed it with a
probe: inject an error after one good batch and the buffer comes back holding
668 bytes. The next partition's `BufBatchWriter` then seeks to the end and
appends after those bytes, and its `flush` writes them into that partition's
byte range, so a reader sees a corrupt block rather than an error.
Nothing reaches that state today, because every error here aborts the task
and `check_scratch` catches it in debug builds. But it is the one invariant in
this design that the code does not enforce, and the release-build failure mode
is silent wrong data rather than a crash. Both earlier shapes were immune for
free: the owned buffer died with the writer, and the move-in/move-out version
in `3b4425cb` lost recycling rather than corrupting anything. The borrow shape
I pushed you toward is what introduced it, so this one is on me.
Would you mind draining the scratch when the write loop or `flush` errors,
in both `SpillWriter::write` and the `Multi` arm of `finish_partition`? Binding
the loop and flush to a `result` and calling `recycled_buffer.clear()` before
propagating would make the invariant true by construction instead of by
assertion.
--
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]
