andygrove commented on code in PR #5003:
URL: https://github.com/apache/datafusion-comet/pull/5003#discussion_r3639959321
##########
native/shuffle/src/shuffle_writer.rs:
##########
@@ -924,6 +924,80 @@ mod test {
.sum()
}
+ /// Batches that are already at least `batch_size` rows should bypass the
+ /// coalescer, being written one-block-per-batch, while smaller batches are
+ /// still coalesced. In both cases the data must roundtrip exactly, in
order.
+ #[test]
+ #[cfg_attr(miri, ignore)]
+ fn test_full_batches_bypass_coalescer() {
+ use crate::writers::BufBatchWriter;
+ use arrow::array::Int32Array;
+
+ let schema = Arc::new(Schema::new(vec![Field::new("v",
DataType::Int32, false)]));
+ let batch_size = 100;
+
+ let make_batch = |start: i32, len: i32| {
+ let values: Vec<i32> = (start..start + len).collect();
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(Int32Array::from(values)) as Arc<dyn Array>],
+ )
+ .unwrap()
+ };
+
+ // Sequence: two full batches (bypass), a partial batch (buffered),
then a
+ // full batch (must not bypass because the coalescer still holds
buffered rows).
+ let inputs = vec![
+ make_batch(0, batch_size), // rows 0..100 -> bypass
+ make_batch(batch_size, batch_size), // rows 100..200 -> bypass
+ make_batch(200, 30), // rows 200..230 -> buffered
+ make_batch(230, batch_size), // rows 230..330 -> coalesced
with buffered 30
+ ];
+
+ let codec = CompressionCodec::Lz4Frame;
+ let encode_time = Time::default();
+ let write_time = Time::default();
+
+ let mut output = Vec::new();
+ {
+ let mut writer = ShuffleBlockWriter::try_new(schema.as_ref(),
codec.clone()).unwrap();
+ let mut buf_writer = BufBatchWriter::new(
+ &mut writer,
+ Cursor::new(&mut output),
+ 1024 * 1024,
+ batch_size as usize,
+ );
+ for batch in &inputs {
+ buf_writer.write(batch, &encode_time, &write_time).unwrap();
+ }
+ buf_writer.flush(&encode_time, &write_time).unwrap();
+ }
+
+ let blocks = read_all_ipc_batches(&output);
+
+ // The two full batches are emitted verbatim as their own blocks, then
the
+ // partial+full pair coalesces into a full block plus a 30-row tail
block.
+ let block_rows: Vec<usize> = blocks.iter().map(|b|
b.num_rows()).collect();
+ assert_eq!(
+ block_rows,
+ vec![100, 100, 100, 30],
Review Comment:
Thanks for the review @0lai0. I added a test based on your feedback.
--
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]
